Merge pull request #161 from get9/bsdp-message-type-helper

BSDP: Add MessageTypeFromPacket helper

2 files changed, 92 insertions, 0 deletions

diff --git a/dhcpv4/bsdp/bsdp.go b/dhcpv4/bsdp/bsdp.go
index e15c59c..172399b 100644
--- a/dhcpv4/bsdp/bsdp.go
+++ b/dhcpv4/bsdp/bsdp.go
@@ -51,6 +51,29 @@ func needsReplyPort(replyPort uint16) bool {
 	return replyPort != 0 && replyPort != dhcpv4.ClientPort
 }
 
+// MessageTypeFromPacket extracts the BSDP message type (LIST, SELECT) from the
+// vendor-specific options and returns it. If the message type option cannot be
+// found, returns false.
+func MessageTypeFromPacket(packet *dhcpv4.DHCPv4) (MessageType, bool) {
+	var (
+		messageType MessageType
+		vendorOpts  *OptVendorSpecificInformation
+		err         error
+	)
+	for _, opt := range packet.GetOption(dhcpv4.OptionVendorSpecificInformation) {
+		if vendorOpts, err = ParseOptVendorSpecificInformation(opt.ToBytes()); err != nil {
+			return messageType, false
+		}
+		if o := vendorOpts.GetOneOption(OptionMessageType); o != nil {
+			if optMessageType, ok := o.(*OptMessageType); ok {
+				return optMessageType.Type, true
+			}
+			return messageType, false
+		}
+	}
+	return messageType, false
+}
+
 // NewInformListForInterface creates a new INFORM packet for interface ifname
 // with configuration options specified by config.
 func NewInformListForInterface(ifname string, replyPort uint16) (*dhcpv4.DHCPv4, error) {
diff --git a/dhcpv4/bsdp/bsdp_test.go b/dhcpv4/bsdp/bsdp_test.go
index 7da4dfa..c9c868a 100644
--- a/dhcpv4/bsdp/bsdp_test.go
+++ b/dhcpv4/bsdp/bsdp_test.go
@@ -371,3 +371,72 @@ func TestNewReplyForInformSelect(t *testing.T) {
 	RequireHasOption(t, vendorOpts, &OptMessageType{Type: MessageTypeSelect})
 	RequireHasOption(t, vendorOpts, &OptSelectedBootImageID{ID: images[0].ID})
 }
+
+func TestMessageTypeForPacket(t *testing.T) {
+	var (
+		pkt            *dhcpv4.DHCPv4
+		gotMessageType MessageType
+		gotOK          bool
+	)
+
+	testcases := []struct {
+		tcName          string
+		opts            []dhcpv4.Option
+		wantOK          bool
+		wantMessageType MessageType
+	}{
+		{
+			tcName: "No options",
+			opts:   []dhcpv4.Option{},
+			wantOK: false,
+		},
+		{
+			tcName: "Some options, no vendor opts",
+			opts: []dhcpv4.Option{
+				&dhcpv4.OptHostName{HostName: "foobar1234"},
+			},
+			wantOK: false,
+		},
+		{
+			tcName: "Vendor opts, no message type",
+			opts: []dhcpv4.Option{
+				&dhcpv4.OptHostName{HostName: "foobar1234"},
+				&OptVendorSpecificInformation{
+					Options: []dhcpv4.Option{
+						&OptVersion{Version: Version1_1},
+					},
+				},
+			},
+			wantOK: false,
+		},
+		{
+			tcName: "Vendor opts, with message type",
+			opts: []dhcpv4.Option{
+				&dhcpv4.OptHostName{HostName: "foobar1234"},
+				&OptVendorSpecificInformation{
+					Options: []dhcpv4.Option{
+						&OptVersion{Version: Version1_1},
+						&OptMessageType{Type: MessageTypeList},
+					},
+				},
+			},
+			wantOK:          true,
+			wantMessageType: MessageTypeList,
+		},
+	}
+	for _, tt := range testcases {
+		t.Run(tt.tcName, func(t *testing.T) {
+			pkt, _ = dhcpv4.New()
+			for _, opt := range tt.opts {
+				pkt.AddOption(opt)
+			}
+			gotMessageType, gotOK = MessageTypeFromPacket(pkt)
+			if tt.wantOK {
+				require.True(t, gotOK)
+				require.Equal(t, tt.wantMessageType, gotMessageType)
+			} else {
+				require.False(t, gotOK)
+			}
+		})
+	}
+}