v6: add RelayOptions with RelayMsg

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

6 files changed, 63 insertions, 66 deletions

diff --git a/dhcpv6/dhcpv6.go b/dhcpv6/dhcpv6.go
index a8f7bdc..07995e0 100644
--- a/dhcpv6/dhcpv6.go
+++ b/dhcpv6/dhcpv6.go
@@ -119,15 +119,10 @@ func DecapsulateRelay(l DHCPv6) (DHCPv6, error) {
 	if !l.IsRelay() {
 		return l, nil
 	}
-	opt := l.GetOneOption(OptionRelayMsg)
-	if opt == nil {
-		return nil, fmt.Errorf("malformed Relay message: no OptRelayMsg found")
+	if rm := l.(*RelayMessage).Options.RelayMessage(); rm != nil {
+		return rm, nil
 	}
-	relayOpt := opt.(*OptRelayMsg)
-	if relayOpt.RelayMessage() == nil {
-		return nil, fmt.Errorf("malformed Relay message: encapsulated message is empty")
-	}
-	return relayOpt.RelayMessage(), nil
+	return nil, fmt.Errorf("malformed Relay message: no embedded message found")
 }
 
 // DecapsulateRelayIndex extracts the content of a relay message. It takes an
@@ -180,8 +175,7 @@ func EncapsulateRelay(d DHCPv6, mType MessageType, linkAddr, peerAddr net.IP) (*
 	} else {
 		outer.HopCount = 0
 	}
-	orm := OptRelayMsg{relayMessage: d}
-	outer.AddOption(&orm)
+	outer.AddOption(OptRelayMessage(d))
 	return &outer, nil
 }
 
diff --git a/dhcpv6/dhcpv6relay.go b/dhcpv6/dhcpv6relay.go
index 209cc71..359e6cd 100644
--- a/dhcpv6/dhcpv6relay.go
+++ b/dhcpv6/dhcpv6relay.go
@@ -10,6 +10,26 @@ import (
 
 const RelayHeaderSize = 34
 
+// RelayOptions are the options valid for RelayForw and RelayRepl messages.
+//
+// RFC 3315 Appendix B defines them to be InterfaceID and RelayMsg options; RFC
+// 4649 also adds the RemoteID option.
+type RelayOptions struct {
+	Options
+}
+
+// RelayMessage returns the message embedded.
+func (ro RelayOptions) RelayMessage() DHCPv6 {
+	opt := ro.Options.GetOne(OptionRelayMsg)
+	if opt == nil {
+		return nil
+	}
+	if relayOpt, ok := opt.(*optRelayMsg); ok {
+		return relayOpt.Msg
+	}
+	return nil
+}
+
 // RelayMessage is a DHCPv6 relay agent message as defined by RFC 3315 Section
 // 7.
 type RelayMessage struct {
@@ -17,7 +37,7 @@ type RelayMessage struct {
 	HopCount    uint8
 	LinkAddr    net.IP
 	PeerAddr    net.IP
-	Options     Options
+	Options     RelayOptions
 }
 
 func write16(b *uio.Lexer, ip net.IP) {
@@ -38,7 +58,7 @@ func (r *RelayMessage) Type() MessageType {
 func (r *RelayMessage) String() string {
 	ret := fmt.Sprintf(
 		"RelayMessage(messageType=%s hopcount=%d, linkaddr=%s, peeraddr=%s, %d options)",
-		r.Type(), r.HopCount, r.LinkAddr, r.PeerAddr, len(r.Options),
+		r.Type(), r.HopCount, r.LinkAddr, r.PeerAddr, len(r.Options.Options),
 	)
 	return ret
 }
diff --git a/dhcpv6/dhcpv6relay_test.go b/dhcpv6/dhcpv6relay_test.go
index c383487..3e37e63 100644
--- a/dhcpv6/dhcpv6relay_test.go
+++ b/dhcpv6/dhcpv6relay_test.go
@@ -31,7 +31,7 @@ func TestRelayMessage(t *testing.T) {
 	if pa := r.PeerAddr; !pa.Equal(ma) {
 		t.Fatalf("Invalid peer address. Expected %v, got %v", ma, pa)
 	}
-	if opts := r.Options; len(opts) != 0 {
+	if opts := r.Options.Options; len(opts) != 0 {
 		t.Fatalf("Invalid options. Expected none, got %v", opts)
 	}
 }
@@ -70,18 +70,14 @@ func TestRelayMessageToBytes(t *testing.T) {
 		LinkAddr:    net.IPv6interfacelocalallnodes,
 		PeerAddr:    net.IPv6linklocalallrouters,
 	}
-	opt := OptRelayMsg{
-		relayMessage: &Message{
-			MessageType:   MessageTypeSolicit,
-			TransactionID: TransactionID{0xaa, 0xbb, 0xcc},
-			Options: MessageOptions{
-				Options: []Option{
-					OptElapsedTime(0),
-				},
-			},
-		},
-	}
-	r.AddOption(&opt)
+	opt := OptRelayMessage(&Message{
+		MessageType:   MessageTypeSolicit,
+		TransactionID: TransactionID{0xaa, 0xbb, 0xcc},
+		Options: MessageOptions{[]Option{
+			OptElapsedTime(0),
+		}},
+	})
+	r.AddOption(opt)
 	relayBytes := r.ToBytes()
 	if !bytes.Equal(expected, relayBytes) {
 		t.Fatalf("Invalid ToBytes result. Expected %v, got %v", expected, relayBytes)
@@ -101,9 +97,7 @@ func TestNewRelayRepFromRelayForw(t *testing.T) {
 	s, err := NewMessage()
 	require.NoError(t, err)
 	s.AddOption(OptClientID(Duid{}))
-	orm := OptRelayMsg{}
-	orm.SetRelayMessage(s)
-	rf.AddOption(&orm)
+	rf.AddOption(OptRelayMessage(s))
 
 	a, err := NewAdvertiseFromSolicit(s)
 	require.NoError(t, err)
diff --git a/dhcpv6/option_relaymsg.go b/dhcpv6/option_relaymsg.go
index 3d3d139..48d9716 100644
--- a/dhcpv6/option_relaymsg.go
+++ b/dhcpv6/option_relaymsg.go
@@ -1,42 +1,39 @@
 package dhcpv6
 
-// This module defines the OptRelayMsg structure.
+// This module defines the optRelayMsg structure.
 // https://www.ietf.org/rfc/rfc3315.txt
 
 import (
 	"fmt"
 )
 
-type OptRelayMsg struct {
-	relayMessage DHCPv6
+// OptRelayMessage embeds a message in a relay option.
+func OptRelayMessage(msg DHCPv6) Option {
+	return &optRelayMsg{Msg: msg}
 }
 
-func (op *OptRelayMsg) Code() OptionCode {
-	return OptionRelayMsg
-}
-
-func (op *OptRelayMsg) ToBytes() []byte {
-	return op.relayMessage.ToBytes()
+type optRelayMsg struct {
+	Msg DHCPv6
 }
 
-func (op *OptRelayMsg) RelayMessage() DHCPv6 {
-	return op.relayMessage
+func (op *optRelayMsg) Code() OptionCode {
+	return OptionRelayMsg
 }
 
-func (op *OptRelayMsg) SetRelayMessage(relayMessage DHCPv6) {
-	op.relayMessage = relayMessage
+func (op *optRelayMsg) ToBytes() []byte {
+	return op.Msg.ToBytes()
 }
 
-func (op *OptRelayMsg) String() string {
-	return fmt.Sprintf("OptRelayMsg{relaymsg=%v}", op.relayMessage)
+func (op *optRelayMsg) String() string {
+	return fmt.Sprintf("RelayMsg: %v", op.Msg)
 }
 
-// build an OptRelayMsg structure from a sequence of bytes.
+// build an optRelayMsg structure from a sequence of bytes.
 // The input data does not include option code and length bytes.
-func ParseOptRelayMsg(data []byte) (*OptRelayMsg, error) {
+func parseOptRelayMsg(data []byte) (*optRelayMsg, error) {
 	var err error
-	var opt OptRelayMsg
-	opt.relayMessage, err = FromBytes(data)
+	var opt optRelayMsg
+	opt.Msg, err = FromBytes(data)
 	if err != nil {
 		return nil, err
 	}
diff --git a/dhcpv6/option_relaymsg_test.go b/dhcpv6/option_relaymsg_test.go
index a39439b..b898a75 100644
--- a/dhcpv6/option_relaymsg_test.go
+++ b/dhcpv6/option_relaymsg_test.go
@@ -9,7 +9,7 @@ import (
 )
 
 func TestRelayMsgParseOptRelayMsg(t *testing.T) {
-	opt, err := ParseOptRelayMsg([]byte{
+	opt, err := parseOptRelayMsg([]byte{
 		1,                // MessageTypeSolicit
 		0xaa, 0xbb, 0xcc, // transaction ID
 		0, 8, // option: elapsed time
@@ -77,22 +77,14 @@ func TestRelayMsgParseOptRelayMsgSingleEncapsulation(t *testing.T) {
 	if mType := r.Type(); mType != MessageTypeRelayForward {
 		t.Fatalf("Invalid messge type for relay. Expected %v, got %v", MessageTypeRelayForward, mType)
 	}
-	if len(r.Options) != 1 {
-		t.Fatalf("Invalid number of options. Expected 1, got %v", len(r.Options))
+	if len(r.Options.Options) != 1 {
+		t.Fatalf("Invalid number of options. Expected 1, got %v", len(r.Options.Options))
 	}
-	if code := r.Options[0].Code(); code != OptionRelayMsg {
-		t.Fatalf("Invalid option code. Expected OptionRelayMsg (%v), got %v",
-			OptionRelayMsg, code,
-		)
-	}
-	opt := r.Options[0]
-	ro, ok := opt.(*OptRelayMsg)
-	if !ok {
-		t.Fatalf("Invalid option type. Expected OptRelayMsg, got %v",
-			reflect.TypeOf(ro),
-		)
+	ro := r.Options.RelayMessage()
+	if ro == nil {
+		t.Fatalf("No relay message")
 	}
-	innerDHCP, ok := ro.RelayMessage().(*Message)
+	innerDHCP, ok := ro.(*Message)
 	if !ok {
 		t.Fatalf("Invalid relay message type. Expected Message, got %v",
 			reflect.TypeOf(innerDHCP),
@@ -156,7 +148,7 @@ func TestSample(t *testing.T) {
 }
 
 func TestRelayMsgParseOptRelayMsgTooShort(t *testing.T) {
-	_, err := ParseOptRelayMsg([]byte{
+	_, err := parseOptRelayMsg([]byte{
 		1,                // MessageTypeSolicit
 		0xaa, 0xbb, 0xcc, // transaction ID
 		0, 8, // option: elapsed time
@@ -166,7 +158,7 @@ func TestRelayMsgParseOptRelayMsgTooShort(t *testing.T) {
 }
 
 func TestRelayMsgString(t *testing.T) {
-	opt, err := ParseOptRelayMsg([]byte{
+	opt, err := parseOptRelayMsg([]byte{
 		1,                // MessageTypeSolicit
 		0xaa, 0xbb, 0xcc, // transaction ID
 		0, 8, // option: elapsed time
@@ -177,7 +169,7 @@ func TestRelayMsgString(t *testing.T) {
 	require.Contains(
 		t,
 		opt.String(),
-		"relaymsg=Message",
+		"RelayMsg: Message",
 		"String() should contain the relaymsg contents",
 	)
 }
diff --git a/dhcpv6/options.go b/dhcpv6/options.go
index 87a33ff..d0bde7c 100644
--- a/dhcpv6/options.go
+++ b/dhcpv6/options.go
@@ -54,7 +54,7 @@ func ParseOption(code OptionCode, optData []byte) (Option, error) {
 	case OptionElapsedTime:
 		opt, err = parseOptElapsedTime(optData)
 	case OptionRelayMsg:
-		opt, err = ParseOptRelayMsg(optData)
+		opt, err = parseOptRelayMsg(optData)
 	case OptionStatusCode:
 		opt, err = ParseOptStatusCode(optData)
 	case OptionUserClass: