v6: add RemoteID getter to RelayOptions

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

2 files changed, 15 insertions, 15 deletions

diff --git a/dhcpv6/ztpv6/parse_remote_id.go b/dhcpv6/ztpv6/parse_remote_id.go
index 5991e96..9de9bdf 100644
--- a/dhcpv6/ztpv6/parse_remote_id.go
+++ b/dhcpv6/ztpv6/parse_remote_id.go
@@ -25,20 +25,22 @@ type CircuitID struct {
 }
 
 // ParseRemoteId will parse the RemoteId Option data for Vendor Specific data
-func ParseRemoteId(packet dhcpv6.DHCPv6) (*CircuitID, error) {
+func ParseRemoteID(packet dhcpv6.DHCPv6) (*CircuitID, error) {
 	// Need to decapsulate the packet after multiple relays in order to reach RemoteId data
 	inner, err := dhcpv6.DecapsulateRelayIndex(packet, -1)
 	if err != nil {
 		return nil, fmt.Errorf("failed to decapsulate relay index: %v", err)
 	}
 
-	if rid := inner.GetOneOption(dhcpv6.OptionRemoteID); rid != nil {
-		remoteID := string(rid.(*dhcpv6.OptRemoteId).RemoteID())
-		circ, err := matchCircuitId(remoteID)
-		if err != nil {
-			return nil, err
+	if rm, ok := inner.(*dhcpv6.RelayMessage); ok {
+		if rid := rm.Options.RemoteID(); rid != nil {
+			remoteID := string(rid.RemoteID)
+			circ, err := matchCircuitId(remoteID)
+			if err != nil {
+				return nil, err
+			}
+			return circ, nil
 		}
-		return circ, nil
 	}
 	return nil, errors.New("failed to parse RemoteID option data")
 }
diff --git a/dhcpv6/ztpv6/parse_remote_id_test.go b/dhcpv6/ztpv6/parse_remote_id_test.go
index 3f4c02e..48d20bc 100644
--- a/dhcpv6/ztpv6/parse_remote_id_test.go
+++ b/dhcpv6/ztpv6/parse_remote_id_test.go
@@ -64,16 +64,14 @@ func TestParseRemoteID(t *testing.T) {
 	}
 	for _, tc := range tt {
 		t.Run(tc.name, func(t *testing.T) {
-			packet, err := dhcpv6.NewMessage()
-			if err != nil {
-				t.Fatalf("failed to creat dhcpv6 packet object: %v", err)
+			m := &dhcpv6.RelayMessage{
+				MessageType: dhcpv6.MessageTypeRelayForward,
 			}
-			opt := dhcpv6.OptRemoteId{}
-			opt.SetRemoteID(tc.circuit)
-			opt.SetEnterpriseNumber(1234)
-			packet.AddOption(&opt)
+			// Has to be a well-formed relay message with the OptRelayMsg.
+			m.Options.Add(dhcpv6.OptRelayMessage(&dhcpv6.Message{}))
+			m.Options.Add(&dhcpv6.OptRemoteID{RemoteID: tc.circuit, EnterpriseNumber: 1234})
 
-			circuit, err := ParseRemoteId(packet)
+			circuit, err := ParseRemoteID(m)
 			if err != nil && !tc.fail {
 				t.Errorf("unexpected failure: %v", err)
 			}