chore: Refactor package structure

2024-01-04 18:49:54 +01:00
parent 13677eaf54
commit 001517eb15
5 changed files with 89 additions and 79 deletions
--- 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)
 	}
 }
--- a/pkg/leases/leases.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
--- a/pkg/messages/messages.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)
 }
--- a/pkg/messages/options.go
+++ b/pkg/messages/options.go
@@ -1,4 +1,4 @@
-package main
+package messages

 type Option struct {
 	Code OptionCode
--- a/pkg/types/types.go
+++ b/pkg/types/types.go
@@ -0,0 +1,5 @@
+package types
+
+type TxId [4]byte
+type Ipv4Addr [4]byte
+type HwAddr [16]byte