forked from krolaw/dhcp4
-
Notifications
You must be signed in to change notification settings - Fork 0
/
example_test.go
107 lines (96 loc) · 3.17 KB
/
example_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
// Example of minimal DHCP server:
package dhcp4_test
import (
dhcp "github.com/krolaw/dhcp4"
"log"
"math/rand"
"net"
"time"
)
// Example using DHCP with a single network interface device
func ExampleHandler() {
serverIP := net.IP{172, 30, 0, 1}
handler := &DHCPHandler{
ip: serverIP,
leaseDuration: 2 * time.Hour,
start: net.IP{172, 30, 0, 2},
leaseRange: 50,
leases: make(map[int]lease, 10),
options: dhcp.Options{
dhcp.OptionSubnetMask: []byte{255, 255, 240, 0},
dhcp.OptionRouter: []byte(serverIP), // Presuming Server is also your router
dhcp.OptionDomainNameServer: []byte(serverIP), // Presuming Server is also your DNS server
},
}
log.Fatal(dhcp.ListenAndServe(handler))
// log.Fatal(dhcp.ListenAndServeIf("eth0",handler)) // Select interface on multi interface device
}
type lease struct {
nic string // Client's CHAddr
expiry time.Time // When the lease expires
}
type DHCPHandler struct {
ip net.IP // Server IP to use
options dhcp.Options // Options to send to DHCP Clients
start net.IP // Start of IP range to distribute
leaseRange int // Number of IPs to distribute (starting from start)
leaseDuration time.Duration // Lease period
leases map[int]lease // Map to keep track of leases
}
func (h *DHCPHandler) ServeDHCP(p dhcp.Packet, msgType dhcp.MessageType, options dhcp.Options) (d dhcp.Packet) {
switch msgType {
case dhcp.Discover:
free, nic := -1, p.CHAddr().String()
for i, v := range h.leases { // Find previous lease
if v.nic == nic {
free = i
goto reply
}
}
if free = h.freeLease(); free == -1 {
return
}
reply:
return dhcp.ReplyPacket(p, dhcp.Offer, h.ip, dhcp.IPAdd(h.start, free), h.leaseDuration,
h.options.SelectOrderOrAll(options[dhcp.OptionParameterRequestList]))
case dhcp.Request:
if server, ok := options[dhcp.OptionServerIdentifier]; ok && !net.IP(server).Equal(h.ip) {
return nil // Message not for this dhcp server
}
reqIP := net.IP(options[dhcp.OptionRequestedIPAddress])
if reqIP == nil {
reqIP = net.IP(p.CIAddr())
}
if len(reqIP) == 4 && !reqIP.Equal(net.IPv4zero) {
if leaseNum := dhcp.IPRange(h.start, reqIP) - 1; leaseNum >= 0 && leaseNum < h.leaseRange {
if l, exists := h.leases[leaseNum]; !exists || l.nic == p.CHAddr().String() {
h.leases[leaseNum] = lease{nic: p.CHAddr().String(), expiry: time.Now().Add(h.leaseDuration)}
return dhcp.ReplyPacket(p, dhcp.ACK, h.ip, net.IP(options[dhcp.OptionRequestedIPAddress]), h.leaseDuration,
h.options.SelectOrderOrAll(options[dhcp.OptionParameterRequestList]))
}
}
}
return dhcp.ReplyPacket(p, dhcp.NAK, h.ip, nil, 0, nil)
case dhcp.Release, dhcp.Decline:
nic := p.CHAddr().String()
for i, v := range h.leases {
if v.nic == nic {
delete(h.leases, i)
break
}
}
}
return nil
}
func (h *DHCPHandler) freeLease() int {
now := time.Now()
b := rand.Intn(h.leaseRange) // Try random first
for _, v := range [][]int{[]int{b, h.leaseRange}, []int{0, b}} {
for i := v[0]; i < v[1]; i++ {
if l, ok := h.leases[i]; !ok || l.expiry.Before(now) {
return i
}
}
}
return -1
}