diff --git a/lease.go b/lease.go new file mode 100644 index 0000000..6b69d7e --- /dev/null +++ b/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) +} diff --git a/main.go b/main.go index fbaf5b0..d689f12 100644 --- a/main.go +++ b/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 } diff --git a/types.go b/types.go index 3becada..708652c 100644 --- a/types.go +++ b/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 }