feat: Working lease manager

lease.go

@@ -0,0 +1,67 @@
+package main
+
+import (
+	"errors"
+	"fmt"
+	"time"
+
+	"go.uber.org/zap"
+)
+
+type ILeaseManager interface {
+	Request(xid TxId, clientAddr HwAddr) (*Lease, error)
+	Release(l Lease) error
+	Lookup(xid TxId) error
+}
+
+type LeaseManager struct {
+	leases map[TxId]Lease
+	count  byte
+}
+
+func NewLeaseManager() *LeaseManager {
+	m := &LeaseManager{
+		leases: make(map[TxId]Lease),
+	}
+	return m
+}
+
+func (m *LeaseManager) Request(xid TxId, clientAddr HwAddr) (*Lease, error) {
+	zap.S().Debugf("LeaseManager.Request(%v, %v)", xid, clientAddr)
+
+	if m.count > 254 {
+		return nil, errors.New("lease address pool exhausted")
+	}
+
+	assigned := [4]byte{10, 0, 0, m.count + 2}
+	lease := Lease{
+		TransactionId: xid,
+		TTL:           time.Now().Add(1 * time.Hour),
+		State:         Requested,
+		ClientAddr:    clientAddr,
+		Address:       assigned,
+	}
+	m.leases[xid] = lease
+
+	m.count += 1
+
+	zap.S().Debug("lease=", lease)
+	return &lease, nil
+}
+
+func (m LeaseManager) Release(l Lease) error {
+	_, found := m.leases[l.TransactionId]
+	if !found {
+		return fmt.Errorf("invalid lease %v", l)
+	}
+	delete(m.leases, l.TransactionId)
+	return nil
+}
+
+func (m LeaseManager) Lookup(xid TxId) (*Lease, error) {
+	lease, found := m.leases[xid]
+	if found {
+		return &lease, nil
+	}
+	return nil, fmt.Errorf("no lease found for xid %v", xid)
+}

main.go

@@ -4,12 +4,13 @@ import (
 	"bytes"
 	"encoding/binary"
 	"net"
-	"time"
 	"unsafe"
 
 	"go.uber.org/zap"
 )
 
+var manager = NewLeaseManager()
+
 var leases []Lease = make([]Lease, 0)
 
 func (dhcp DhcpMessage) Debug(s *zap.SugaredLogger) {
@@ -59,10 +60,12 @@ out:
 	return options
 }
 
-func prepareOffer(request DhcpMessage) DhcpMessage {
+func prepareOffer(request DhcpMessage, lease Lease) DhcpMessage {
 	var sname [64]byte
 	var file [128]byte
+	var siaddr Ipv4Addr
 	copy(sname[:], "go-dhcp-server")
+	copy(siaddr[:], net.IPv4(10, 0, 0, 1).To4())
 
 	dhcp := DhcpMessage{
 		Op:     BOOTREPLY,
@@ -71,20 +74,20 @@ func prepareOffer(request DhcpMessage) DhcpMessage {
 		Hops:   0,
 		Secs:   0,
 		Flags:  request.Flags,
-		Xid:    request.Xid,
-		Siaddr: binary.BigEndian.Uint32(net.IPv4(10, 0, 0, 1).To4()),
-		Ciaddr: 0,
-		Yiaddr: binary.BigEndian.Uint32(net.IPv4(10, 0, 0, 2).To4()),
+		Xid:    lease.TransactionId,
+		Siaddr: siaddr,
+		Ciaddr: [4]byte{0, 0, 0, 0},
+		Yiaddr: lease.Address,
 		Giaddr: request.Giaddr,
 		Chaddr: request.Chaddr,
-		Magic:  [4]byte{0x63, 0x82, 0x53, 0x63},
+		Magic:  magic,
 		Sname:  sname,
 		File:   file,
 	}
 	return dhcp
 }
 
-func prepareAck(request DhcpMessage) DhcpMessage {
+func prepareAck(request DhcpMessage, lease Lease) DhcpMessage {
 	var sname [64]byte
 	var file [128]byte
 	copy(sname[:], "go-dhcp-server")
@@ -97,12 +100,12 @@ func prepareAck(request DhcpMessage) DhcpMessage {
 		Secs:   0,
 		Flags:  0,
 		Xid:    request.Xid,
-		Siaddr: binary.BigEndian.Uint32(net.IPv4(10, 0, 0, 1).To4()),
+		Siaddr: lease.Address,
 		Ciaddr: request.Ciaddr,
-		Yiaddr: binary.BigEndian.Uint32(net.IPv4(10, 0, 0, 2).To4()),
+		Yiaddr: lease.Address,
 		Giaddr: request.Giaddr,
 		Chaddr: request.Chaddr,
-		Magic:  [4]byte{0x63, 0x82, 0x53, 0x63},
+		Magic:  magic,
 		Sname:  sname,
 		File:   file,
 	}
@@ -151,7 +154,11 @@ func sendMessage(conn *net.UDPConn, message DhcpMessage, options []Option) error
 func handleOffer(dhcp DhcpMessage, remote *net.UDPAddr) error {
 	s := zap.S()
 
-	offer := prepareOffer(dhcp)
+	lease, err := manager.Request(dhcp.Xid, dhcp.Chaddr)
+	if err != nil {
+		panic(err)
+	}
+	offer := prepareOffer(dhcp, *lease)
 	localAddr, _ := net.ResolveUDPAddr("udp", "172.17.0.1:68")
 	clientAddr, _ := net.ResolveUDPAddr("udp", "255.255.255.255:68")
 	conn, err := net.DialUDP("udp", localAddr, clientAddr)
@@ -174,7 +181,12 @@ func handleOffer(dhcp DhcpMessage, remote *net.UDPAddr) error {
 func handleAck(dhcp DhcpMessage, remote *net.UDPAddr) error {
 	s := zap.S()
 
-	ack := prepareAck(dhcp)
+	lease, err := manager.Lookup(dhcp.Xid)
+	if err != nil {
+		panic(err)
+	}
+
+	ack := prepareAck(dhcp, *lease)
 
 	options := []Option{
 		{
@@ -199,15 +211,11 @@ func handleAck(dhcp DhcpMessage, remote *net.UDPAddr) error {
 	s.Info("Sending DHCPACK")
 	sendMessage(conn, ack, options)
 
-	lease := Lease{
-		TransactionId: ack.Xid,
-		TTL:           time.Now().Add(1 * time.Hour),
-		State:         Requested,
-		ClientAddr:    ack.Chaddr,
-		Address:       int2ip(ack.Yiaddr),
+	lease, err = manager.Request(dhcp.Xid, dhcp.Chaddr)
+	if err != nil {
+		return err
 	}
-	leases = append(leases, lease)
-
+	s.Debug("Found lease", lease)
 	return nil
 }
 

types.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"net"
 	"time"
 )
 
@@ -197,20 +196,25 @@ const (
 )
 
 type LeaseState uint8
+type TxId [4]byte
+type Ipv4Addr [4]byte
+type HwAddr [16]byte
+
+var magic = [4]byte{0x63, 0x82, 0x53, 0x63}
 
 type DhcpMessage struct {
 	Op     DhcpOp
 	Htype  MessageType
 	Hlen   byte
 	Hops   byte
-	Xid    uint32
+	Xid    TxId
 	Secs   uint16
 	Flags  uint16
-	Ciaddr uint32
-	Yiaddr uint32
-	Siaddr uint32
-	Giaddr uint32
-	Chaddr [16]byte
+	Ciaddr Ipv4Addr
+	Yiaddr Ipv4Addr
+	Siaddr Ipv4Addr
+	Giaddr Ipv4Addr
+	Chaddr HwAddr
 	Sname  [64]byte
 	File   [128]byte
 	Magic  [4]byte
@@ -222,9 +226,9 @@ const (
 )
 
 type Lease struct {
-	TransactionId uint32
+	TransactionId TxId
 	TTL           time.Time
 	State         LeaseState
-	ClientAddr    [16]byte
-	Address       net.IP
+	ClientAddr    HwAddr
+	Address       Ipv4Addr
 }