Fixed OptRelayMsg.ToBytes and added tests

2 files changed, 47 insertions, 1 deletions

diff --git a/dhcpv6/dhcpv6relay_test.go b/dhcpv6/dhcpv6relay_test.go
index 983a882..99704b9 100644
--- a/dhcpv6/dhcpv6relay_test.go
+++ b/dhcpv6/dhcpv6relay_test.go
@@ -63,3 +63,44 @@ func TestDHCPv6RelaySettersAndGetters(t *testing.T) {
 		t.Fatalf("Invalid options. Expected %v, got %v", oneOpt, opts)
 	}
 }
+
+func TestDHCPv6RelayToBytes(t *testing.T) {
+	expected := []byte{
+		12,                                             // RELAY_FORW
+		1,                                              // hop count
+		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, // link addr
+		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, // peer addr
+		// option relay message
+		0, 9, // relay msg
+		0, 10, // option length
+		// inner dhcp solicit
+		1,                // SOLICIT
+		0xaa, 0xbb, 0xcc, // transaction ID
+		// inner option - elapsed time
+		0, 8, // elapsed time
+		0, 2, // length
+		0, 0,
+	}
+	r := DHCPv6Relay{
+		messageType: RELAY_FORW,
+		hopCount:    1,
+		linkAddr:    net.IPv6loopback,
+		peerAddr:    net.IPv6loopback,
+	}
+	opt := OptRelayMsg{
+		relayMessage: &DHCPv6Message{
+			messageType:   SOLICIT,
+			transactionID: 0xaabbcc,
+			options: []Option{
+				&OptElapsedTime{
+					elapsedTime: 0,
+				},
+			},
+		},
+	}
+	r.AddOption(&opt)
+	relayBytes := r.ToBytes()
+	if !bytes.Equal(expected, relayBytes) {
+		t.Fatalf("Invalid ToBytes result. Expected %v, got %v", expected, relayBytes)
+	}
+}
diff --git a/dhcpv6/option_relaymsg.go b/dhcpv6/option_relaymsg.go
index 3b5f865..98b9e84 100644
--- a/dhcpv6/option_relaymsg.go
+++ b/dhcpv6/option_relaymsg.go
@@ -4,6 +4,7 @@ package dhcpv6
 // https://www.ietf.org/rfc/rfc3315.txt
 
 import (
+	"encoding/binary"
 	"fmt"
 )
 
@@ -16,7 +17,11 @@ func (op *OptRelayMsg) Code() OptionCode {
 }
 
 func (op *OptRelayMsg) ToBytes() []byte {
-	return op.relayMessage.ToBytes()
+	buf := make([]byte, 4)
+	binary.BigEndian.PutUint16(buf[0:2], uint16(OPTION_RELAY_MSG))
+	binary.BigEndian.PutUint16(buf[2:4], uint16(op.Length()))
+	buf = append(buf, op.relayMessage.ToBytes()...)
+	return buf
 }
 
 func (op *OptRelayMsg) RelayMessage() DHCPv6 {