Added WithClientId, WithServerId and tests (#56)

2 files changed, 56 insertions, 1 deletions

diff --git a/dhcpv6/modifiers.go b/dhcpv6/modifiers.go
index 08745bc..54285cd 100644
--- a/dhcpv6/modifiers.go
+++ b/dhcpv6/modifiers.go
@@ -4,6 +4,24 @@ import (
 	"log"
 )
 
+// WithClientID adds a client ID option to a DHCPv6 packet
+func WithClientID(duid Duid) Modifier {
+	return func(d DHCPv6) DHCPv6 {
+		cid := OptClientId{Cid: duid}
+		d.UpdateOption(&cid)
+		return d
+	}
+}
+
+// WithServerID adds a client ID option to a DHCPv6 packet
+func WithServerID(duid Duid) Modifier {
+	return func(d DHCPv6) DHCPv6 {
+		sid := OptServerId{Sid: duid}
+		d.UpdateOption(&sid)
+		return d
+	}
+}
+
 // WithNetboot adds bootfile URL and bootfile param options to a DHCPv6 packet.
 func WithNetboot(d DHCPv6) DHCPv6 {
 	msg, ok := d.(*DHCPv6Message)
@@ -26,7 +44,7 @@ func WithNetboot(d DHCPv6) DHCPv6 {
 
 // WithUserClass adds a user class option to the packet
 func WithUserClass(uc []byte) Modifier {
-    // TODO let the user specify multiple user classes
+	// TODO let the user specify multiple user classes
 	return func(d DHCPv6) DHCPv6 {
 		ouc := OptUserClass{UserClasses: [][]byte{uc}}
 		d.AddOption(&ouc)
diff --git a/dhcpv6/modifiers_test.go b/dhcpv6/modifiers_test.go
new file mode 100644
index 0000000..da1ef56
--- /dev/null
+++ b/dhcpv6/modifiers_test.go
@@ -0,0 +1,37 @@
+package dhcpv6
+
+import (
+	"net"
+	"testing"
+
+	"github.com/insomniacslk/dhcp/iana"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithClientID(t *testing.T) {
+	duid := Duid{
+		Type:          DUID_LL,
+		HwType:        iana.HwTypeEthernet,
+		LinkLayerAddr: net.HardwareAddr([]byte{0xfa, 0xce, 0xb0, 0x00, 0x00, 0x0c}),
+	}
+	m, err := NewMessage(WithClientID(duid))
+	require.NoError(t, err)
+	opt := m.GetOneOption(OPTION_CLIENTID)
+	require.NotNil(t, opt)
+	cid := opt.(*OptClientId)
+	require.Equal(t, cid.Cid, duid)
+}
+
+func TestWithServerID(t *testing.T) {
+	duid := Duid{
+		Type:          DUID_LL,
+		HwType:        iana.HwTypeEthernet,
+		LinkLayerAddr: net.HardwareAddr([]byte{0xfa, 0xce, 0xb0, 0x00, 0x00, 0x0c}),
+	}
+	m, err := NewMessage(WithServerID(duid))
+	require.NoError(t, err)
+	opt := m.GetOneOption(OPTION_SERVERID)
+	require.NotNil(t, opt)
+	sid := opt.(*OptServerId)
+	require.Equal(t, sid.Sid, duid)
+}