From e164cff02b0f20a03c0c3050eb8f6e4a812f461b Mon Sep 17 00:00:00 2001
From: Simon Zolin <s.zolin@adguard.com>
Date: Fri, 1 Mar 2019 20:03:22 +0300
Subject: [PATCH] + dhcp: test

---
 dhcpd/dhcpd.go      |   8 +++
 dhcpd/dhcpd_test.go | 115 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 123 insertions(+)
 create mode 100644 dhcpd/dhcpd_test.go

diff --git a/dhcpd/dhcpd.go b/dhcpd/dhcpd.go
index d6f1c528..b9affb5a 100644
--- a/dhcpd/dhcpd.go
+++ b/dhcpd/dhcpd.go
@@ -368,3 +368,11 @@ func (s *Server) Leases() []*Lease {
 	s.RUnlock()
 	return result
 }
+
+// Reset internal state
+func (s *Server) reset() {
+	s.Lock()
+	s.leases = nil
+	s.Unlock()
+	s.IPpool = make(map[[4]byte]net.HardwareAddr)
+}
diff --git a/dhcpd/dhcpd_test.go b/dhcpd/dhcpd_test.go
new file mode 100644
index 00000000..10ca1b59
--- /dev/null
+++ b/dhcpd/dhcpd_test.go
@@ -0,0 +1,115 @@
+package dhcpd
+
+import (
+	"bytes"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/krolaw/dhcp4"
+)
+
+func check(t *testing.T, result bool, msg string) {
+	if !result {
+		t.Fatal(msg)
+	}
+}
+
+// Tests performed:
+// . Handle Discover message (lease reserve)
+// . Handle Request message (lease commit)
+func TestDHCP(t *testing.T) {
+	var s = Server{}
+	var p, p2 dhcp4.Packet
+	var hw net.HardwareAddr
+	var lease *Lease
+	var opt dhcp4.Options
+
+	s.leaseStart = []byte{1, 1, 1, 1}
+	s.leaseStop = []byte{1, 1, 1, 2}
+	s.leaseTime = 5 * time.Second
+	s.leaseOptions = dhcp4.Options{}
+	s.ipnet = &net.IPNet{
+		IP:   []byte{1, 2, 3, 4},
+		Mask: []byte{0xff, 0xff, 0xff, 0xff},
+	}
+
+	p = make(dhcp4.Packet, 241)
+
+	// Reserve an IP
+	hw = []byte{1, 2, 3, 4, 5, 6}
+	p.SetCHAddr(hw)
+	lease, _ = s.reserveLease(p)
+	check(t, bytes.Equal(lease.HWAddr, hw), "lease.HWAddr")
+	check(t, bytes.Equal(lease.IP, []byte{1, 1, 1, 1}), "lease.IP")
+
+	// Try to reserve another IP for the same machine - no new IP must be reserved
+	hw = []byte{1, 2, 3, 4, 5, 6}
+	p.SetCHAddr(hw)
+	lease, _ = s.reserveLease(p)
+	check(t, bytes.Equal(lease.HWAddr, hw), "lease.HWAddr")
+	check(t, bytes.Equal(lease.IP, []byte{1, 1, 1, 1}), "lease.IP")
+
+	// Reserve an IP - the next IP from the range
+	hw = []byte{2, 2, 3, 4, 5, 6}
+	p.SetCHAddr(hw)
+	lease, _ = s.reserveLease(p)
+	check(t, bytes.Equal(lease.HWAddr, hw), "lease.HWAddr")
+	check(t, bytes.Equal(lease.IP, []byte{1, 1, 1, 2}), "lease.IP")
+
+	// Reserve an IP - we have no more available IPs
+	p.SetCHAddr([]byte{3, 2, 3, 4, 5, 6})
+	lease, _ = s.reserveLease(p)
+	check(t, lease == nil, "lease == nil")
+
+	// Decline request for a lease which doesn't match our internal state
+	hw = []byte{1, 2, 3, 4, 5, 6}
+	p.SetCHAddr(hw)
+	p.SetCIAddr([]byte{0, 0, 0, 0})
+	opt = make(dhcp4.Options, 10)
+	// ask a different IP
+	opt[dhcp4.OptionRequestedIPAddress] = []byte{1, 1, 1, 2}
+	p2 = s.handleDHCP4Request(p, opt)
+	opt = p2.ParseOptions()
+	check(t, bytes.Equal(opt[dhcp4.OptionDHCPMessageType], []byte{byte(dhcp4.NAK)}), "dhcp4.NAK")
+
+	// Commit the previously reserved lease
+	hw = []byte{1, 2, 3, 4, 5, 6}
+	p.SetCHAddr(hw)
+	p.SetCIAddr([]byte{0, 0, 0, 0})
+	opt = make(dhcp4.Options, 10)
+	opt[dhcp4.OptionRequestedIPAddress] = []byte{1, 1, 1, 1}
+	p2 = s.handleDHCP4Request(p, opt)
+	opt = p2.ParseOptions()
+	check(t, bytes.Equal(opt[dhcp4.OptionDHCPMessageType], []byte{byte(dhcp4.ACK)}), "dhcp4.ACK")
+	check(t, bytes.Equal(p2.YIAddr(), []byte{1, 1, 1, 1}), "p2.YIAddr")
+	check(t, bytes.Equal(p2.CHAddr(), hw), "p2.CHAddr")
+	check(t, bytes.Equal(opt[dhcp4.OptionIPAddressLeaseTime], dhcp4.OptionsLeaseTime(5*time.Second)), "OptionIPAddressLeaseTime")
+	check(t, bytes.Equal(opt[dhcp4.OptionServerIdentifier], s.ipnet.IP), "OptionServerIdentifier")
+
+	s.reset()
+	misc(t, &s)
+}
+
+// Small tests that don't require a static server's state
+func misc(t *testing.T, s *Server) {
+	var p, p2 dhcp4.Packet
+	var hw net.HardwareAddr
+	var opt dhcp4.Options
+
+	p = make(dhcp4.Packet, 241)
+
+	// Commit a lease for an IP without prior Discover-Offer packets
+	hw = []byte{2, 2, 3, 4, 5, 6}
+	p.SetCHAddr(hw)
+	p.SetCIAddr([]byte{0, 0, 0, 0})
+	opt = make(dhcp4.Options, 10)
+	opt[dhcp4.OptionRequestedIPAddress] = []byte{1, 1, 1, 1}
+	p2 = s.handleDHCP4Request(p, opt)
+	opt = p2.ParseOptions()
+	check(t, bytes.Equal(opt[dhcp4.OptionDHCPMessageType], []byte{byte(dhcp4.ACK)}), "dhcp4.ACK")
+	check(t, bytes.Equal(p2.YIAddr(), []byte{1, 1, 1, 1}), "p2.YIAddr")
+	check(t, bytes.Equal(p2.CHAddr(), hw), "p2.CHAddr")
+	check(t, bytes.Equal(opt[dhcp4.OptionIPAddressLeaseTime], dhcp4.OptionsLeaseTime(5*time.Second)), "OptionIPAddressLeaseTime")
+	check(t, bytes.Equal(opt[dhcp4.OptionServerIdentifier], s.ipnet.IP), "OptionServerIdentifier")
+}