DHCPv4 relay modifier (#149)

2 files changed, 28 insertions, 0 deletions

diff --git a/dhcpv4/modifiers.go b/dhcpv4/modifiers.go
index 2a9c35b..01d63ac 100644
--- a/dhcpv4/modifiers.go
+++ b/dhcpv4/modifiers.go
@@ -1,5 +1,9 @@
 package dhcpv4
 
+import (
+	"net"
+)
+
 // 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
@@ -63,3 +67,14 @@ func WithRequestedOptions(optionCodes ...OptionCode) Modifier {
 		return d
 	}
 }
+
+// WithRelay adds parameters required for DHCPv4 to be relayed by the relay
+// server with given ip
+func WithRelay(ip net.IP) Modifier {
+	return func(d *DHCPv4) *DHCPv4 {
+		d.SetUnicast()
+		d.SetGatewayIPAddr(ip)
+		d.SetHopCount(1)
+		return d
+	}
+}
diff --git a/dhcpv4/modifiers_test.go b/dhcpv4/modifiers_test.go
index c40088c..c20558e 100644
--- a/dhcpv4/modifiers_test.go
+++ b/dhcpv4/modifiers_test.go
@@ -1,6 +1,7 @@
 package dhcpv4
 
 import (
+	"net"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -86,3 +87,15 @@ func TestWithRequestedOptions(t *testing.T) {
 	opts = o.(*OptParameterRequestList)
 	require.ElementsMatch(t, opts.RequestedOpts, []OptionCode{OptionFQDN, OptionHostName})
 }
+
+func TestWithRelay(t *testing.T) {
+	d, err := New()
+	require.NoError(t, err)
+	ip := net.ParseIP("10.0.0.1")
+	require.NotNil(t, ip)
+	d = WithRelay(ip)(d)
+	require.NotNil(t, d)
+	require.True(t, d.IsUnicast(), "expected unicast")
+	require.Equal(t, ip, d.GatewayIPAddr())
+	require.Equal(t, uint8(1), d.HopCount())
+}