dhcpv4: add option code list type for simpler modifiers.

4 files changed, 59 insertions, 95 deletions

diff --git a/dhcpv4/modifiers.go b/dhcpv4/modifiers.go
index e7f6576..cd80bc9 100644
--- a/dhcpv4/modifiers.go
+++ b/dhcpv4/modifiers.go
@@ -47,60 +47,26 @@ func WithOption(opt Option) Modifier {
 // rfc compliant or not. More details in issue #113
 func WithUserClass(uc []byte, rfc bool) Modifier {
 	// TODO let the user specify multiple user classes
-	return func(d *DHCPv4) *DHCPv4 {
-		ouc := OptUserClass{
-			UserClasses: [][]byte{uc},
-			Rfc3004:     rfc,
-		}
-		d.UpdateOption(&ouc)
-		return d
-	}
+	return WithOption(&OptUserClass{
+		UserClasses: [][]byte{uc},
+		Rfc3004:     rfc,
+	})
 }
 
 // WithNetboot adds bootfile URL and bootfile param options to a DHCPv4 packet.
 func WithNetboot(d *DHCPv4) *DHCPv4 {
-	params := d.GetOneOption(OptionParameterRequestList)
-
-	var (
-		OptParams                 *OptParameterRequestList
-		foundOptionTFTPServerName bool
-		foundOptionBootfileName   bool
-	)
-	if params != nil {
-		OptParams = params.(*OptParameterRequestList)
-		for _, option := range OptParams.RequestedOpts {
-			if option == OptionTFTPServerName {
-				foundOptionTFTPServerName = true
-			} else if option == OptionBootfileName {
-				foundOptionBootfileName = true
-			}
-		}
-		if !foundOptionTFTPServerName {
-			OptParams.RequestedOpts = append(OptParams.RequestedOpts, OptionTFTPServerName)
-		}
-		if !foundOptionBootfileName {
-			OptParams.RequestedOpts = append(OptParams.RequestedOpts, OptionBootfileName)
-		}
-	} else {
-		OptParams = &OptParameterRequestList{
-			RequestedOpts: []OptionCode{OptionTFTPServerName, OptionBootfileName},
-		}
-		d.UpdateOption(OptParams)
-	}
-	return d
+	return WithRequestedOptions(OptionTFTPServerName, OptionBootfileName)(d)
 }
 
-// WithRequestedOptions adds requested options to the packet
+// WithRequestedOptions adds requested options to the packet.
 func WithRequestedOptions(optionCodes ...OptionCode) Modifier {
 	return func(d *DHCPv4) *DHCPv4 {
 		params := d.GetOneOption(OptionParameterRequestList)
 		if params == nil {
-			params = &OptParameterRequestList{}
-			d.UpdateOption(params)
-		}
-		opts := params.(*OptParameterRequestList)
-		for _, optionCode := range optionCodes {
-			opts.RequestedOpts = append(opts.RequestedOpts, optionCode)
+			d.UpdateOption(&OptParameterRequestList{OptionCodeList(optionCodes)})
+		} else {
+			opts := params.(*OptParameterRequestList)
+			opts.RequestedOpts.Add(optionCodes...)
 		}
 		return d
 	}
@@ -112,59 +78,34 @@ func WithRelay(ip net.IP) Modifier {
 	return func(d *DHCPv4) *DHCPv4 {
 		d.SetUnicast()
 		d.GatewayIPAddr = ip
-		d.HopCount = 1
+		d.HopCount += 1
 		return d
 	}
 }
 
 // WithNetmask adds or updates an OptSubnetMask
 func WithNetmask(mask net.IPMask) Modifier {
-	return func(d *DHCPv4) *DHCPv4 {
-		osm := OptSubnetMask{
-			SubnetMask: mask,
-		}
-		d.UpdateOption(&osm)
-		return d
-	}
+	return WithOption(&OptSubnetMask{SubnetMask: mask})
 }
 
 // WithLeaseTime adds or updates an OptIPAddressLeaseTime
 func WithLeaseTime(leaseTime uint32) Modifier {
-	return func(d *DHCPv4) *DHCPv4 {
-		olt := OptIPAddressLeaseTime{
-			LeaseTime: leaseTime,
-		}
-		d.UpdateOption(&olt)
-		return d
-	}
+	return WithOption(&OptIPAddressLeaseTime{LeaseTime: leaseTime})
 }
 
 // WithDNS adds or updates an OptionDomainNameServer
 func WithDNS(dnses ...net.IP) Modifier {
-	return func(d *DHCPv4) *DHCPv4 {
-		odns := OptDomainNameServer{NameServers: dnses}
-		d.UpdateOption(&odns)
-		return d
-	}
+	return WithOption(&OptDomainNameServer{NameServers: dnses})
 }
 
 // WithDomainSearchList adds or updates an OptionDomainSearch
 func WithDomainSearchList(searchList ...string) Modifier {
-	return func(d *DHCPv4) *DHCPv4 {
-		labels := rfc1035label.Labels{
-			Labels: searchList,
-		}
-		odsl := OptDomainSearch{DomainSearch: &labels}
-		d.UpdateOption(&odsl)
-		return d
-	}
+	return WithOption(&OptDomainSearch{DomainSearch: &rfc1035label.Labels{
+		Labels: searchList,
+	}})
 }
 
 // WithRouter adds or updates an OptionRouter
 func WithRouter(routers ...net.IP) Modifier {
-	return func(d *DHCPv4) *DHCPv4 {
-		ortr := OptRouter{Routers: routers}
-		d.UpdateOption(&ortr)
-		return d
-	}
+	return WithOption(&OptRouter{Routers: routers})
 }
diff --git a/dhcpv4/modifiers_test.go b/dhcpv4/modifiers_test.go
index e653817..d9bb3c7 100644
--- a/dhcpv4/modifiers_test.go
+++ b/dhcpv4/modifiers_test.go
@@ -69,7 +69,7 @@ func TestUserClassModifierRFC(t *testing.T) {
 func TestWithNetboot(t *testing.T) {
 	d, _ := New()
 	d = WithNetboot(d)
-	require.Equal(t, "Parameter Request List -> [TFTP Server Name, Bootfile Name]", d.Options[0].String())
+	require.Equal(t, "Parameter Request List -> TFTP Server Name, Bootfile Name", d.Options[0].String())
 }
 
 func TestWithNetbootExistingTFTP(t *testing.T) {
@@ -79,7 +79,7 @@ func TestWithNetbootExistingTFTP(t *testing.T) {
 	}
 	d.UpdateOption(OptParams)
 	d = WithNetboot(d)
-	require.Equal(t, "Parameter Request List -> [TFTP Server Name, Bootfile Name]", d.Options[0].String())
+	require.Equal(t, "Parameter Request List -> TFTP Server Name, Bootfile Name", d.Options[0].String())
 }
 
 func TestWithNetbootExistingBootfileName(t *testing.T) {
@@ -89,7 +89,7 @@ func TestWithNetbootExistingBootfileName(t *testing.T) {
 	}
 	d.UpdateOption(OptParams)
 	d = WithNetboot(d)
-	require.Equal(t, "Parameter Request List -> [Bootfile Name, TFTP Server Name]", d.Options[0].String())
+	require.Equal(t, "Parameter Request List -> Bootfile Name, TFTP Server Name", d.Options[0].String())
 }
 
 func TestWithNetbootExistingBoth(t *testing.T) {
@@ -99,7 +99,7 @@ func TestWithNetbootExistingBoth(t *testing.T) {
 	}
 	d.UpdateOption(OptParams)
 	d = WithNetboot(d)
-	require.Equal(t, "Parameter Request List -> [Bootfile Name, TFTP Server Name]", d.Options[0].String())
+	require.Equal(t, "Parameter Request List -> Bootfile Name, TFTP Server Name", d.Options[0].String())
 }
 
 func TestWithRequestedOptions(t *testing.T) {
diff --git a/dhcpv4/option_parameter_request_list.go b/dhcpv4/option_parameter_request_list.go
index d82f1ea..750c957 100644
--- a/dhcpv4/option_parameter_request_list.go
+++ b/dhcpv4/option_parameter_request_list.go
@@ -7,18 +7,49 @@ import (
 	"github.com/u-root/u-root/pkg/uio"
 )
 
+// OptionCodeList is a list of DHCP option codes.
+type OptionCodeList []OptionCode
+
+// Has returns whether c is in the list.
+func (ol OptionCodeList) Has(c OptionCode) bool {
+	for _, code := range ol {
+		if code == c {
+			return true
+		}
+	}
+	return false
+}
+
+// Add adds option codes in cs to ol.
+func (ol *OptionCodeList) Add(cs ...OptionCode) {
+	for _, c := range cs {
+		if !ol.Has(c) {
+			*ol = append(*ol, c)
+		}
+	}
+}
+
+// String returns a human-readable string for the option names.
+func (ol OptionCodeList) String() string {
+	var names []string
+	for _, code := range ol {
+		names = append(names, code.String())
+	}
+	return strings.Join(names, ", ")
+}
+
 // OptParameterRequestList implements the parameter request list option
 // described by RFC 2132, Section 9.8.
 type OptParameterRequestList struct {
-	RequestedOpts []OptionCode
+	RequestedOpts OptionCodeList
 }
 
 // ParseOptParameterRequestList returns a new OptParameterRequestList from a
 // byte stream, or error if any.
 func ParseOptParameterRequestList(data []byte) (*OptParameterRequestList, error) {
 	buf := uio.NewBigEndianBuffer(data)
-	requestedOpts := make([]OptionCode, 0, buf.Len())
-	for buf.Len() > 0 {
+	requestedOpts := make(OptionCodeList, 0, buf.Len())
+	for buf.Has(1) {
 		requestedOpts = append(requestedOpts, optionCode(buf.Read8()))
 	}
 	return &OptParameterRequestList{RequestedOpts: requestedOpts}, buf.Error()
@@ -40,13 +71,5 @@ func (o *OptParameterRequestList) ToBytes() []byte {
 
 // String returns a human-readable string for this option.
 func (o *OptParameterRequestList) String() string {
-	var optNames []string
-	for _, ro := range o.RequestedOpts {
-		name := ro.String()
-		if name == "Unknown" {
-			name += fmt.Sprintf("%s (%v)", name, ro)
-		}
-		optNames = append(optNames, name)
-	}
-	return fmt.Sprintf("Parameter Request List -> [%v]", strings.Join(optNames, ", "))
+	return fmt.Sprintf("Parameter Request List -> %s", o.RequestedOpts)
 }
diff --git a/dhcpv4/option_parameter_request_list_test.go b/dhcpv4/option_parameter_request_list_test.go
index 0590da8..a09aaad 100644
--- a/dhcpv4/option_parameter_request_list_test.go
+++ b/dhcpv4/option_parameter_request_list_test.go
@@ -14,7 +14,7 @@ func TestOptParameterRequestListInterfaceMethods(t *testing.T) {
 	expectedBytes := []byte{67, 5}
 	require.Equal(t, expectedBytes, o.ToBytes(), "ToBytes")
 
-	expectedString := "Parameter Request List -> [Bootfile Name, Name Server]"
+	expectedString := "Parameter Request List -> Bootfile Name, Name Server"
 	require.Equal(t, expectedString, o.String(), "String")
 }
 
@@ -25,6 +25,6 @@ func TestParseOptParameterRequestList(t *testing.T) {
 	)
 	o, err = ParseOptParameterRequestList([]byte{67, 5})
 	require.NoError(t, err)
-	expectedOpts := []OptionCode{OptionBootfileName, OptionNameServer}
+	expectedOpts := OptionCodeList{OptionBootfileName, OptionNameServer}
 	require.Equal(t, expectedOpts, o.RequestedOpts)
 }