From 001517eb15e17f0ca8adfa5e144a3f7611a099a6 Mon Sep 17 00:00:00 2001 From: Adrian Rumpold Date: Thu, 4 Jan 2024 18:49:54 +0100 Subject: [PATCH] chore: Refactor package structure --- main.go | 80 ++++++++++++--------------- lease.go => pkg/leases/leases.go | 30 +++++++--- types.go => pkg/messages/messages.go | 51 ++++++++--------- options.go => pkg/messages/options.go | 2 +- pkg/types/types.go | 5 ++ 5 files changed, 89 insertions(+), 79 deletions(-) rename lease.go => pkg/leases/leases.go (62%) rename types.go => pkg/messages/messages.go (92%) rename options.go => pkg/messages/options.go (94%) create mode 100644 pkg/types/types.go diff --git a/main.go b/main.go index d689f12..19fe422 100644 --- a/main.go +++ b/main.go @@ -6,43 +6,31 @@ import ( "net" "unsafe" + "github.com/adrianokf/go-dhcp/pkg/leases" + "github.com/adrianokf/go-dhcp/pkg/messages" + "github.com/adrianokf/go-dhcp/pkg/types" "go.uber.org/zap" ) -var manager = NewLeaseManager() +var manager = leases.NewLeaseManager() -var leases []Lease = make([]Lease, 0) +var magic = [4]byte{0x63, 0x82, 0x53, 0x63} -func (dhcp DhcpMessage) Debug(s *zap.SugaredLogger) { - s.Debugf("op=%x, htype=%x, hlen=%x, hops=%x", dhcp.Op, dhcp.Htype, dhcp.Hlen, dhcp.Hops) - s.Debugf("xid=%x", dhcp.Xid) - s.Debugf("secs=%d, flags=%x", dhcp.Secs, dhcp.Flags) - - s.Debugf("ciaddr=%x", dhcp.Ciaddr) - s.Debugf("siaddr=%x", dhcp.Siaddr) - s.Debugf("giaddr=%x", dhcp.Giaddr) - s.Debugf("chaddr=%x", dhcp.Chaddr) - - s.Debugf("sname=%s", string(dhcp.Sname[:])) - s.Debugf("file=%s", string(dhcp.File[:])) - s.Debugf("magic=%x", dhcp.Magic) -} - -func parseOptions(data []byte) Options { +func parseOptions(data []byte) messages.Options { s := zap.S() i := 0 - options := make(Options) + options := make(messages.Options) out: for i < len(data) { - code := OptionCode(data[i]) + code := messages.OptionCode(data[i]) switch code { - case OptionEnd: + case messages.OptionEnd: s.Debug("Found END option at offset ", i) break out - case OptionPad: + case messages.OptionPad: s.Debug("Found padding option at offset ", i) i += 1 continue @@ -52,7 +40,7 @@ out: payload := data[i+2 : i+2+size] s.Debugf("code=%d, size=%d, payload=%x", code, size, payload) - options[code] = Option{Code: code, Data: data[i+1 : i+2+size]} + options[code] = messages.Option{Code: code, Data: data[i+1 : i+2+size]} i += size + 2 } @@ -60,15 +48,15 @@ out: return options } -func prepareOffer(request DhcpMessage, lease Lease) DhcpMessage { +func prepareOffer(request messages.DhcpMessage, lease leases.Lease) messages.DhcpMessage { var sname [64]byte var file [128]byte - var siaddr Ipv4Addr + var siaddr types.Ipv4Addr copy(sname[:], "go-dhcp-server") copy(siaddr[:], net.IPv4(10, 0, 0, 1).To4()) - dhcp := DhcpMessage{ - Op: BOOTREPLY, + dhcp := messages.DhcpMessage{ + Op: messages.BOOTREPLY, Htype: 1, // Ethernet Hlen: 6, // Ethernet address length Hops: 0, @@ -87,13 +75,13 @@ func prepareOffer(request DhcpMessage, lease Lease) DhcpMessage { return dhcp } -func prepareAck(request DhcpMessage, lease Lease) DhcpMessage { +func prepareAck(request messages.DhcpMessage, lease leases.Lease) messages.DhcpMessage { var sname [64]byte var file [128]byte copy(sname[:], "go-dhcp-server") - dhcp := DhcpMessage{ - Op: BOOTREPLY, + dhcp := messages.DhcpMessage{ + Op: messages.BOOTREPLY, Htype: 1, // Ethernet Hlen: 6, // Ethernet address length Hops: 0, @@ -115,7 +103,7 @@ func prepareAck(request DhcpMessage, lease Lease) DhcpMessage { // sendMessage transmits a DHCP message with options via a UDP connection // The end option (code 255) is automatically appended and does not need to // be passed explicitly. -func sendMessage(conn *net.UDPConn, message DhcpMessage, options []Option) error { +func sendMessage(conn *net.UDPConn, message messages.DhcpMessage, options []messages.Option) error { buf := make([]byte, 0) w := bytes.NewBuffer(buf) err := binary.Write(w, binary.BigEndian, message) @@ -136,7 +124,7 @@ func sendMessage(conn *net.UDPConn, message DhcpMessage, options []Option) error // Automatically add END option, so the caller doesn't // need to specificy it for every invocation. - err = w.WriteByte(byte(OptionEnd)) + err = w.WriteByte(byte(messages.OptionEnd)) if err != nil { return err } @@ -151,7 +139,7 @@ func sendMessage(conn *net.UDPConn, message DhcpMessage, options []Option) error return nil } -func handleOffer(dhcp DhcpMessage, remote *net.UDPAddr) error { +func handleOffer(dhcp messages.DhcpMessage, remote *net.UDPAddr) error { s := zap.S() lease, err := manager.Request(dhcp.Xid, dhcp.Chaddr) @@ -168,17 +156,17 @@ func handleOffer(dhcp DhcpMessage, remote *net.UDPAddr) error { defer conn.Close() s.Info("Sending DHCPOFFER...") - options := []Option{ + options := []messages.Option{ { - Code: OptionDHCPMessageType, - Data: []byte{1, byte(MessageTypeOffer)}, + Code: messages.OptionDHCPMessageType, + Data: []byte{1, byte(messages.MessageTypeOffer)}, }, } sendMessage(conn, offer, options) return nil } -func handleAck(dhcp DhcpMessage, remote *net.UDPAddr) error { +func handleAck(dhcp messages.DhcpMessage, remote *net.UDPAddr) error { s := zap.S() lease, err := manager.Lookup(dhcp.Xid) @@ -188,13 +176,13 @@ func handleAck(dhcp DhcpMessage, remote *net.UDPAddr) error { ack := prepareAck(dhcp, *lease) - options := []Option{ + options := []messages.Option{ { - Code: OptionDHCPMessageType, - Data: []byte{1, byte(MessageTypeAck)}, + Code: messages.OptionDHCPMessageType, + Data: []byte{1, byte(messages.MessageTypeAck)}, }, { - Code: OptionIPAddressLeaseTime, + Code: messages.OptionIPAddressLeaseTime, Data: append([]byte{4}, u32tob(3600)...), }, } @@ -224,7 +212,7 @@ func handleMsg(data []byte, remote *net.UDPAddr) { s.Debugf("Connection from client %v", remote.IP) - var dhcp DhcpMessage + var dhcp messages.DhcpMessage reader := bytes.NewReader(data) binary.Read(reader, binary.BigEndian, &dhcp) dhcp.Debug(s) @@ -238,14 +226,14 @@ func handleMsg(data []byte, remote *net.UDPAddr) { s.Debug("Raw options data:", optData) options := parseOptions(optData) - dhcpMsgType := options[OptionDHCPMessageType] + dhcpMsgType := options[messages.OptionDHCPMessageType] s.Info("DHCP message type ", dhcpMsgType) - switch MessageType(dhcpMsgType.Data[1]) { - case MessageTypeDiscover: + switch messages.MessageType(dhcpMsgType.Data[1]) { + case messages.MessageTypeDiscover: go handleOffer(dhcp, remote) - case MessageTypeRequest: + case messages.MessageTypeRequest: go handleAck(dhcp, remote) } } diff --git a/lease.go b/pkg/leases/leases.go similarity index 62% rename from lease.go rename to pkg/leases/leases.go index 6b69d7e..fe02d1e 100644 --- a/lease.go +++ b/pkg/leases/leases.go @@ -1,32 +1,48 @@ -package main +package leases import ( "errors" "fmt" "time" + "github.com/adrianokf/go-dhcp/pkg/types" "go.uber.org/zap" ) +type LeaseState uint8 + +const ( + Offered LeaseState = iota + Requested +) + +type Lease struct { + TransactionId types.TxId + TTL time.Time + State LeaseState + ClientAddr types.HwAddr + Address types.Ipv4Addr +} + type ILeaseManager interface { - Request(xid TxId, clientAddr HwAddr) (*Lease, error) + Request(xid types.TxId, clientAddr types.HwAddr) (*Lease, error) Release(l Lease) error - Lookup(xid TxId) error + Lookup(xid types.TxId) error } type LeaseManager struct { - leases map[TxId]Lease + leases map[types.TxId]Lease count byte } func NewLeaseManager() *LeaseManager { m := &LeaseManager{ - leases: make(map[TxId]Lease), + leases: make(map[types.TxId]Lease), } return m } -func (m *LeaseManager) Request(xid TxId, clientAddr HwAddr) (*Lease, error) { +func (m *LeaseManager) Request(xid types.TxId, clientAddr types.HwAddr) (*Lease, error) { zap.S().Debugf("LeaseManager.Request(%v, %v)", xid, clientAddr) if m.count > 254 { @@ -58,7 +74,7 @@ func (m LeaseManager) Release(l Lease) error { return nil } -func (m LeaseManager) Lookup(xid TxId) (*Lease, error) { +func (m LeaseManager) Lookup(xid types.TxId) (*Lease, error) { lease, found := m.leases[xid] if found { return &lease, nil diff --git a/types.go b/pkg/messages/messages.go similarity index 92% rename from types.go rename to pkg/messages/messages.go index 708652c..6e9cb80 100644 --- a/types.go +++ b/pkg/messages/messages.go @@ -1,9 +1,15 @@ -package main +package messages import ( - "time" + "github.com/adrianokf/go-dhcp/pkg/types" + "go.uber.org/zap" ) +type Message interface { + Marshal() ([]byte, error) + Unmarshal([]byte) (Message, error) +} + type OptionCode byte // DHCPv4 Options @@ -195,40 +201,35 @@ const ( MessageTypeInform ) -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 TxId + Xid types.TxId Secs uint16 Flags uint16 - Ciaddr Ipv4Addr - Yiaddr Ipv4Addr - Siaddr Ipv4Addr - Giaddr Ipv4Addr - Chaddr HwAddr + Ciaddr types.Ipv4Addr + Yiaddr types.Ipv4Addr + Siaddr types.Ipv4Addr + Giaddr types.Ipv4Addr + Chaddr types.HwAddr Sname [64]byte File [128]byte Magic [4]byte } -const ( - Offered LeaseState = iota - Requested -) +func (dhcp DhcpMessage) Debug(s *zap.SugaredLogger) { + s.Debugf("op=%x, htype=%x, hlen=%x, hops=%x", dhcp.Op, dhcp.Htype, dhcp.Hlen, dhcp.Hops) + s.Debugf("xid=%x", dhcp.Xid) + s.Debugf("secs=%d, flags=%x", dhcp.Secs, dhcp.Flags) -type Lease struct { - TransactionId TxId - TTL time.Time - State LeaseState - ClientAddr HwAddr - Address Ipv4Addr + s.Debugf("ciaddr=%x", dhcp.Ciaddr) + s.Debugf("siaddr=%x", dhcp.Siaddr) + s.Debugf("giaddr=%x", dhcp.Giaddr) + s.Debugf("chaddr=%x", dhcp.Chaddr) + + s.Debugf("sname=%s", string(dhcp.Sname[:])) + s.Debugf("file=%s", string(dhcp.File[:])) + s.Debugf("magic=%x", dhcp.Magic) } diff --git a/options.go b/pkg/messages/options.go similarity index 94% rename from options.go rename to pkg/messages/options.go index 646257b..71bec0c 100644 --- a/options.go +++ b/pkg/messages/options.go @@ -1,4 +1,4 @@ -package main +package messages type Option struct { Code OptionCode diff --git a/pkg/types/types.go b/pkg/types/types.go new file mode 100644 index 0000000..448da90 --- /dev/null +++ b/pkg/types/types.go @@ -0,0 +1,5 @@ +package types + +type TxId [4]byte +type Ipv4Addr [4]byte +type HwAddr [16]byte