nclient6: allow non-rapid responses to RapidSolicit

Signed-off-by: Chris Koch <chrisko@google.com>

1 files changed, 24 insertions, 6 deletions

diff --git a/dhcpv6/nclient6/client.go b/dhcpv6/nclient6/client.go
index a90d050..74cdf5c 100644
--- a/dhcpv6/nclient6/client.go
+++ b/dhcpv6/nclient6/client.go
@@ -294,11 +294,17 @@ func WithDebugLogger() ClientOpt {
 type Matcher func(*dhcpv6.Message) bool
 
 // IsMessageType returns a matcher that checks for the message type.
-//
-// If t is MessageTypeNone, all packets are matched.
-func IsMessageType(t dhcpv6.MessageType) Matcher {
+func IsMessageType(t dhcpv6.MessageType, tt ...dhcpv6.MessageType) Matcher {
 	return func(p *dhcpv6.Message) bool {
-		return p.MessageType == t || t == dhcpv6.MessageTypeNone
+		if p.MessageType == t {
+			return true
+		}
+		for _, mt := range tt {
+			if p.MessageType == mt {
+				return true
+			}
+		}
+		return false
 	}
 }
 
@@ -309,11 +315,23 @@ func (c *Client) RapidSolicit(ctx context.Context, modifiers ...dhcpv6.Modifier)
 	if err != nil {
 		return nil, err
 	}
-	msg, err := c.SendAndRead(ctx, c.serverAddr, solicit, IsMessageType(dhcpv6.MessageTypeReply))
+	msg, err := c.SendAndRead(ctx, c.serverAddr, solicit, IsMessageType(dhcpv6.MessageTypeReply, dhcpv6.MessageTypeAdvertise))
 	if err != nil {
 		return nil, err
 	}
-	return msg, nil
+
+	switch msg.MessageType {
+	case dhcpv6.MessageTypeReply:
+		// We got RapidCommitted.
+		return msg, nil
+
+	case dhcpv6.MessageTypeAdvertise:
+		// We didn't get RapidCommitted. Request regular lease.
+		return c.Request(ctx, msg, modifiers...)
+
+	default:
+		return nil, fmt.Errorf("invalid message type: cannot happen")
+	}
 }
 
 // Solicit sends a solicitation message and returns the first valid