fix OptUserClass (#117)

4 files changed, 45 insertions, 9 deletions

diff --git a/dhcpv4/dhcpv4_test.go b/dhcpv4/dhcpv4_test.go
index 7e5f083..217c801 100644
--- a/dhcpv4/dhcpv4_test.go
+++ b/dhcpv4/dhcpv4_test.go
@@ -363,7 +363,7 @@ func TestDHCPv4RequestFromOfferWithModifier(t *testing.T) {
 	require.NoError(t, err)
 	require.NotEqual(t, (*MessageType)(nil), *req.MessageType())
 	require.Equal(t, MessageTypeRequest, *req.MessageType())
-	require.Equal(t, "OptUserClass{userclass=[linuxboot]}", req.options[3].String())
+	require.Equal(t, "User Class Information -> linuxboot", req.options[3].String())
 }
 
 func TestNewReplyFromRequest(t *testing.T) {
@@ -385,7 +385,7 @@ func TestNewReplyFromRequestWithModifier(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, discover.TransactionID(), reply.TransactionID())
 	require.Equal(t, discover.GatewayIPAddr(), reply.GatewayIPAddr())
-	require.Equal(t, "OptUserClass{userclass=[linuxboot]}", reply.options[0].String())
+	require.Equal(t, "User Class Information -> linuxboot", reply.options[0].String())
 }
 
 func TestDHCPv4MessageTypeNil(t *testing.T) {
diff --git a/dhcpv4/modifiers_test.go b/dhcpv4/modifiers_test.go
index 415a4ea..2e249ee 100644
--- a/dhcpv4/modifiers_test.go
+++ b/dhcpv4/modifiers_test.go
@@ -10,5 +10,5 @@ func TestUserClassModifier(t *testing.T) {
 	d, _ := New()
 	userClass := WithUserClass([]byte("linuxboot"))
 	d = userClass(d)
-	require.Equal(t, "OptUserClass{userclass=[linuxboot]}", d.options[0].String())
-}
\ No newline at end of file
+	require.Equal(t, "User Class Information -> linuxboot", d.options[0].String())
+}
diff --git a/dhcpv4/option_userclass.go b/dhcpv4/option_userclass.go
index fc088fa..d6ddabc 100644
--- a/dhcpv4/option_userclass.go
+++ b/dhcpv4/option_userclass.go
@@ -12,6 +12,7 @@ import (
 // OptUserClass represents an option encapsulating User Classes.
 type OptUserClass struct {
 	UserClasses [][]byte
+	Rfc3004 bool
 }
 
 // Code returns the option code
@@ -22,6 +23,9 @@ func (op *OptUserClass) Code() OptionCode {
 // ToBytes serializes the option and returns it as a sequence of bytes
 func (op *OptUserClass) ToBytes() []byte {
 	buf := []byte{byte(op.Code()), byte(op.Length())}
+	if !op.Rfc3004 {
+		return append(buf, op.UserClasses[0]...)
+	}
 	for _, uc := range op.UserClasses {
 		buf = append(buf, byte(len(uc)))
 		buf = append(buf, uc...)
@@ -32,6 +36,9 @@ func (op *OptUserClass) ToBytes() []byte {
 // Length returns the option length
 func (op *OptUserClass) Length() int {
 	ret := 0
+	if !op.Rfc3004 {
+		return len(op.UserClasses[0])
+	}
 	for _, uc := range op.UserClasses {
 		ret += 1 + len(uc)
 	}
@@ -40,10 +47,14 @@ func (op *OptUserClass) Length() int {
 
 func (op *OptUserClass) String() string {
 	ucStrings := make([]string, 0, len(op.UserClasses))
-	for _, uc := range op.UserClasses {
-		ucStrings = append(ucStrings, string(uc))
+	if !op.Rfc3004 {
+		ucStrings = append(ucStrings, string(op.UserClasses[0]))
+	} else {
+		for _, uc := range op.UserClasses {
+			ucStrings = append(ucStrings, string(uc))
+		}
 	}
-	return fmt.Sprintf("OptUserClass{userclass=[%s]}", strings.Join(ucStrings, ", "))
+	return fmt.Sprintf("User Class Information -> %v", strings.Join(ucStrings, ", "))
 }
 
 // ParseOptUserClass returns a new OptUserClass from a byte stream or
@@ -51,7 +62,7 @@ func (op *OptUserClass) String() string {
 func ParseOptUserClass(data []byte) (*OptUserClass, error) {
 	opt := OptUserClass{}
 
-	if len(data) < 4 {
+	if len(data) < 3 {
 		return nil, ErrShortByteStream
 	}
 	code := OptionCode(data[0])
@@ -82,7 +93,7 @@ func ParseOptUserClass(data []byte) (*OptUserClass, error) {
 		opt.UserClasses = append(opt.UserClasses, data[:totalLength])
 		return &opt, nil
 	}
-
+	opt.Rfc3004 = true
 	for i := 0; i < totalLength; {
 		ucLen := int(data[i])
 		if ucLen == 0 {
diff --git a/dhcpv4/option_userclass_test.go b/dhcpv4/option_userclass_test.go
index 02b6f94..f6039df 100644
--- a/dhcpv4/option_userclass_test.go
+++ b/dhcpv4/option_userclass_test.go
@@ -9,6 +9,7 @@ import (
 func TestOptUserClassToBytes(t *testing.T) {
 	opt := OptUserClass{
 		UserClasses: [][]byte{[]byte("linuxboot")},
+		Rfc3004: true,
 	}
 	data := opt.ToBytes()
 	expected := []byte{
@@ -19,6 +20,19 @@ func TestOptUserClassToBytes(t *testing.T) {
 	require.Equal(t, expected, data)
 }
 
+func TestOptUserClassMicrosoftToBytes(t *testing.T) {
+	opt := OptUserClass{
+		UserClasses: [][]byte{[]byte("linuxboot")},
+	}
+	data := opt.ToBytes()
+	expected := []byte{
+		77, // OptionUserClass
+		9,  // length
+		'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+	}
+	require.Equal(t, expected, data)
+}
+
 func TestParseOptUserClassMultiple(t *testing.T) {
 	expected := []byte{
 		77, 15,
@@ -48,6 +62,16 @@ func TestParseOptUserClassMicrosoft(t *testing.T) {
 	require.Equal(t, []byte("linuxboot"), opt.UserClasses[0])
 }
 
+func TestParseOptUserClassMicrosoftShort(t *testing.T) {
+	expected := []byte{
+		77, 1, 'l',
+	}
+	opt, err := ParseOptUserClass(expected)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(opt.UserClasses))
+	require.Equal(t, []byte("l"), opt.UserClasses[0])
+}
+
 func TestParseOptUserClassMicrosoftLongerThanLength(t *testing.T) {
 	expected := []byte{
 		77, 9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't', 'X',
@@ -74,6 +98,7 @@ func TestOptUserClassToBytesMultiple(t *testing.T) {
 			[]byte("linuxboot"),
 			[]byte("test"),
 		},
+		Rfc3004: true,
 	}
 	data := opt.ToBytes()
 	expected := []byte{