[dhcpv4] simplify userclass handling (#249)

11 files changed, 162 insertions, 208 deletions

diff --git a/dhcpv4/dhcpv4.go b/dhcpv4/dhcpv4.go
index 1e6c26e..0373764 100644
--- a/dhcpv4/dhcpv4.go
+++ b/dhcpv4/dhcpv4.go
@@ -666,16 +666,16 @@ func (d *DHCPv4) SubnetMask() net.IPMask {
 // UserClass returns the user class if present.
 //
 // The user class information option is defined by RFC 3004.
-func (d *DHCPv4) UserClass() *UserClass {
+func (d *DHCPv4) UserClass() []string {
 	v := d.Options.Get(OptionUserClassInformation)
 	if v == nil {
 		return nil
 	}
-	var uc UserClass
+	var uc Strings
 	if err := uc.FromBytes(v); err != nil {
-		return nil
+		return []string{GetString(OptionUserClassInformation, d.Options)}
 	}
-	return &uc
+	return uc
 }
 
 // VIVC returns the vendor-identifying vendor class option if present.
diff --git a/dhcpv4/dhcpv4_test.go b/dhcpv4/dhcpv4_test.go
index a3e2a8a..5fb5934 100644
--- a/dhcpv4/dhcpv4_test.go
+++ b/dhcpv4/dhcpv4_test.go
@@ -237,7 +237,7 @@ func TestDHCPv4NewRequestFromOfferWithModifier(t *testing.T) {
 	require.NoError(t, err)
 	offer.UpdateOption(OptMessageType(MessageTypeOffer))
 	offer.UpdateOption(OptServerIdentifier(net.IPv4(192, 168, 0, 1)))
-	userClass := WithUserClass([]byte("linuxboot"), false)
+	userClass := WithUserClass("linuxboot", false)
 	req, err := NewRequestFromOffer(offer, userClass)
 	require.NoError(t, err)
 	require.Equal(t, MessageTypeRequest, req.MessageType())
@@ -257,7 +257,7 @@ func TestNewReplyFromRequestWithModifier(t *testing.T) {
 	discover, err := New()
 	require.NoError(t, err)
 	discover.GatewayIPAddr = net.IPv4(192, 168, 0, 1)
-	userClass := WithUserClass([]byte("linuxboot"), false)
+	userClass := WithUserClass("linuxboot", false)
 	reply, err := NewReplyFromRequest(discover, userClass)
 	require.NoError(t, err)
 	require.Equal(t, discover.TransactionID, reply.TransactionID)
diff --git a/dhcpv4/modifiers.go b/dhcpv4/modifiers.go
index 41ad4f4..1659b21 100644
--- a/dhcpv4/modifiers.go
+++ b/dhcpv4/modifiers.go
@@ -88,11 +88,11 @@ func WithOption(opt Option) Modifier {
 // WithUserClass adds a user class option to the packet.
 // The rfc parameter allows you to specify if the userclass should be
 // rfc compliant or not. More details in issue #113
-func WithUserClass(uc []byte, rfc bool) Modifier {
+func WithUserClass(uc string, rfc bool) Modifier {
 	// TODO let the user specify multiple user classes
 	return func(d *DHCPv4) {
 		if rfc {
-			d.UpdateOption(OptRFC3004UserClass([][]byte{uc}))
+			d.UpdateOption(OptRFC3004UserClass([]string{uc}))
 		} else {
 			d.UpdateOption(OptUserClass(uc))
 		}
diff --git a/dhcpv4/modifiers_test.go b/dhcpv4/modifiers_test.go
index 9b6c8ca..3617deb 100644
--- a/dhcpv4/modifiers_test.go
+++ b/dhcpv4/modifiers_test.go
@@ -45,7 +45,7 @@ func TestWithOptionModifier(t *testing.T) {
 }
 
 func TestUserClassModifier(t *testing.T) {
-	d, err := New(WithUserClass([]byte("linuxboot"), false))
+	d, err := New(WithUserClass("linuxboot", false))
 	require.NoError(t, err)
 
 	expected := []byte{
@@ -55,7 +55,7 @@ func TestUserClassModifier(t *testing.T) {
 }
 
 func TestUserClassModifierRFC(t *testing.T) {
-	d, err := New(WithUserClass([]byte("linuxboot"), true))
+	d, err := New(WithUserClass("linuxboot", true))
 	require.NoError(t, err)
 
 	expected := []byte{
diff --git a/dhcpv4/option_string.go b/dhcpv4/option_string.go
index b505575..289319b 100644
--- a/dhcpv4/option_string.go
+++ b/dhcpv4/option_string.go
@@ -72,3 +72,8 @@ func OptTFTPServerName(name string) Option {
 func OptClassIdentifier(name string) Option {
 	return Option{Code: OptionClassIdentifier, Value: String(name)}
 }
+
+// OptUserClass returns a new DHCPv4 User Class option.
+func OptUserClass(name string) Option {
+	return Option{Code: OptionUserClassInformation, Value: String(name)}
+}
diff --git a/dhcpv4/option_string_test.go b/dhcpv4/option_string_test.go
index 4dbb77d..768d0fe 100644
--- a/dhcpv4/option_string_test.go
+++ b/dhcpv4/option_string_test.go
@@ -95,3 +95,21 @@ func TestParseOptClassIdentifier(t *testing.T) {
 	m, _ = New()
 	require.Equal(t, "", m.ClassIdentifier())
 }
+
+func TestOptUserClass(t *testing.T) {
+	o := OptUserClass("linuxboot")
+	require.Equal(t, OptionUserClassInformation, o.Code, "Code")
+	expected := []byte{
+		'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+	}
+	require.Equal(t, expected, o.Value.ToBytes(), "ToBytes")
+	require.Equal(t, "User Class Information: linuxboot", o.String())
+}
+
+func TestParseOptUserClass(t *testing.T) {
+	m, _ := New(WithUserClass("linuxboot", false))
+	require.Equal(t, []string{"linuxboot"}, m.UserClass())
+
+	m, _ = New()
+	require.Equal(t, 0, len(m.UserClass()))
+}
diff --git a/dhcpv4/option_strings.go b/dhcpv4/option_strings.go
new file mode 100644
index 0000000..dafa9f2
--- /dev/null
+++ b/dhcpv4/option_strings.go
@@ -0,0 +1,55 @@
+package dhcpv4
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/u-root/u-root/pkg/uio"
+)
+
+// Strings represents an option encapsulating a list of strings in IPv4 DHCP as
+// specified in RFC 3004
+//
+// Strings implements the OptionValue type.
+type Strings []string
+
+// FromBytes parses Strings from a DHCP packet as specified by RFC 3004.
+func (o *Strings) FromBytes(data []byte) error {
+	buf := uio.NewBigEndianBuffer(data)
+	if buf.Len() == 0 {
+		return fmt.Errorf("Strings DHCP option must always list at least one String")
+	}
+
+	*o = make(Strings, 0)
+	for buf.Has(1) {
+		ucLen := buf.Read8()
+		if ucLen == 0 {
+			return fmt.Errorf("DHCP Strings must have length greater than 0")
+		}
+		*o = append(*o, string(buf.CopyN(int(ucLen))))
+	}
+	return buf.FinError()
+}
+
+// ToBytes marshals Strings to a DHCP packet as specified by RFC 3004.
+func (o Strings) ToBytes() []byte {
+	buf := uio.NewBigEndianBuffer(nil)
+	for _, uc := range o {
+		buf.Write8(uint8(len(uc)))
+		buf.WriteBytes([]byte(uc))
+	}
+	return buf.Data()
+}
+
+// String returns a human-readable representation of a list of Strings.
+func (o Strings) String() string {
+	return strings.Join(o, ", ")
+}
+
+// OptRFC3004UserClass returns a new user class option according to RFC 3004.
+func OptRFC3004UserClass(v []string) Option {
+	return Option{
+		Code:  OptionUserClassInformation,
+		Value: Strings(v),
+	}
+}
diff --git a/dhcpv4/option_strings_test.go b/dhcpv4/option_strings_test.go
new file mode 100644
index 0000000..cd5d712
--- /dev/null
+++ b/dhcpv4/option_strings_test.go
@@ -0,0 +1,68 @@
+package dhcpv4
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseStringsMultiple(t *testing.T) {
+	var opt Strings
+	expected := []byte{
+		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+		4, 't', 'e', 's', 't',
+	}
+	err := opt.FromBytes(expected)
+	require.NoError(t, err)
+	require.Equal(t, len(opt), 2)
+	require.Equal(t, "linuxboot", opt[0])
+	require.Equal(t, "test", opt[1])
+}
+
+func TestParseStringsNone(t *testing.T) {
+	var opt Strings
+	expected := []byte{}
+	err := opt.FromBytes(expected)
+	require.Error(t, err)
+}
+
+func TestParseStrings(t *testing.T) {
+	var opt Strings
+	expected := []byte{
+		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+	}
+	err := opt.FromBytes(expected)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(opt))
+	require.Equal(t, "linuxboot", opt[0])
+}
+
+func TestParseStringsZeroLength(t *testing.T) {
+	var opt Strings
+	err := opt.FromBytes([]byte{0, 0})
+	require.Error(t, err)
+}
+
+func TestOptRFC3004UserClass(t *testing.T) {
+	opt := OptRFC3004UserClass(Strings([]string{"linuxboot"}))
+	data := opt.Value.ToBytes()
+	expected := []byte{
+		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+	}
+	require.Equal(t, expected, data)
+}
+
+func TestOptRFC3004UserClassMultiple(t *testing.T) {
+	opt := OptRFC3004UserClass(
+		[]string{
+			"linuxboot",
+			"test",
+		},
+	)
+	data := opt.Value.ToBytes()
+	expected := []byte{
+		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
+		4, 't', 'e', 's', 't',
+	}
+	require.Equal(t, expected, data)
+}
diff --git a/dhcpv4/option_userclass.go b/dhcpv4/option_userclass.go
deleted file mode 100644
index 69ceeba..0000000
--- a/dhcpv4/option_userclass.go
+++ /dev/null
@@ -1,98 +0,0 @@
-package dhcpv4
-
-import (
-	"errors"
-	"fmt"
-	"strings"
-
-	"github.com/u-root/u-root/pkg/uio"
-)
-
-// UserClass implements the user class option described by RFC 3004.
-type UserClass struct {
-	UserClasses [][]byte
-	RFC3004     bool
-}
-
-// OptUserClass returns a new user class option.
-func OptUserClass(v []byte) Option {
-	return Option{
-		Code: OptionUserClassInformation,
-		Value: &UserClass{
-			UserClasses: [][]byte{v},
-			RFC3004:     false,
-		},
-	}
-}
-
-// OptRFC3004UserClass returns a new user class option according to RFC 3004.
-func OptRFC3004UserClass(v [][]byte) Option {
-	return Option{
-		Code: OptionUserClassInformation,
-		Value: &UserClass{
-			UserClasses: v,
-			RFC3004:     true,
-		},
-	}
-}
-
-// ToBytes serializes the option and returns it as a sequence of bytes
-func (op *UserClass) ToBytes() []byte {
-	buf := uio.NewBigEndianBuffer(nil)
-	if !op.RFC3004 {
-		buf.WriteBytes(op.UserClasses[0])
-	} else {
-		for _, uc := range op.UserClasses {
-			buf.Write8(uint8(len(uc)))
-			buf.WriteBytes(uc)
-		}
-	}
-	return buf.Data()
-}
-
-// String returns a human-readable user class.
-func (op *UserClass) String() string {
-	ucStrings := make([]string, 0, len(op.UserClasses))
-	if !op.RFC3004 {
-		ucStrings = append(ucStrings, string(op.UserClasses[0]))
-	} else {
-		for _, uc := range op.UserClasses {
-			ucStrings = append(ucStrings, string(uc))
-		}
-	}
-	return strings.Join(ucStrings, ", ")
-}
-
-// FromBytes parses data into op.
-func (op *UserClass) FromBytes(data []byte) error {
-	buf := uio.NewBigEndianBuffer(data)
-
-	// Check if option is Microsoft style instead of RFC compliant, issue #113
-
-	// User-class options are, according to RFC3004, supposed to contain a set
-	// of strings each with length UC_Len_i. Here we check that this is so,
-	// by seeing if all the UC_Len_i lengths are consistent with the overall
-	// option length. If the lengths don't add up, we assume that the option
-	// is a single string and non RFC3004 compliant
-	var counting int
-	for counting < buf.Len() {
-		// UC_Len_i does not include itself so add 1
-		counting += int(data[counting]) + 1
-	}
-	if counting != buf.Len() {
-		op.UserClasses = append(op.UserClasses, data)
-		return nil
-	}
-	op.RFC3004 = true
-	for buf.Has(1) {
-		ucLen := buf.Read8()
-		if ucLen == 0 {
-			return fmt.Errorf("DHCP user class must have length greater than 0")
-		}
-		op.UserClasses = append(op.UserClasses, buf.CopyN(int(ucLen)))
-	}
-	if len(op.UserClasses) == 0 {
-		return errors.New("ParseOptUserClass: at least one user class is required")
-	}
-	return buf.FinError()
-}
diff --git a/dhcpv4/option_userclass_test.go b/dhcpv4/option_userclass_test.go
deleted file mode 100644
index 149fb92..0000000
--- a/dhcpv4/option_userclass_test.go
+++ /dev/null
@@ -1,99 +0,0 @@
-package dhcpv4
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/require"
-)
-
-func TestOptUserClassToBytes(t *testing.T) {
-	opt := OptRFC3004UserClass([][]byte{[]byte("linuxboot")})
-	data := opt.Value.ToBytes()
-	expected := []byte{
-		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
-	}
-	require.Equal(t, expected, data)
-}
-
-func TestOptUserClassMicrosoftToBytes(t *testing.T) {
-	opt := OptUserClass([]byte("linuxboot"))
-	data := opt.Value.ToBytes()
-	expected := []byte{
-		'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
-	}
-	require.Equal(t, expected, data)
-}
-
-func TestParseOptUserClassMultiple(t *testing.T) {
-	var opt UserClass
-	expected := []byte{
-		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
-		4, 't', 'e', 's', 't',
-	}
-	err := opt.FromBytes(expected)
-	require.NoError(t, err)
-	require.Equal(t, len(opt.UserClasses), 2)
-	require.Equal(t, []byte("linuxboot"), opt.UserClasses[0])
-	require.Equal(t, []byte("test"), opt.UserClasses[1])
-}
-
-func TestParseOptUserClassNone(t *testing.T) {
-	var opt UserClass
-	expected := []byte{}
-	err := opt.FromBytes(expected)
-	require.Error(t, err)
-}
-
-func TestParseOptUserClassMicrosoft(t *testing.T) {
-	var opt UserClass
-	expected := []byte{
-		'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
-	}
-	err := opt.FromBytes(expected)
-	require.NoError(t, err)
-	require.Equal(t, 1, len(opt.UserClasses))
-	require.Equal(t, []byte("linuxboot"), opt.UserClasses[0])
-}
-
-func TestParseOptUserClassMicrosoftShort(t *testing.T) {
-	var opt UserClass
-	expected := []byte{
-		'l',
-	}
-	err := opt.FromBytes(expected)
-	require.NoError(t, err)
-	require.Equal(t, 1, len(opt.UserClasses))
-	require.Equal(t, []byte("l"), opt.UserClasses[0])
-}
-
-func TestParseOptUserClass(t *testing.T) {
-	var opt UserClass
-	expected := []byte{
-		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
-	}
-	err := opt.FromBytes(expected)
-	require.NoError(t, err)
-	require.Equal(t, 1, len(opt.UserClasses))
-	require.Equal(t, []byte("linuxboot"), opt.UserClasses[0])
-}
-
-func TestOptUserClassToBytesMultiple(t *testing.T) {
-	opt := OptRFC3004UserClass(
-		[][]byte{
-			[]byte("linuxboot"),
-			[]byte("test"),
-		},
-	)
-	data := opt.Value.ToBytes()
-	expected := []byte{
-		9, 'l', 'i', 'n', 'u', 'x', 'b', 'o', 'o', 't',
-		4, 't', 'e', 's', 't',
-	}
-	require.Equal(t, expected, data)
-}
-
-func TestParseOptUserClassZeroLength(t *testing.T) {
-	var opt UserClass
-	err := opt.FromBytes([]byte{0, 0})
-	require.Error(t, err)
-}
diff --git a/dhcpv4/options.go b/dhcpv4/options.go
index 4c70743..68ce665 100644
--- a/dhcpv4/options.go
+++ b/dhcpv4/options.go
@@ -330,7 +330,12 @@ func getOption(code OptionCode, data []byte, vendorDecoder OptionDecoder) fmt.St
 		d = &u
 
 	case OptionUserClassInformation:
-		d = &UserClass{}
+		var s Strings
+		d = &s
+		if s.FromBytes(data) != nil {
+			var s String
+			d = &s
+		}
 
 	case OptionVendorIdentifyingVendorClass:
 		d = &VIVCIdentifiers{}