dhcpv6: option userclass: supporting multiple user classes (#44)

dhcpv6: option userclass: supporting multiple user classes

2 files changed, 35 insertions, 12 deletions

diff --git a/dhcpv6/modifiers.go b/dhcpv6/modifiers.go
index 75ebfb0..08745bc 100644
--- a/dhcpv6/modifiers.go
+++ b/dhcpv6/modifiers.go
@@ -26,8 +26,9 @@ 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
 	return func(d DHCPv6) DHCPv6 {
-		ouc := OptUserClass{UserClass: uc}
+		ouc := OptUserClass{UserClasses: [][]byte{uc}}
 		d.AddOption(&ouc)
 		return d
 	}
diff --git a/dhcpv6/option_userclass.go b/dhcpv6/option_userclass.go
index 68dc298..0bae97f 100644
--- a/dhcpv6/option_userclass.go
+++ b/dhcpv6/option_userclass.go
@@ -5,12 +5,14 @@ package dhcpv6
 
 import (
 	"encoding/binary"
+	"errors"
 	"fmt"
+	"strings"
 )
 
 // OptUserClass represent a DHCPv6 User Class option
 type OptUserClass struct {
-	UserClass []byte
+	UserClasses [][]byte
 }
 
 // Code returns the option code
@@ -20,32 +22,52 @@ func (op *OptUserClass) Code() OptionCode {
 
 // ToBytes serializes the option and returns it as a sequence of bytes
 func (op *OptUserClass) ToBytes() []byte {
-	buf := make([]byte, 6)
+	buf := make([]byte, 4)
 	binary.BigEndian.PutUint16(buf[0:2], uint16(OPTION_USER_CLASS))
 	binary.BigEndian.PutUint16(buf[2:4], uint16(op.Length()))
-	// user-class-data has an internal data length field too..
-	binary.BigEndian.PutUint16(buf[4:6], uint16(len(op.UserClass)))
-	buf = append(buf, op.UserClass...)
+	u16 := make([]byte, 2)
+	for _, uc := range op.UserClasses {
+		binary.BigEndian.PutUint16(u16, uint16(len(uc)))
+		buf = append(buf, u16...)
+		buf = append(buf, uc...)
+	}
 	return buf
 }
 
 // Length returns the option length
 func (op *OptUserClass) Length() int {
-	return 2 + len(op.UserClass)
+	ret := 0
+	for _, uc := range op.UserClasses {
+		ret += 2 + len(uc)
+	}
+	return ret
 }
 
 func (op *OptUserClass) String() string {
-	return fmt.Sprintf("OptUserClass{userclass=%s}", string(op.UserClass))
+	ucStrings := make([]string, 0)
+	for _, uc := range op.UserClasses {
+		ucStrings = append(ucStrings, string(uc))
+	}
+	return fmt.Sprintf("OptUserClass{userclass=[%s]}", strings.Join(ucStrings, ", "))
 }
 
 // ParseOptUserClass builds an OptUserClass structure from a sequence of
 // bytes. The input data does not include option code and length bytes.
 func ParseOptUserClass(data []byte) (*OptUserClass, error) {
 	opt := OptUserClass{}
-	dataLen := int(binary.BigEndian.Uint16(data[:2]))
-	if dataLen != len(data)-2 {
-		return nil, fmt.Errorf("ParseOptUserClass: declared data length does not match actual length: %d != %d", dataLen, len(data)-2)
+	for {
+		if len(data) == 0 {
+			break
+		}
+		if len(data) < 2 {
+			return nil, errors.New("ParseOptUserClass: short data: missing length field")
+		}
+		ucLen := int(binary.BigEndian.Uint16(data[2:]))
+		if len(data) < ucLen+2 {
+			return nil, fmt.Errorf("ParseOptUserClass: short data: less than %d bytes", ucLen+2)
+		}
+		opt.UserClasses = append(opt.UserClasses, data[2:ucLen])
+		data = data[2+ucLen:]
 	}
-	opt.UserClass = append(opt.UserClass, data[2:]...)
 	return &opt, nil
 }