NetBootServerCore/NetBootServerCore/Components/NetBoot/Dhcp/DhcpServer.cs

using NetBootServerCore.Components.Logging;
using NetBootServerCore.Components.NetBoot.Dhcp.Pools;
using NetBootServerCore.Components.NetBoot.Dhcp.Protocol;
using SharpPcap;
using SharpPcap.LibPcap;
using System.Net;
using System.Net.Sockets;
using static NetBootServerCore.Components.NetBoot.Dhcp.Protocol.Enums;
using ProtocolType = System.Net.Sockets.ProtocolType;

namespace NetBootServerCore.Components.NetBoot.Dhcp
{
    public class DhcpServer : NbsComponentBase
    {
        private const int Port = 67;
        private readonly IPEndPoint _broadcastTargetAddress = new IPEndPoint(IPAddress.Broadcast, Port + 1);
        private readonly IConfigurationSection _config;

        private IPEndPoint _listenOn;

        private IPAddress _poolStart;
        private int _poolSize;
        private TimeSpan _leaseTime;

        private IPAddress _subNet;
        private IPAddress[] _gateway;
        private IPAddress[] _dns;

        private IPAddress _broadcastAddress;

        private IPAddress _tftp;
        private string _fileName;
        private string _fileNameForIpxe;

        private LibPcapLiveDevice _nic;
        private ARP _arp;
        public AddressPool Pool { get; private set; }

        private bool _isRunning;
        private Task _task;

        private Socket _server;
        private bool _enable;

        public DhcpServer(IConfigurationSection config) : base(new Logger("DHCP")) => _config = config;

        private void ProcessRequest(byte[] buffer)
        {
            var pdu = new DhcpPdu(buffer);

            var clientMac = pdu.ClientMacAddressHex;

            Logger.LogDebug($"Request from {clientMac}");

            //extract options before clear
            var vendor = pdu.Vendor;
            var userClass = pdu.UserClass;
            var hostName = pdu.HostName;
            var requestMessageType = pdu.MessageType;

            Logger.LogDebug($"Message type: {requestMessageType}");
            Logger.LogDebug($"Vendor: {vendor ?? "(null)"}");
            Logger.LogDebug($"UserClass: {userClass ?? "(null)"}");
            Logger.LogDebug($"HostName: {hostName ?? "(null)"}");

            AddressPoolSlot aSlot = null;
            {
                var timeFilter = DateTime.Now.Add(-_leaseTime);
                var dSlot = Pool.Slots.FirstOrDefault(p => p.Mac == clientMac && (p.LastUpdateTime.HasValue == false || p.LastUpdateTime.Value > timeFilter));
                if (dSlot != null)
                {
                    var eMac = _arp.Resolve(dSlot.Address)?.GetAddressBytes().ToHypHexString();
                    if (dSlot.Mac == eMac || eMac == null) aSlot = dSlot;
                }

                if (aSlot == null)
                {
                    foreach (var slot in Pool.Slots)
                    {
                        if (slot.LastUpdateTime.HasValue && slot.LastUpdateTime > timeFilter && slot.Mac != clientMac) continue;

                        var physicalAddress = _arp.Resolve(slot.Address);
                        if (physicalAddress == null)
                        {
                            aSlot = slot;
                            break;
                        }

                        var macHex = physicalAddress.GetAddressBytes().ToHypHexString();
                        if (slot.Mac != macHex)
                        {
                            slot.IsDirty = true;
                            slot.LastUpdateTime = DateTime.Now;
                        }
                        else if (slot.Mac == null)
                        {
                            slot.Mac = macHex;
                            slot.LastUpdateTime = DateTime.Now;
                        }
                    }
                }
            }

            // ready for reply
            pdu.OpCode = DhcpOpCode.BootReply;
            pdu.Options.Clear();

            if (aSlot == null)
            {
                pdu.MessageType = DhcpMessageType.Nak;
                Logger.LogWarn($"Allocate fail {clientMac}");
            }
            else
            {
                if (requestMessageType == DhcpMessageType.Discover)
                    pdu.MessageType = DhcpMessageType.Offer;
                else if (requestMessageType == DhcpMessageType.Request)
                    pdu.MessageType = DhcpMessageType.Ack;
                else
                {
                    Logger.LogWarn($"Unexpected message type: {requestMessageType}({(byte)requestMessageType})");
                    return;
                }

                aSlot.Mac = clientMac;
                aSlot.HostName = hostName;
                aSlot.Vendor = vendor;
                aSlot.LastUpdateTime = DateTime.Now;
                aSlot.IsDirty = false;

                pdu.YourIpAddress = aSlot.Address;

                pdu.Routers = _gateway;
                pdu.SubNetMask = _subNet;
                pdu.DnsServers = _dns;
                pdu.BroadcastAddress = _broadcastAddress;
                pdu.LeaseTime = _leaseTime;

                pdu.DhcpServerIdentifier = _listenOn.Address;
                pdu.TftpServer = _tftp.ToString();
                pdu.NextServerIpAddress = _tftp;

                if (userClass == "iPXE") pdu.BootFileName = _fileNameForIpxe;
                else
                {
                    if (vendor == null || !vendor.StartsWith("PXEClient")) pdu.BootFileName = _fileName;
                    else
                    {
                        var parts = vendor.Split(':');
                        if (parts.Length < 3) pdu.BootFileName = _fileName;
                        else
                        {
                            var prefix = string.Join(":", parts[..3]);
                            switch (prefix)
                            {
                                //case "PXEClient:Arch:00002":
                                //case "PXEClient:Arch:00006":
                                //  UEFI x86
                                //  break;

                                //case "PXEClient:Arch:00007":
                                //case "PXEClient:Arch:00008":
                                //case "PXEClient:Arch:00009":
                                //  UEFI x64
                                //  break;

                                //case "PXEClient:Arch:00000": //Legacy
                                default:
                                    pdu.BootFileName = _fileName;
                                    break;
                            }
                        }
                    }
                }
            }

            var len = pdu.WriteToBuffer(buffer);
            _server.SendTo(buffer, 0, len, SocketFlags.None, _broadcastTargetAddress);

            Logger.LogVerbose($"{clientMac} -> {pdu.YourIpAddress} {pdu.MessageType}");
        }

        protected override void OnInit()
        {
            _enable = bool.Parse(_config["Enable"]);
            if (!_enable)
            {
                Logger.LogInfo("Disabled by config");
                return;
            }

            _listenOn = new IPEndPoint(IPAddress.Parse(_config["ListenOn"]), Port);

            _poolStart = IPAddress.Parse(_config["PoolStart"]);
            _poolSize = int.Parse(_config["PoolSize"]);
            _leaseTime = TimeSpan.FromHours(int.Parse(_config["LeaseTimeInHours"]));

            _subNet = IPAddress.Parse(_config["SubNet"]);
            _gateway = _config.GetSection("Gateway").Get<string[]>().Select(IPAddress.Parse).ToArray();
            _dns = _config.GetSection("DNS").Get<string[]>().Select(IPAddress.Parse).ToArray();

            _broadcastAddress = IPAddress.Parse(_config["BroadcastAddress"]);

            _tftp = IPAddress.Parse(_config["Tftp"]);
            _fileName = _config["FileName"];
            _fileNameForIpxe = _config["FileName_iPXE"];

            Pool = new AddressPool(_poolStart, _poolSize);

            _nic?.Dispose();
            _nic = LibPcapLiveDeviceList.Instance.FirstOrDefault(p => p.Addresses.Any(q => Equals(q.Addr.ipAddress, _listenOn.Address)));
            _arp = new ARP(_nic);
        }

        protected override void OnStart()
        {
            if (!_enable)
            {
                Logger.LogInfo("Disabled by config");
                return;
            }

            _isRunning = true;
            _task = Task.Run(() =>
            {
                _server = new Socket(_listenOn.AddressFamily, SocketType.Dgram, ProtocolType.Udp)
                {
                    EnableBroadcast = true,
                    SendBufferSize = 1500,
                    ReceiveBufferSize = 1500
                };
                _server.Bind(_listenOn);
                var buffer = new byte[1500];
                EndPoint remoteEndPoint = new IPEndPoint(IPAddress.Any, 68);
                while (_isRunning)
                {
                    if (!_server.Poll(500 * 1000, SelectMode.SelectRead)) continue;

                    _server.ReceiveFrom(buffer, 0, buffer.Length, SocketFlags.None, ref remoteEndPoint);
                    try
                    {
                        ProcessRequest(buffer);
                    }
                    catch (Exception e)
                    {
                        Logger.LogError(nameof(ProcessRequest), e);
                    }
                }
            });
        }

        protected override void OnStop()
        {
            if (!_enable)
            {
                Logger.LogInfo("Disabled by config");
                return;
            }

            _isRunning = false;
            _task.Wait(10 * 1000);
        }
    }
}