v6: add RemoteID getter to RelayOptions

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

7 files changed, 58 insertions, 74 deletions

diff --git a/dhcpv6/dhcpv6relay.go b/dhcpv6/dhcpv6relay.go
index 04a3b94..9c08e00 100644
--- a/dhcpv6/dhcpv6relay.go
+++ b/dhcpv6/dhcpv6relay.go
@@ -42,6 +42,18 @@ func (ro RelayOptions) InterfaceID() []byte {
 	return nil
 }
 
+// RemoteID returns the remote ID in this relay message.
+func (ro RelayOptions) RemoteID() *OptRemoteID {
+	opt := ro.Options.GetOne(OptionRemoteID)
+	if opt == nil {
+		return nil
+	}
+	if rid, ok := opt.(*OptRemoteID); ok {
+		return rid
+	}
+	return nil
+}
+
 // RelayMessage is a DHCPv6 relay agent message as defined by RFC 3315 Section
 // 7.
 type RelayMessage struct {
diff --git a/dhcpv6/dhcpv6relay_test.go b/dhcpv6/dhcpv6relay_test.go
index 9bfe2e2..113842c 100644
--- a/dhcpv6/dhcpv6relay_test.go
+++ b/dhcpv6/dhcpv6relay_test.go
@@ -91,7 +91,7 @@ func TestNewRelayRepFromRelayForw(t *testing.T) {
 	rf.PeerAddr = net.IPv6linklocalallrouters
 	rf.LinkAddr = net.IPv6interfacelocalallnodes
 	rf.AddOption(OptInterfaceID(nil))
-	rf.AddOption(&OptRemoteId{})
+	rf.AddOption(&OptRemoteID{})
 
 	// create the inner message
 	s, err := NewMessage()
diff --git a/dhcpv6/option_remoteid.go b/dhcpv6/option_remoteid.go
index e6803c6..494da20 100644
--- a/dhcpv6/option_remoteid.go
+++ b/dhcpv6/option_remoteid.go
@@ -6,54 +6,37 @@ import (
 	"github.com/u-root/u-root/pkg/uio"
 )
 
-// OptRemoteId implemens the Remote ID option.
-//
-// https://www.ietf.org/rfc/rfc4649.txt
-type OptRemoteId struct {
-	enterpriseNumber uint32
-	remoteId         []byte
+// OptRemoteID implemens the Remote ID option as defined by RFC 4649.
+type OptRemoteID struct {
+	EnterpriseNumber uint32
+	RemoteID         []byte
 }
 
-func (op *OptRemoteId) Code() OptionCode {
+// Code implements Option.Code.
+func (*OptRemoteID) Code() OptionCode {
 	return OptionRemoteID
 }
 
 // ToBytes serializes this option to a byte stream.
-func (op *OptRemoteId) ToBytes() []byte {
+func (op *OptRemoteID) ToBytes() []byte {
 	buf := uio.NewBigEndianBuffer(nil)
-	buf.Write32(op.enterpriseNumber)
-	buf.WriteBytes(op.remoteId)
+	buf.Write32(op.EnterpriseNumber)
+	buf.WriteBytes(op.RemoteID)
 	return buf.Data()
 }
 
-func (op *OptRemoteId) EnterpriseNumber() uint32 {
-	return op.enterpriseNumber
-}
-
-func (op *OptRemoteId) SetEnterpriseNumber(enterpriseNumber uint32) {
-	op.enterpriseNumber = enterpriseNumber
-}
-
-func (op *OptRemoteId) RemoteID() []byte {
-	return op.remoteId
-}
-
-func (op *OptRemoteId) SetRemoteID(remoteId []byte) {
-	op.remoteId = append([]byte(nil), remoteId...)
-}
-
-func (op *OptRemoteId) String() string {
-	return fmt.Sprintf("OptRemoteId{enterprisenum=%v, remoteid=%v}",
-		op.enterpriseNumber, op.remoteId,
+func (op *OptRemoteID) String() string {
+	return fmt.Sprintf("RemoteID: EnterpriseNumber %d RemoteID %v",
+		op.EnterpriseNumber, op.RemoteID,
 	)
 }
 
 // ParseOptRemoteId builds an OptRemoteId structure from a sequence of bytes.
 // The input data does not include option code and length bytes.
-func ParseOptRemoteId(data []byte) (*OptRemoteId, error) {
-	var opt OptRemoteId
+func ParseOptRemoteID(data []byte) (*OptRemoteID, error) {
+	var opt OptRemoteID
 	buf := uio.NewBigEndianBuffer(data)
-	opt.enterpriseNumber = buf.Read32()
-	opt.remoteId = buf.ReadAll()
+	opt.EnterpriseNumber = buf.Read32()
+	opt.RemoteID = buf.ReadAll()
 	return &opt, buf.FinError()
 }
diff --git a/dhcpv6/option_remoteid_test.go b/dhcpv6/option_remoteid_test.go
index b8b90f4..5c1ee52 100644
--- a/dhcpv6/option_remoteid_test.go
+++ b/dhcpv6/option_remoteid_test.go
@@ -7,27 +7,27 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestOptRemoteId(t *testing.T) {
+func TestOptRemoteID(t *testing.T) {
 	expected := []byte{0xaa, 0xbb, 0xcc, 0xdd}
 	remoteId := []byte("DSLAM01 eth2/1/01/21")
 	expected = append(expected, remoteId...)
-	opt, err := ParseOptRemoteId(expected)
+	opt, err := ParseOptRemoteID(expected)
 	if err != nil {
 		t.Fatal(err)
 	}
-	if en := opt.EnterpriseNumber(); en != 0xaabbccdd {
+	if en := opt.EnterpriseNumber; en != 0xaabbccdd {
 		t.Fatalf("Invalid Enterprise Number. Expected 0xaabbccdd, got %v", en)
 	}
-	if rid := opt.RemoteID(); !bytes.Equal(rid, remoteId) {
+	if rid := opt.RemoteID; !bytes.Equal(rid, remoteId) {
 		t.Fatalf("Invalid Remote ID. Expected %v, got %v", expected, rid)
 	}
 }
 
-func TestOptRemoteIdToBytes(t *testing.T) {
+func TestOptRemoteIDToBytes(t *testing.T) {
 	remoteId := []byte("DSLAM01 eth2/1/01/21")
 	expected := append([]byte{0, 0, 0, 0}, remoteId...)
-	opt := OptRemoteId{
-		remoteId: remoteId,
+	opt := OptRemoteID{
+		RemoteID: remoteId,
 	}
 	toBytes := opt.ToBytes()
 	if !bytes.Equal(toBytes, expected) {
@@ -35,41 +35,30 @@ func TestOptRemoteIdToBytes(t *testing.T) {
 	}
 }
 
-func TestOptRemoteIdSet(t *testing.T) {
-	enterpriseNumber := uint32(12345)
-	remoteID := []byte("DSLAM01 eth2/1/01/21")
-	opt := OptRemoteId{}
-	opt.SetEnterpriseNumber(enterpriseNumber)
-	opt.SetRemoteID(remoteID)
-
-	require.Equal(t, uint32(12345), opt.EnterpriseNumber())
-	require.Equal(t, []byte("DSLAM01 eth2/1/01/21"), opt.RemoteID())
-}
-
-func TestOptRemoteIdParseOptRemoteIdTooShort(t *testing.T) {
+func TestOptRemoteIDParseOptRemoteIDTooShort(t *testing.T) {
 	buf := []byte{0xaa, 0xbb, 0xcc}
-	_, err := ParseOptRemoteId(buf)
+	_, err := ParseOptRemoteID(buf)
 	require.Error(t, err, "A short option should return an error")
 }
 
-func TestOptRemoteIdString(t *testing.T) {
+func TestOptRemoteIDString(t *testing.T) {
 	buf := []byte{0xaa, 0xbb, 0xcc, 0xdd}
 	remoteId := []byte("Test1234")
 	buf = append(buf, remoteId...)
 
-	opt, err := ParseOptRemoteId(buf)
+	opt, err := ParseOptRemoteID(buf)
 	require.NoError(t, err)
 	str := opt.String()
 	require.Contains(
 		t,
 		str,
-		"enterprisenum=2864434397",
+		"EnterpriseNumber 2864434397",
 		"String() should contain the enterprisenum",
 	)
 	require.Contains(
 		t,
 		str,
-		"remoteid=[84 101 115 116 49 50 51 52]",
+		"RemoteID [84 101 115 116 49 50 51 52]",
 		"String() should contain the remoteid bytes",
 	)
 }
diff --git a/dhcpv6/options.go b/dhcpv6/options.go
index f91bcc4..752ca2c 100644
--- a/dhcpv6/options.go
+++ b/dhcpv6/options.go
@@ -74,7 +74,7 @@ func ParseOption(code OptionCode, optData []byte) (Option, error) {
 	case OptionIAPrefix:
 		opt, err = ParseOptIAPrefix(optData)
 	case OptionRemoteID:
-		opt, err = ParseOptRemoteId(optData)
+		opt, err = ParseOptRemoteID(optData)
 	case OptionFQDN:
 		opt, err = ParseOptFQDN(optData)
 	case OptionBootfileURL:
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)
 			}