Added DHCPv6.Length() and improved relay msg parsing

6 files changed, 77 insertions, 4 deletions

diff --git a/dhcpv6/dhcpv6.go b/dhcpv6/dhcpv6.go
index fe096a6..97eb401 100644
--- a/dhcpv6/dhcpv6.go
+++ b/dhcpv6/dhcpv6.go
@@ -8,6 +8,7 @@ type DHCPv6 interface {
 	Type() MessageType
 	ToBytes() []byte
 	Summary() string
+	Length() int
 }
 
 func FromBytes(data []byte) (DHCPv6, error) {
@@ -28,16 +29,22 @@ func FromBytes(data []byte) (DHCPv6, error) {
 		return nil, fmt.Errorf("Invalid header size: shorter than %v bytes", headerSize)
 	}
 	if isRelay {
+		var (
+			linkAddr, peerAddr []byte
+		)
 		d := DHCPv6Relay{
 			messageType: messageType,
 			hopCount:    uint8(data[1]),
-			linkAddr:    append(data[2:18]),
-			peerAddr:    append(data[18:34]),
 		}
+		linkAddr = append(linkAddr, data[2:18]...)
+		d.linkAddr = linkAddr
+		peerAddr = append(peerAddr, data[18:34]...)
+		d.peerAddr = peerAddr
 		options, err := OptionsFromBytes(data[34:])
 		if err != nil {
 			return nil, err
 		}
+		// TODO fail if no OptRelayMessage is present
 		d.options = options
 		return &d, nil
 	} else {
diff --git a/dhcpv6/dhcpv6message.go b/dhcpv6/dhcpv6message.go
index 417c46d..b781d6e 100644
--- a/dhcpv6/dhcpv6message.go
+++ b/dhcpv6/dhcpv6message.go
@@ -182,3 +182,11 @@ func (d *DHCPv6Message) ToBytes() []byte {
 	}
 	return ret
 }
+
+func (d *DHCPv6Message) Length() int {
+	mLen := 4
+	for _, opt := range d.options {
+		mLen += opt.Length() + 4 // +4 for opt code and opt len
+	}
+	return mLen
+}
diff --git a/dhcpv6/dhcpv6relay.go b/dhcpv6/dhcpv6relay.go
index 12b677a..b8d635e 100644
--- a/dhcpv6/dhcpv6relay.go
+++ b/dhcpv6/dhcpv6relay.go
@@ -47,8 +47,16 @@ func (r *DHCPv6Relay) ToBytes() []byte {
 	copy(ret[2:18], r.peerAddr)
 	copy(ret[18:34], r.linkAddr)
 	for _, opt := range r.options {
-		ret = append(opt.ToBytes())
+		ret = append(ret, opt.ToBytes()...)
 	}
 
 	return ret
 }
+
+func (r *DHCPv6Relay) Length() int {
+	mLen := RelayHeaderSize
+	for _, opt := range r.options {
+		mLen += opt.Length()
+	}
+	return mLen
+}
diff --git a/dhcpv6/option_relaymsg.go b/dhcpv6/option_relaymsg.go
index b0a69a5..3b5f865 100644
--- a/dhcpv6/option_relaymsg.go
+++ b/dhcpv6/option_relaymsg.go
@@ -28,7 +28,7 @@ func (op *OptRelayMsg) SetRelayMessage(relayMessage DHCPv6) {
 }
 
 func (op *OptRelayMsg) Length() int {
-	return len(op.relayMessage.ToBytes())
+	return op.relayMessage.Length()
 }
 
 func (op *OptRelayMsg) String() string {
diff --git a/dhcpv6/option_relaymsg_test.go b/dhcpv6/option_relaymsg_test.go
new file mode 100644
index 0000000..95fb9fb
--- /dev/null
+++ b/dhcpv6/option_relaymsg_test.go
@@ -0,0 +1,47 @@
+package dhcpv6
+
+import (
+	"testing"
+)
+
+func TestRelayMsgParseOptRelayMsg(t *testing.T) {
+	opt, err := ParseOptRelayMsg([]byte{
+		1,                // SOLICIT
+		0xaa, 0xbb, 0xcc, // transaction ID
+		0, 8, // option: elapsed time
+		0, 2, // option length
+		0, 0, // option value
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if code := opt.Code(); code != OPTION_RELAY_MSG {
+		t.Fatalf("Invalid option code. Expected OPTION_RELAY_MSG (%v), got %v",
+			OPTION_RELAY_MSG, code,
+		)
+	}
+}
+
+func TestRelayMsgOptionsFromBytes(t *testing.T) {
+	opts, err := OptionsFromBytes([]byte{
+		0, 9, // option: relay message
+		0, 10, // relayed message length
+		1,                // SOLICIT
+		0xaa, 0xbb, 0xcc, // transaction ID
+		0, 8, // option: elapsed time
+		0, 2, // option length
+		0, 0, // option value
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(opts) != 1 {
+		t.Fatalf("Invalid number of options. Expected 1, got %v", len(opts))
+	}
+	opt := opts[0]
+	if code := opt.Code(); code != OPTION_RELAY_MSG {
+		t.Fatalf("Invalid option code. Expected OPTION_RELAY_MSG (%v), got %v",
+			OPTION_RELAY_MSG, code,
+		)
+	}
+}
diff --git a/dhcpv6/options.go b/dhcpv6/options.go
index 2c5842c..2fcc9e9 100644
--- a/dhcpv6/options.go
+++ b/dhcpv6/options.go
@@ -96,6 +96,9 @@ func ParseOption(dataStart []byte) (Option, error) {
 	if err != nil {
 		return nil, err
 	}
+	if length != opt.Length() {
+		return nil, fmt.Errorf("different lengths for opt %v: %v != %v", code, length, opt.Length())
+	}
 	return opt, nil
 }