Skip to content

Commit 7314c9b

Browse files
jrgpjrgp
committed
Remove more boilerplate + add more tests
- Remove boilerplate around the common IP option setting logic in request handler for DHCPOFFER/DHCPACK - Add tests to verify the correct IP options get set - Set types to byte for some constants to avoid the need for type casts
1 parent c6cb16d commit 7314c9b

File tree

4 files changed

+81
-30
lines changed

4 files changed

+81
-30
lines changed

‎constants.go

+10-10
Original file line numberDiff line numberDiff line change
@@ -4,22 +4,22 @@ package main
44
// DHCP Op types
55
//
66
const (
7-
BOOT_REQUEST = 1
8-
BOOT_REPLY = 2
7+
BOOT_REQUEST byte = 1
8+
BOOT_REPLY byte = 2
99
)
1010

1111
//
1212
// DHCP Message types
1313
//
1414
const (
15-
DHCPDISCOVER = 1 // Implemented
16-
DHCPOFFER = 2 // Implemented
17-
DHCPREQUEST = 3 // Implemented
18-
DHCPDECLINE = 4
19-
DHCPACK = 5 // Implemented
20-
DHCPNAK = 6 // Implemented
21-
DHCPRELEASE = 7 // Implemented
22-
DHCPINFORM = 8
15+
DHCPDISCOVER byte = 1 // Implemented
16+
DHCPOFFER byte = 2 // Implemented
17+
DHCPREQUEST byte = 3 // Implemented
18+
DHCPDECLINE byte = 4
19+
DHCPACK byte = 5 // Implemented
20+
DHCPNAK byte = 6 // Implemented
21+
DHCPRELEASE byte = 7 // Implemented
22+
DHCPINFORM byte = 8
2323
)
2424

2525
var opNames = map[byte]string{

‎options.go

+47
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"fmt"
1010
"io"
1111
"log"
12+
"net"
1213
)
1314

1415
// Represent a single DHCP option
@@ -56,6 +57,9 @@ func (o *Options) Dump() {
5657
}
5758
}
5859

60+
//
61+
// Abstract away boilerplate for common getting operations
62+
//
5963
func (o *Options) Get(code byte) (Option, bool) {
6064
option, ok := o.data[code]
6165
return option, ok
@@ -70,6 +74,48 @@ func (o *Options) GetByte(code byte) byte {
7074
return 0
7175
}
7276

77+
// Primarily used in testing
78+
func (o *Options) GetFixedV4s(code byte) []FixedV4 {
79+
if option, ok := o.data[code]; ok {
80+
if len(option.Data)%4 == 0 {
81+
ips := make([]FixedV4, len(option.Data)/4)
82+
for i := 0; i < len(ips); i++ {
83+
if ip, err := BytesToFixedV4(option.Data[i*4 : i*4+4]); err == nil {
84+
ips[i] = ip
85+
}
86+
}
87+
return ips
88+
}
89+
}
90+
return nil
91+
}
92+
93+
//
94+
// Abstract away boilerplate for common IP setting operations
95+
//
96+
func (o *Options) SetIPs(code byte, ips ...net.IP) {
97+
if len(ips) == 0 {
98+
return
99+
}
100+
data := make([]byte, 0, 4*len(ips))
101+
for _, ip := range ips {
102+
data = append(data, IpToFixedV4(ip).Bytes()...)
103+
}
104+
o.Set(code, data)
105+
}
106+
107+
func (o *Options) SetFixedV4s(code byte, ips ...FixedV4) {
108+
if len(ips) == 0 {
109+
return
110+
}
111+
data := make([]byte, 0, 4*len(ips))
112+
for _, ip := range ips {
113+
data = append(data, ip.Bytes()...)
114+
}
115+
o.Set(code, data)
116+
}
117+
118+
// Set a single option
73119
func (o *Options) Set(code byte, data []byte) {
74120
option := Option{
75121
Data: data,
@@ -87,6 +133,7 @@ func (o *Options) Set(code byte, data []byte) {
87133
o.data[code] = option
88134
}
89135

136+
// Encode all options, including sentinel, to buf
90137
func (o *Options) Encode(buf *bytes.Buffer) error {
91138
// Need the sentinel value at the end
92139
if len(o.order) > 0 && o.order[len(o.order)-1] != OPTION_SENTINEL {

‎request.go

+4-12
Original file line numberDiff line numberDiff line change
@@ -118,31 +118,23 @@ func (r *RequestHandler) SendLeaseInfo(lease *Lease, op byte) *DHCPMessage {
118118
options.Set(OPTION_MESSAGE_TYPE, []byte{op})
119119

120120
// Netmask option
121-
options.Set(OPTION_SUBNET, IpToFixedV4(r.pool.Netmask).Bytes())
121+
options.SetIPs(OPTION_SUBNET, r.pool.Netmask)
122122

123123
// Router (defgw)
124124
if len(r.pool.Router) > 0 {
125-
bytes := make([]byte, 0, 4*len(r.pool.Router))
126-
for _, ip := range r.pool.Router {
127-
bytes = append(bytes, IpToFixedV4(ip).Bytes()...)
128-
}
129-
options.Set(OPTION_ROUTER, bytes)
125+
options.SetIPs(OPTION_ROUTER, r.pool.Router...)
130126
}
131127

132128
// DNS servers
133129
if len(r.pool.Dns) > 0 {
134-
bytes := make([]byte, 0, 4*len(r.pool.Dns))
135-
for _, ip := range r.pool.Dns {
136-
bytes = append(bytes, IpToFixedV4(ip).Bytes()...)
137-
}
138-
options.Set(OPTION_DNS_SERVER, bytes)
130+
options.SetIPs(OPTION_DNS_SERVER, r.pool.Dns...)
139131
}
140132

141133
// Lease time
142134
options.Set(OPTION_LEASE_TIME, long2bytes(uint32(r.pool.LeaseTime.Seconds())))
143135

144136
// DHCP server
145-
options.Set(OPTION_SERVER_ID, r.pool.MyIp.Bytes())
137+
options.SetFixedV4s(OPTION_SERVER_ID, r.pool.MyIp)
146138

147139
return &DHCPMessage{header, options}
148140
}

‎request_test.go

+20-8
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@ func TestDhcpDiscover(t *testing.T) {
1212
pool.Start = net.ParseIP("10.0.0.10")
1313
pool.End = net.ParseIP("10.0.0.20")
1414
pool.Netmask = net.ParseIP("255.255.255.0")
15+
pool.Router = []net.IP{net.ParseIP("10.0.0.1")}
16+
pool.MyIp = IpToFixedV4(net.ParseIP("10.0.0.254"))
17+
pool.Dns = []net.IP{net.ParseIP("1.1.1.1"), net.ParseIP("1.0.0.1")}
1518

1619
//
1720
// DHCPREQUEST targeting an unknown lease. Should get a NAK back
@@ -28,7 +31,8 @@ func TestDhcpDiscover(t *testing.T) {
2831
handler := NewRequestHandler(message, pool)
2932
response := handler.Handle()
3033

31-
require.Equal(t, byte(DHCPNAK), response.Options.GetByte(OPTION_MESSAGE_TYPE))
34+
require.Equal(t, BOOT_REPLY, response.Header.Op)
35+
require.Equal(t, DHCPNAK, response.Options.GetByte(OPTION_MESSAGE_TYPE))
3236

3337
//
3438
// DISCOVER. Should get back a lease.
@@ -39,14 +43,20 @@ func TestDhcpDiscover(t *testing.T) {
3943

4044
message, err = ParseDhcpMessage(b)
4145
require.Nil(t, err)
42-
require.Equal(t, byte(DHCPDISCOVER), message.Options.GetByte(OPTION_MESSAGE_TYPE))
46+
require.Equal(t, DHCPDISCOVER, message.Options.GetByte(OPTION_MESSAGE_TYPE))
4347

4448
handler = NewRequestHandler(message, pool)
4549
response = handler.Handle()
4650

47-
require.Equal(t, byte(DHCPOFFER), response.Options.GetByte(OPTION_MESSAGE_TYPE))
51+
require.Equal(t, BOOT_REPLY, response.Header.Op)
52+
require.Equal(t, DHCPOFFER, response.Options.GetByte(OPTION_MESSAGE_TYPE))
4853
require.Equal(t, IpToFixedV4(net.ParseIP("10.0.0.10")), response.Header.YourAddr)
4954

55+
require.Equal(t, []FixedV4{IpToFixedV4(net.ParseIP("255.255.255.0"))}, response.Options.GetFixedV4s(OPTION_SUBNET))
56+
require.Equal(t, []FixedV4{IpToFixedV4(net.ParseIP("10.0.0.1"))}, response.Options.GetFixedV4s(OPTION_ROUTER))
57+
require.Equal(t, []FixedV4{IpToFixedV4(net.ParseIP("1.1.1.1")), IpToFixedV4(net.ParseIP("1.0.0.1"))}, response.Options.GetFixedV4s(OPTION_DNS_SERVER))
58+
require.Equal(t, []FixedV4{IpToFixedV4(net.ParseIP("10.0.0.254"))}, response.Options.GetFixedV4s(OPTION_SERVER_ID))
59+
5060
// Pool should have a lease for this mac
5161
lease, ok := pool.TouchLeaseByMac(message.Header.Mac)
5262
require.True(t, ok)
@@ -61,12 +71,13 @@ func TestDhcpDiscover(t *testing.T) {
6171

6272
message, err = ParseDhcpMessage(b)
6373
require.Nil(t, err)
64-
require.Equal(t, byte(DHCPREQUEST), message.Options.GetByte(OPTION_MESSAGE_TYPE))
74+
require.Equal(t, BOOT_REPLY, response.Header.Op)
75+
require.Equal(t, DHCPREQUEST, message.Options.GetByte(OPTION_MESSAGE_TYPE))
6576

6677
handler = NewRequestHandler(message, pool)
6778
response = handler.Handle()
6879

69-
require.Equal(t, byte(DHCPNAK), response.Options.GetByte(OPTION_MESSAGE_TYPE))
80+
require.Equal(t, DHCPNAK, response.Options.GetByte(OPTION_MESSAGE_TYPE))
7081

7182
//
7283
// A request should now get back an ACK
@@ -77,12 +88,13 @@ func TestDhcpDiscover(t *testing.T) {
7788

7889
message, err = ParseDhcpMessage(b)
7990
require.Nil(t, err)
80-
require.Equal(t, byte(DHCPREQUEST), message.Options.GetByte(OPTION_MESSAGE_TYPE))
91+
require.Equal(t, DHCPREQUEST, message.Options.GetByte(OPTION_MESSAGE_TYPE))
8192

8293
handler = NewRequestHandler(message, pool)
8394
response = handler.Handle()
8495

85-
require.Equal(t, byte(DHCPACK), response.Options.GetByte(OPTION_MESSAGE_TYPE))
96+
require.Equal(t, BOOT_REPLY, response.Header.Op)
97+
require.Equal(t, DHCPACK, response.Options.GetByte(OPTION_MESSAGE_TYPE))
8698

8799
//
88100
// Do a DHCPRELEASE
@@ -93,7 +105,7 @@ func TestDhcpDiscover(t *testing.T) {
93105

94106
message, err = ParseDhcpMessage(b)
95107
require.Nil(t, err)
96-
require.Equal(t, byte(DHCPRELEASE), message.Options.GetByte(OPTION_MESSAGE_TYPE))
108+
require.Equal(t, DHCPRELEASE, message.Options.GetByte(OPTION_MESSAGE_TYPE))
97109

98110
handler = NewRequestHandler(message, pool)
99111
response = handler.Handle()

0 commit comments

Comments
 (0)