5 files changed, 142 insertions, 3 deletions
diff --git a/dhcpv4/dhcpv4.go b/dhcpv4/dhcpv4.go
index 94b8351..2931565 100644
--- a/dhcpv4/dhcpv4.go
+++ b/dhcpv4/dhcpv4.go
@@ -612,6 +612,20 @@ func (d *DHCPv4) AddOption(option Option) {
 	}
 }
 
+// UpdateOption updates the existing options with the passed option, adding it
+// at the end if not present already
+func (d *DHCPv4) UpdateOption(option Option) {
+	for idx, opt := range d.options {
+		if opt.Code() == option.Code() {
+			d.options[idx] = option
+			// don't look further
+			return
+		}
+	}
+	// if not found, add it
+	d.AddOption(option)
+}
+
 // MessageType returns the message type, trying to extract it from the
 // OptMessageType option. It returns nil if the message type cannot be extracted
 func (d *DHCPv4) MessageType() *MessageType {
diff --git a/dhcpv4/dhcpv4_test.go b/dhcpv4/dhcpv4_test.go
index 059ae0c..283e728 100644
--- a/dhcpv4/dhcpv4_test.go
+++ b/dhcpv4/dhcpv4_test.go
@@ -347,9 +347,7 @@ func TestGetOption(t *testing.T) {
 
 func TestAddOption(t *testing.T) {
 	d, err := New()
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 
 	hostnameOpt := &OptionGeneric{OptionCode: OptionHostName, Data: []byte("darkstar")}
 	bootFileOpt1 := &OptionGeneric{OptionCode: OptionBootfileName, Data: []byte("boot.img")}
@@ -363,6 +361,23 @@ func TestAddOption(t *testing.T) {
 	require.Equal(t, options[3].Code(), OptionEnd)
 }
 
+func TestUpdateOption(t *testing.T) {
+	d, err := New()
+	require.NoError(t, err)
+	require.Equal(t, 1, len(d.options))
+	require.Equal(t, OptionEnd, d.options[0].Code())
+	// test that it will add the option since it's missing
+	d.UpdateOption(&OptDomainName{DomainName: "slackware.it"})
+	require.Equal(t, 2, len(d.options))
+	require.Equal(t, OptionDomainName, d.options[0].Code())
+	require.Equal(t, OptionEnd, d.options[1].Code())
+	// test that it won't add another option of the same type
+	d.UpdateOption(&OptDomainName{DomainName: "slackware.it"})
+	require.Equal(t, 2, len(d.options))
+	require.Equal(t, OptionDomainName, d.options[0].Code())
+	require.Equal(t, OptionEnd, d.options[1].Code())
+}
+
 func TestStrippedOptions(t *testing.T) {
 	// Normal set of options that terminate with OptionEnd.
 	d, err := New()
diff --git a/dhcpv4/modifiers.go b/dhcpv4/modifiers.go
index 188d91f..033718e 100644
--- a/dhcpv4/modifiers.go
+++ b/dhcpv4/modifiers.go
@@ -2,6 +2,8 @@ package dhcpv4
 
 import (
 	"net"
+
+	"github.com/insomniacslk/dhcp/rfc1035label"
 )
 
 // WithTransactionID sets the Transaction ID for the DHCPv4 packet
@@ -114,3 +116,55 @@ func WithRelay(ip net.IP) Modifier {
 		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
+	}
+}
+
+// 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
+	}
+}
+
+// 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
+	}
+}
+
+// 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
+	}
+}
+
+// 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
+	}
+}
diff --git a/dhcpv4/modifiers_test.go b/dhcpv4/modifiers_test.go
index ce4ff38..f50e40b 100644
--- a/dhcpv4/modifiers_test.go
+++ b/dhcpv4/modifiers_test.go
@@ -136,3 +136,56 @@ func TestWithRelay(t *testing.T) {
 	require.Equal(t, ip, d.GatewayIPAddr())
 	require.Equal(t, uint8(1), d.HopCount())
 }
+
+func TestWithNetmask(t *testing.T) {
+	d := &DHCPv4{}
+	d = WithNetmask(net.IPv4Mask(255, 255, 255, 0))(d)
+	require.Equal(t, 1, len(d.options))
+	require.Equal(t, OptionSubnetMask, d.options[0].Code())
+	osm := d.options[0].(*OptSubnetMask)
+	require.Equal(t, net.IPv4Mask(255, 255, 255, 0), osm.SubnetMask)
+}
+
+func TestWithLeaseTime(t *testing.T) {
+	d := &DHCPv4{}
+	d = WithLeaseTime(uint32(3600))(d)
+	require.Equal(t, 1, len(d.options))
+	require.Equal(t, OptionIPAddressLeaseTime, d.options[0].Code())
+	olt := d.options[0].(*OptIPAddressLeaseTime)
+	require.Equal(t, uint32(3600), olt.LeaseTime)
+}
+
+func TestWithDNS(t *testing.T) {
+	d := &DHCPv4{}
+	d = WithDNS(net.ParseIP("10.0.0.1"), net.ParseIP("10.0.0.2"))(d)
+	require.Equal(t, 1, len(d.options))
+	require.Equal(t, OptionDomainNameServer, d.options[0].Code())
+	olt := d.options[0].(*OptDomainNameServer)
+	require.Equal(t, 2, len(olt.NameServers))
+	require.Equal(t, net.ParseIP("10.0.0.1"), olt.NameServers[0])
+	require.Equal(t, net.ParseIP("10.0.0.2"), olt.NameServers[1])
+	require.NotEqual(t, net.ParseIP("10.0.0.1"), olt.NameServers[1])
+}
+
+func TestWithDomainSearchList(t *testing.T) {
+	d := &DHCPv4{}
+	d = WithDomainSearchList("slackware.it", "dhcp.slackware.it")(d)
+	require.Equal(t, 1, len(d.options))
+	osl := d.options[0].(*OptDomainSearch)
+	require.Equal(t, OptionDNSDomainSearchList, osl.Code())
+	require.NotNil(t, osl.DomainSearch)
+	require.Equal(t, 2, len(osl.DomainSearch.Labels))
+	require.Equal(t, "slackware.it", osl.DomainSearch.Labels[0])
+	require.Equal(t, "dhcp.slackware.it", osl.DomainSearch.Labels[1])
+}
+
+func TestWithRouter(t *testing.T) {
+	d := &DHCPv4{}
+	rtr := net.ParseIP("10.0.0.254")
+	d = WithRouter(rtr)(d)
+	require.Equal(t, 1, len(d.options))
+	ortr := d.options[0].(*OptRouter)
+	require.Equal(t, OptionRouter, ortr.Code())
+	require.Equal(t, 1, len(ortr.Routers))
+	require.Equal(t, rtr, ortr.Routers[0])
+}
diff --git a/dhcpv4/option_domain_search.go b/dhcpv4/option_domain_search.go
index c640c0f..9c24eea 100644
--- a/dhcpv4/option_domain_search.go
+++ b/dhcpv4/option_domain_search.go
@@ -9,6 +9,9 @@ import (
 	"github.com/insomniacslk/dhcp/rfc1035label"
 )
 
+// FIXME rename OptDomainSearch to OptDomainSearchList, and DomainSearch to
+// SearchList, for consistency with the equivalent v6 option
+
 // OptDomainSearch represents an option encapsulating a domain search list.
 type OptDomainSearch struct {
 	DomainSearch *rfc1035label.Labels