DHCPv6 client now waits for specific packet types (#38)

Before this patch the client just took whatever DHCP message it
received, without checking if it's what it wants. Now it checks for the
message type.
If the message is not a relay, also check the transaction ID

2 files changed, 48 insertions, 9 deletions

diff --git a/dhcpv6/client.go b/dhcpv6/client.go
index ec1e09d..f46854c 100644
--- a/dhcpv6/client.go
+++ b/dhcpv6/client.go
@@ -59,10 +59,22 @@ func (c *Client) Exchange(ifname string, solicit DHCPv6) ([]DHCPv6, error) {
 	return conversation, nil
 }
 
-func (c *Client) sendReceive(ifname string, packet DHCPv6) (DHCPv6, error) {
+func (c *Client) sendReceive(ifname string, packet DHCPv6, expectedType MessageType) (DHCPv6, error) {
 	if packet == nil {
 		return nil, fmt.Errorf("Packet to send cannot be nil")
 	}
+	if expectedType == MSGTYPE_NONE {
+		// infer the expected type from the packet being sent
+		if packet.Type() == SOLICIT {
+			expectedType = ADVERTISE
+		} else if packet.Type() == REQUEST {
+			expectedType = REPLY
+		} else if packet.Type() == RELAY_FORW {
+			expectedType = RELAY_REPL
+		} else if packet.Type() == LEASEQUERY {
+			expectedType = LEASEQUERY_REPLY
+		} // and probably more
+	}
 	// if no LocalAddr is specified, get the interface's link-local address
 	var laddr net.UDPAddr
 	if c.LocalAddr == nil {
@@ -109,13 +121,38 @@ func (c *Client) sendReceive(ifname string, packet DHCPv6) (DHCPv6, error) {
 	buf := make([]byte, maxUDPReceivedPacketSize)
 	oobdata := []byte{} // ignoring oob data
 	conn.SetReadDeadline(time.Now().Add(c.ReadTimeout))
-	n, _, _, _, err := conn.ReadMsgUDP(buf, oobdata)
-	if err != nil {
-		return nil, err
+	var (
+		adv       DHCPv6
+		isMessage bool
+	)
+	msg, ok := packet.(*DHCPv6Message)
+	if ok {
+		isMessage = true
 	}
-	adv, err := FromBytes(buf[:n])
-	if err != nil {
-		return nil, err
+	for {
+		n, _, _, _, err := conn.ReadMsgUDP(buf, oobdata)
+		if err != nil {
+			return nil, err
+		}
+		adv, err = FromBytes(buf[:n])
+		if err != nil {
+			return nil, err
+		}
+		if recvMsg, ok := adv.(*DHCPv6Message); ok && isMessage {
+			// if a regular message, check the transaction ID first
+			// XXX should this unpack relay messages and check the XID of the
+			// inner packet too?
+			if msg.TransactionID() != recvMsg.TransactionID() {
+				// different XID, we don't want this packet for sure
+				continue
+			}
+		}
+		if expectedType == MSGTYPE_NONE {
+			// just take whatever arrived
+			break
+		} else if adv.Type() == expectedType {
+			break
+		}
 	}
 	return adv, nil
 }
@@ -130,7 +167,7 @@ func (c *Client) Solicit(ifname string, solicit DHCPv6) (DHCPv6, DHCPv6, error)
 			return nil, nil, err
 		}
 	}
-	advertise, err := c.sendReceive(ifname, solicit)
+	advertise, err := c.sendReceive(ifname, solicit, MSGTYPE_NONE)
 	return solicit, advertise, err
 }
 
@@ -144,6 +181,6 @@ func (c *Client) Request(ifname string, advertise, request DHCPv6) (DHCPv6, DHCP
 			return nil, nil, err
 		}
 	}
-	reply, err := c.sendReceive(ifname, request)
+	reply, err := c.sendReceive(ifname, request, MSGTYPE_NONE)
 	return request, reply, err
 }
diff --git a/dhcpv6/types.go b/dhcpv6/types.go
index bbe4c5b..795a284 100644
--- a/dhcpv6/types.go
+++ b/dhcpv6/types.go
@@ -5,6 +5,8 @@ package dhcpv6
 type MessageType uint8
 
 const (
+	// MSGTYPE_NONE is used internally and is not part of the RFC
+	MSGTYPE_NONE        MessageType = 0
 	SOLICIT             MessageType = 1
 	ADVERTISE           MessageType = 2
 	REQUEST             MessageType = 3