go-dhcp/cmd/root.go

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/adrianokf/go-dhcp/pkg/util"
	"github.com/spf13/cobra"
	"go.uber.org/zap"
)

var listenInterface string

type Handler struct {
	conn         *net.UDPConn
	leaseManager leases.LeaseManager
}

var magic = [4]byte{0x63, 0x82, 0x53, 0x63}

func parseOptions(data []byte) messages.Options {
	i := 0
	options := make(messages.Options)

out:
	for i < len(data) {
		code := messages.OptionCode(data[i])
		switch code {
		case messages.OptionEnd:
			zap.S().Debug("Found END option at offset ", i)
			break out

		case messages.OptionPad:
			zap.S().Debug("Found padding option at offset ", i)
			i += 1
			continue
		}

		size := int(data[i+1])
		payload := data[i+2 : i+2+size]

		zap.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
	}

	zap.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 (h Handler) sendMessage(remote *net.UDPAddr, message messages.DhcpMessage, options []messages.Option) error {
	// Send packets for 0.0.0.0 to broadcast address (255.255.255.255) instead
	var destination *net.UDPAddr
	if remote.IP.IsUnspecified() {
		destination = &net.UDPAddr{
			IP:   net.ParseIP("255.255.255.255"),
			Port: types.ClientPort,
		}
	} else {
		destination = remote
	}

	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 specify it for every invocation.
	err = w.WriteByte(byte(messages.OptionEnd))
	if err != nil {
		return err
	}

	msg := w.Bytes()
	zap.S().Debugf("Local addr: %s, remote addr: %s ", h.conn.LocalAddr(), destination)
	zap.S().Debug("Msg data: ", msg)
	_, err = h.conn.WriteToUDP(msg, destination)
	if err != nil {
		return err
	}

	return nil
}

func (h Handler) sendOffer(dhcp messages.DhcpMessage, remote *net.UDPAddr) error {
	s := zap.S()

	lease, err := h.leaseManager.Request(dhcp.Xid, dhcp.Chaddr)
	if err != nil {
		panic(err)
	}
	offer := prepareOffer(dhcp, *lease)

	s.Info("Sending DHCPOFFER...")
	options := []messages.Option{
		{
			Code: messages.OptionDHCPMessageType,
			Data: []byte{1, byte(messages.MessageTypeOffer)},
		},
	}
	h.sendMessage(remote, offer, options)
	return nil
}

func (h Handler) sendAck(dhcp messages.DhcpMessage, remote *net.UDPAddr) error {
	s := zap.S()

	lease, err := h.leaseManager.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}, util.U32ToByte(3600)...),
		},
	}
	s.Debug("Options: ", options)

	s.Info("Sending DHCPACK")
	h.sendMessage(remote, ack, options)

	lease, err = h.leaseManager.Request(dhcp.Xid, dhcp.Chaddr)
	if err != nil {
		return err
	}
	s.Debug("Found lease: ", lease)
	return nil
}

func (h Handler) 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 != magic {
		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 h.sendOffer(dhcp, remote)

	case messages.MessageTypeRequest:
		go h.sendAck(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().Debug("Listening on interface ", 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().Panic("No IPv4 address associated with interface ", interfaceName)
		}
		addr, _ = net.ResolveUDPAddr("udp4", ":67")
	}

	var conn *net.UDPConn
	var err error
	if interfaceName != "all" {
		conn, err = NewIPv4UDPConn(interfaceName, addr)
	} else {
		conn, err = net.ListenUDP("udp4", addr)
	}
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	handler := Handler{
		conn:         conn,
		leaseManager: *leases.NewLeaseManager(),
	}

	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 handler.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)
	}
}