feat: Turn into Cobra application, add CLI flag for listen interface

cmd/root.go

@@ -0,0 +1,330 @@
+package cmd
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"net"
+	"os"
+	"unsafe"
+
+	"github.com/adrianokf/go-dhcp/pkg/leases"
+	"github.com/adrianokf/go-dhcp/pkg/messages"
+	"github.com/adrianokf/go-dhcp/pkg/types"
+	"github.com/spf13/cobra"
+	"go.uber.org/zap"
+)
+
+var listenInterface string
+
+var manager = leases.NewLeaseManager()
+
+var magic = [4]byte{0x63, 0x82, 0x53, 0x63}
+
+func parseOptions(data []byte) messages.Options {
+	s := zap.S()
+	i := 0
+
+	options := make(messages.Options)
+
+out:
+	for i < len(data) {
+		code := messages.OptionCode(data[i])
+		switch code {
+		case messages.OptionEnd:
+			s.Debug("Found END option at offset ", i)
+			break out
+
+		case messages.OptionPad:
+			s.Debug("Found padding option at offset ", i)
+			i += 1
+			continue
+		}
+
+		size := int(data[i+1])
+		payload := data[i+2 : i+2+size]
+
+		s.Debugf("code=%d, size=%d, payload=%x", code, size, payload)
+		options[code] = messages.Option{Code: code, Data: data[i+1 : i+2+size]}
+		i += size + 2
+	}
+
+	s.Debugf("Parsed options: ", options)
+	return options
+}
+
+func prepareOffer(request messages.DhcpMessage, lease leases.Lease) messages.DhcpMessage {
+	var sname [64]byte
+	var file [128]byte
+	var siaddr types.Ipv4Addr
+	copy(sname[:], "go-dhcp-server")
+	copy(siaddr[:], net.IPv4(10, 0, 0, 1).To4())
+
+	dhcp := messages.DhcpMessage{
+		Op:     messages.BOOTREPLY,
+		Htype:  1, // Ethernet
+		Hlen:   6, // Ethernet address length
+		Hops:   0,
+		Secs:   0,
+		Flags:  request.Flags,
+		Xid:    lease.TransactionId,
+		Siaddr: siaddr,
+		Ciaddr: [4]byte{0, 0, 0, 0},
+		Yiaddr: lease.Address,
+		Giaddr: request.Giaddr,
+		Chaddr: request.Chaddr,
+		Magic:  magic,
+		Sname:  sname,
+		File:   file,
+	}
+	return dhcp
+}
+
+func prepareAck(request messages.DhcpMessage, lease leases.Lease) messages.DhcpMessage {
+	var sname [64]byte
+	var file [128]byte
+	copy(sname[:], "go-dhcp-server")
+
+	dhcp := messages.DhcpMessage{
+		Op:     messages.BOOTREPLY,
+		Htype:  1, // Ethernet
+		Hlen:   6, // Ethernet address length
+		Hops:   0,
+		Secs:   0,
+		Flags:  0,
+		Xid:    request.Xid,
+		Siaddr: lease.Address,
+		Ciaddr: request.Ciaddr,
+		Yiaddr: lease.Address,
+		Giaddr: request.Giaddr,
+		Chaddr: request.Chaddr,
+		Magic:  magic,
+		Sname:  sname,
+		File:   file,
+	}
+	return dhcp
+}
+
+// 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 messages.DhcpMessage, options []messages.Option) error {
+	buf := make([]byte, 0)
+	w := bytes.NewBuffer(buf)
+	err := binary.Write(w, binary.BigEndian, message)
+	if err != nil {
+		return err
+	}
+
+	for _, v := range options {
+		err = w.WriteByte(byte(v.Code))
+		if err != nil {
+			return err
+		}
+		_, err = w.Write(v.Data)
+		if err != nil {
+			return err
+		}
+	}
+
+	// Automatically add END option, so the caller doesn't
+	// need to specificy it for every invocation.
+	err = w.WriteByte(byte(messages.OptionEnd))
+	if err != nil {
+		return err
+	}
+
+	msg := w.Bytes()
+	zap.S().Debug("Msg", msg)
+	_, err = conn.Write(msg)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func handleOffer(dhcp messages.DhcpMessage, remote *net.UDPAddr) error {
+	s := zap.S()
+
+	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)
+	if err != nil {
+		panic(err)
+	}
+	defer conn.Close()
+
+	s.Info("Sending DHCPOFFER...")
+	options := []messages.Option{
+		{
+			Code: messages.OptionDHCPMessageType,
+			Data: []byte{1, byte(messages.MessageTypeOffer)},
+		},
+	}
+	sendMessage(conn, offer, options)
+	return nil
+}
+
+func handleAck(dhcp messages.DhcpMessage, remote *net.UDPAddr) error {
+	s := zap.S()
+
+	lease, err := manager.Lookup(dhcp.Xid)
+	if err != nil {
+		panic(err)
+	}
+
+	ack := prepareAck(dhcp, *lease)
+
+	options := []messages.Option{
+		{
+			Code: messages.OptionDHCPMessageType,
+			Data: []byte{1, byte(messages.MessageTypeAck)},
+		},
+		{
+			Code: messages.OptionIPAddressLeaseTime,
+			Data: append([]byte{4}, u32tob(3600)...),
+		},
+	}
+	s.Debug("Options: ", options)
+
+	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)
+	if err != nil {
+		panic(err)
+	}
+	defer conn.Close()
+
+	s.Info("Sending DHCPACK")
+	sendMessage(conn, ack, options)
+
+	lease, err = manager.Request(dhcp.Xid, dhcp.Chaddr)
+	if err != nil {
+		return err
+	}
+	s.Debug("Found lease", lease)
+	return nil
+}
+
+func handleMsg(data []byte, remote *net.UDPAddr) {
+	s := zap.S()
+
+	s.Debugf("Connection from client %v", remote.IP)
+
+	var dhcp messages.DhcpMessage
+	reader := bytes.NewReader(data)
+	binary.Read(reader, binary.BigEndian, &dhcp)
+	dhcp.Debug(s)
+
+	if dhcp.Magic != [4]byte{0x63, 0x82, 0x53, 0x63} {
+		panic("Invalid DHCP magic field")
+	}
+
+	optDataOffset := int(unsafe.Sizeof(dhcp))
+	optData := data[optDataOffset:]
+	s.Debug("Raw options data:", optData)
+
+	options := parseOptions(optData)
+	dhcpMsgType := options[messages.OptionDHCPMessageType]
+	s.Info("DHCP message type ", dhcpMsgType)
+
+	switch messages.MessageType(dhcpMsgType.Data[1]) {
+	case messages.MessageTypeDiscover:
+		go handleOffer(dhcp, remote)
+
+	case messages.MessageTypeRequest:
+		go handleAck(dhcp, remote)
+	}
+}
+
+func runServer(interfaceName string) {
+	var addr *net.UDPAddr
+	if interfaceName == "all" {
+		zap.L().Debug("Listening on all interfaces")
+		addr, _ = net.ResolveUDPAddr("udp4", ":67")
+	} else {
+		zap.S().Debugf("Listening on interface %s", interfaceName)
+		iface, err := net.InterfaceByName(interfaceName)
+		if err != nil {
+			panic(err)
+		}
+		addrs, err := iface.Addrs()
+		if err != nil {
+			panic(err)
+		}
+
+		// Find first IPv4 address associated with the interface
+		var ip net.IP = nil
+		for _, ifaddr := range addrs {
+			switch a := ifaddr.(type) {
+			case *net.IPAddr:
+				ip = a.IP
+			case *net.IPNet:
+				ip = a.IP
+			}
+
+			ip = ip.To4()
+			if ip != nil {
+				break
+			}
+		}
+		if ip == nil {
+			zap.S().Panicf("No IPv4 address associated with interface %s", interfaceName)
+		}
+		fmt.Printf("%+v\n", ip)
+		addr, _ = net.ResolveUDPAddr("udp4", ip.String()+":67")
+	}
+	conn, err := net.ListenUDP("udp4", addr)
+	if err != nil {
+		panic(err)
+	}
+	defer conn.Close()
+
+	zap.S().Infof("Listening for incoming connections on %s", addr.String())
+
+	for {
+		buf := make([]byte, 1024)
+		rlen, remote, err := conn.ReadFrom(buf[:])
+		if err != nil {
+			panic(err)
+		}
+
+		// Do stuff with the read bytes
+		remoteAddr, ok := remote.(*net.UDPAddr)
+		if !ok {
+			zap.S().Warn("Not a valid remote IP address: ", remote)
+			continue
+		}
+		go handleMsg(buf[0:rlen], remoteAddr)
+	}
+}
+
+var rootCmd = &cobra.Command{
+	Use:   "go-dhcp",
+	Short: "go-dhcp is a simple DHCP server written in Go",
+	Run: func(cmd *cobra.Command, args []string) {
+		runServer(listenInterface)
+	},
+}
+
+func init() {
+	// Set up logging
+	logger, _ := zap.NewDevelopment()
+	defer logger.Sync() // flushes buffer, if any
+	zap.ReplaceGlobals(logger)
+
+	rootCmd.PersistentFlags().StringVarP(&listenInterface, "interface", "i", "all", "Interface to listen on")
+}
+
+func Execute() {
+	if err := rootCmd.Execute(); err != nil {
+		fmt.Println(err)
+		os.Exit(1)
+	}
+}

cmd/util.go

@@ -1,16 +1,9 @@
-package main
+package cmd
 
 import (
 	"encoding/binary"
-	"net"
 )
 
-func int2ip(nn uint32) net.IP {
-	ip := make(net.IP, 4)
-	binary.BigEndian.PutUint32(ip, nn)
-	return ip
-}
-
 func u32tob(u uint32) []byte {
 	buf := make([]byte, 4)
 	binary.BigEndian.PutUint32(buf, u)

go.mod

@@ -5,6 +5,12 @@ go 1.17
 require go.uber.org/zap v1.26.0
 
 require (
+	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
+)
+
+require (
+	github.com/spf13/cobra v1.8.0
 	go.uber.org/atomic v1.11.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 )

go.sum

@@ -1,8 +1,11 @@
 github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
+github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
+github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
@@ -10,6 +13,11 @@ github.com/pkg/errors v0.8.1 h1:iURUrRGxPUNPdy5/HRSm+Yj6okJ6UtLINN0Q9M4+h3I=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/spf13/cobra v1.8.0 h1:7aJaZx1B85qltLMc546zn58BxxfZdR/W22ej9CFoEf0=
+github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho=
+github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
+github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
@@ -63,3 +71,4 @@ gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b h1:h8qDotaEPuJATrMmW04NCwg7v22aHH28wwpauUhK9Oo=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

main.go

@@ -1,271 +1,9 @@
 package main
 
 import (
-	"bytes"
+	"github.com/adrianokf/go-dhcp/cmd"
-	"encoding/binary"
-	"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 = leases.NewLeaseManager()
-
-var magic = [4]byte{0x63, 0x82, 0x53, 0x63}
-
-func parseOptions(data []byte) messages.Options {
-	s := zap.S()
-	i := 0
-
-	options := make(messages.Options)
-
-out:
-	for i < len(data) {
-		code := messages.OptionCode(data[i])
-		switch code {
-		case messages.OptionEnd:
-			s.Debug("Found END option at offset ", i)
-			break out
-
-		case messages.OptionPad:
-			s.Debug("Found padding option at offset ", i)
-			i += 1
-			continue
-		}
-
-		size := int(data[i+1])
-		payload := data[i+2 : i+2+size]
-
-		s.Debugf("code=%d, size=%d, payload=%x", code, size, payload)
-		options[code] = messages.Option{Code: code, Data: data[i+1 : i+2+size]}
-		i += size + 2
-	}
-
-	s.Debugf("Parsed options: ", options)
-	return options
-}
-
-func prepareOffer(request messages.DhcpMessage, lease leases.Lease) messages.DhcpMessage {
-	var sname [64]byte
-	var file [128]byte
-	var siaddr types.Ipv4Addr
-	copy(sname[:], "go-dhcp-server")
-	copy(siaddr[:], net.IPv4(10, 0, 0, 1).To4())
-
-	dhcp := messages.DhcpMessage{
-		Op:     messages.BOOTREPLY,
-		Htype:  1, // Ethernet
-		Hlen:   6, // Ethernet address length
-		Hops:   0,
-		Secs:   0,
-		Flags:  request.Flags,
-		Xid:    lease.TransactionId,
-		Siaddr: siaddr,
-		Ciaddr: [4]byte{0, 0, 0, 0},
-		Yiaddr: lease.Address,
-		Giaddr: request.Giaddr,
-		Chaddr: request.Chaddr,
-		Magic:  magic,
-		Sname:  sname,
-		File:   file,
-	}
-	return dhcp
-}
-
-func prepareAck(request messages.DhcpMessage, lease leases.Lease) messages.DhcpMessage {
-	var sname [64]byte
-	var file [128]byte
-	copy(sname[:], "go-dhcp-server")
-
-	dhcp := messages.DhcpMessage{
-		Op:     messages.BOOTREPLY,
-		Htype:  1, // Ethernet
-		Hlen:   6, // Ethernet address length
-		Hops:   0,
-		Secs:   0,
-		Flags:  0,
-		Xid:    request.Xid,
-		Siaddr: lease.Address,
-		Ciaddr: request.Ciaddr,
-		Yiaddr: lease.Address,
-		Giaddr: request.Giaddr,
-		Chaddr: request.Chaddr,
-		Magic:  magic,
-		Sname:  sname,
-		File:   file,
-	}
-	return dhcp
-}
-
-// 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 messages.DhcpMessage, options []messages.Option) error {
-	buf := make([]byte, 0)
-	w := bytes.NewBuffer(buf)
-	err := binary.Write(w, binary.BigEndian, message)
-	if err != nil {
-		return err
-	}
-
-	for _, v := range options {
-		err = w.WriteByte(byte(v.Code))
-		if err != nil {
-			return err
-		}
-		_, err = w.Write(v.Data)
-		if err != nil {
-			return err
-		}
-	}
-
-	// Automatically add END option, so the caller doesn't
-	// need to specificy it for every invocation.
-	err = w.WriteByte(byte(messages.OptionEnd))
-	if err != nil {
-		return err
-	}
-
-	msg := w.Bytes()
-	zap.S().Debug("Msg", msg)
-	_, err = conn.Write(msg)
-	if err != nil {
-		return err
-	}
-
-	return nil
-}
-
-func handleOffer(dhcp messages.DhcpMessage, remote *net.UDPAddr) error {
-	s := zap.S()
-
-	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)
-	if err != nil {
-		panic(err)
-	}
-	defer conn.Close()
-
-	s.Info("Sending DHCPOFFER...")
-	options := []messages.Option{
-		{
-			Code: messages.OptionDHCPMessageType,
-			Data: []byte{1, byte(messages.MessageTypeOffer)},
-		},
-	}
-	sendMessage(conn, offer, options)
-	return nil
-}
-
-func handleAck(dhcp messages.DhcpMessage, remote *net.UDPAddr) error {
-	s := zap.S()
-
-	lease, err := manager.Lookup(dhcp.Xid)
-	if err != nil {
-		panic(err)
-	}
-
-	ack := prepareAck(dhcp, *lease)
-
-	options := []messages.Option{
-		{
-			Code: messages.OptionDHCPMessageType,
-			Data: []byte{1, byte(messages.MessageTypeAck)},
-		},
-		{
-			Code: messages.OptionIPAddressLeaseTime,
-			Data: append([]byte{4}, u32tob(3600)...),
-		},
-	}
-	s.Debug("Options: ", options)
-
-	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)
-	if err != nil {
-		panic(err)
-	}
-	defer conn.Close()
-
-	s.Info("Sending DHCPACK")
-	sendMessage(conn, ack, options)
-
-	lease, err = manager.Request(dhcp.Xid, dhcp.Chaddr)
-	if err != nil {
-		return err
-	}
-	s.Debug("Found lease", lease)
-	return nil
-}
-
-func handleMsg(data []byte, remote *net.UDPAddr) {
-	s := zap.S()
-
-	s.Debugf("Connection from client %v", remote.IP)
-
-	var dhcp messages.DhcpMessage
-	reader := bytes.NewReader(data)
-	binary.Read(reader, binary.BigEndian, &dhcp)
-	dhcp.Debug(s)
-
-	if dhcp.Magic != [4]byte{0x63, 0x82, 0x53, 0x63} {
-		panic("Invalid DHCP magic field")
-	}
-
-	optDataOffset := int(unsafe.Sizeof(dhcp))
-	optData := data[optDataOffset:]
-	s.Debug("Raw options data:", optData)
-
-	options := parseOptions(optData)
-	dhcpMsgType := options[messages.OptionDHCPMessageType]
-	s.Info("DHCP message type ", dhcpMsgType)
-
-	switch messages.MessageType(dhcpMsgType.Data[1]) {
-	case messages.MessageTypeDiscover:
-		go handleOffer(dhcp, remote)
-
-	case messages.MessageTypeRequest:
-		go handleAck(dhcp, remote)
-	}
-}
-
 func main() {
-	// Set up logging
+	cmd.Execute()
-	logger, _ := zap.NewDevelopment()
-	defer logger.Sync() // flushes buffer, if any
-	zap.ReplaceGlobals(logger)
-
-	addr, _ := net.ResolveUDPAddr("udp4", ":67")
-	conn, err := net.ListenUDP("udp4", addr)
-	if err != nil {
-		panic(err)
-	}
-	defer conn.Close()
-
-	zap.L().Info("Listening for incoming connections")
-
-	for {
-		buf := make([]byte, 1024)
-		rlen, remote, err := conn.ReadFrom(buf[:])
-		if err != nil {
-			panic(err)
-		}
-
-		// Do stuff with the read bytes
-		remoteAddr, ok := remote.(*net.UDPAddr)
-		if !ok {
-			zap.S().Warn("Not a valid remote IP address: ", remote)
-			continue
-		}
-		go handleMsg(buf[0:rlen], remoteAddr)
-	}
 }