dhcp-dns-manager/internal/dhcp/server.go

package dhcp

import (
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"os"
	"sync"
	"time"
	"dhcp-dns-manager/internal/db"
	"dhcp-dns-manager/internal/config"
	"golang.org/x/sys/unix"
)

// DHCP 消息类型
const (
	MsgDiscover = 1
	MsgOffer    = 2
	MsgRequest  = 3
	MsgDecline  = 4
	MsgACK      = 5
	MsgNAK      = 6
	MsgRelease  = 7
)

// DHCP 选项
const (
	OptionSubnetMask       = 1
	OptionRouter           = 3
	OptionDNS              = 6
	OptionHostname         = 12
	OptionLeaseTime        = 51
	OptionMessageType      = 53
	OptionServerIdentifier = 54
	OptionRequestedIP      = 50
	OptionEnd              = 255
)

type Server struct {
	config         *config.DHCPConfig
	configReloader func() *config.DHCPConfig // optional: reload config from ConfigManager
	db             *db.DB
	leases         map[string]*db.DHCPLease
	staticBindings map[string]db.DHCPStaticBinding
	leaseMutex     sync.RWMutex
	conn           *net.UDPConn
	stopChan       chan struct{}
	serverIP       net.IP
	usedIPs        map[string]string // IP -> MAC
}

func NewServer(cfg *config.DHCPConfig, database *db.DB) *Server {
	return &Server{
		config:         cfg,
		db:             database,
		leases:         make(map[string]*db.DHCPLease),
		staticBindings: make(map[string]db.DHCPStaticBinding),
		usedIPs:        make(map[string]string),
		stopChan:       make(chan struct{}),
		serverIP:       net.ParseIP(cfg.Gateway).To4(),
	}
}

// SetConfigReloader sets a function to reload config dynamically.
// This allows the DHCP server to pick up config changes made via the web UI.
func (s *Server) SetConfigReloader(reloader func() *config.DHCPConfig) {
	s.configReloader = reloader
}

// getConfig returns the current config, reloading if a reloader is set.
func (s *Server) getConfig() *config.DHCPConfig {
	if s.configReloader != nil {
		return s.configReloader()
	}
	return s.config
}

// getServerIP returns the server IP for DHCP options
func (s *Server) getServerIP() net.IP {
	cfg := s.getConfig()
	return net.ParseIP(cfg.Gateway).To4()
}

func (s *Server) Start() error {
	if !s.config.Enabled {
		return nil
	}

	// Load existing leases
	s.loadLeases()
	s.loadStaticBindings()

	// Start lease cleanup goroutine
	go s.cleanupLeases()

	// Start DHCP server on UDP port 67
	// Use a raw connection to properly handle broadcast responses
	conn, err := newBroadcastUDPConn("0.0.0.0", 67)
	if err != nil {
		return fmt.Errorf("failed to listen on UDP 67: %v", err)
	}

	s.conn = conn
	log.Printf("DHCP server listening on 0.0.0.0:67")

	// Handle DHCP requests
	go s.handleDHCP()

	return nil
}

func (s *Server) Stop() {
	if s.conn != nil {
		s.conn.Close()
	}
	close(s.stopChan)
}

func (s *Server) handleDHCP() {
	buf := make([]byte, 1024)

	for {
		select {
		case <-s.stopChan:
			return
		default:
			s.conn.SetReadDeadline(time.Now().Add(1 * time.Second))
			n, remoteAddr, err := s.conn.ReadFromUDP(buf)
			if err != nil {
				continue
			}

			s.processDHCPMessage(buf[:n], remoteAddr)
		}
	}
}

func (s *Server) processDHCPMessage(data []byte, remoteAddr *net.UDPAddr) {
	if len(data) < 240 {
		return
	}

	// Parse DHCP message
	msgType := parseMessageType(data)
	if msgType == 0 {
		return
	}

	// Get client MAC
	clientMAC := formatMAC(data[28:34])

	// Get client IP from packet
	clientIP := net.IP(data[16:20])

	switch msgType {
	case MsgDiscover:
		s.handleDiscover(data, clientMAC, clientIP, remoteAddr)
	case MsgRequest:
		s.handleRequest(data, clientMAC, clientIP, remoteAddr)
	case MsgRelease:
		s.handleRelease(clientMAC)
	case MsgDecline:
		s.handleDecline(clientMAC)
	}
}

func (s *Server) handleDiscover(data []byte, clientMAC string, clientIP net.IP, remoteAddr *net.UDPAddr) {
	// Find or assign IP
	offeredIP := s.assignIP(clientMAC, data)
	if offeredIP == "" {
		log.Printf("DHCP: No available IP for %s", clientMAC)
		return
	}

	// Record the offered IP as a provisional lease BEFORE sending the Offer.
	// This ensures that when the client sends Request back, verifyAssignment
	// will find the lease and return true (instead of NAKing the client).
	s.recordLease(clientMAC, offeredIP, data)

	// Send DHCP Offer
	s.sendOffer(data, clientMAC, offeredIP, remoteAddr)
	log.Printf("DHCP: Offered %s to %s", offeredIP, clientMAC)
}

func (s *Server) handleRequest(data []byte, clientMAC string, clientIP net.IP, remoteAddr *net.UDPAddr) {
	// Get requested IP
	requestedIP := parseRequestedIP(data).String()
	if requestedIP == "<nil>" || requestedIP == "" {
		// Fallback to ciaddr
		requestedIP = clientIP.String()
	}

	// Verify the requested IP
	if s.verifyAssignment(clientMAC, requestedIP) {
		// Send DHCP ACK
		s.sendACK(data, clientMAC, requestedIP, remoteAddr)
		log.Printf("DHCP: ACK %s to %s", requestedIP, clientMAC)

		// Record lease
		s.recordLease(clientMAC, requestedIP, data)
	} else {
		// Send DHCP NAK
		s.sendNAK(data, remoteAddr)
		log.Printf("DHCP: NAK for %s (requested %s)", clientMAC, requestedIP)
	}
}

func (s *Server) handleRelease(clientMAC string) {
	s.leaseMutex.Lock()
	defer s.leaseMutex.Unlock()

	if lease, exists := s.leases[clientMAC]; exists {
		log.Printf("DHCP: Released %s for %s", lease.IP, clientMAC)
		delete(s.leases, clientMAC)
		delete(s.usedIPs, lease.IP)
	}
}

func (s *Server) handleDecline(clientMAC string) {
	s.leaseMutex.Lock()
	defer s.leaseMutex.Unlock()

	if lease, exists := s.leases[clientMAC]; exists {
		log.Printf("DHCP: Declined %s for %s", lease.IP, clientMAC)
		delete(s.leases, clientMAC)
		delete(s.usedIPs, lease.IP)
	}
}

func (s *Server) assignIP(clientMAC string, data []byte) string {
	s.leaseMutex.Lock()
	defer s.leaseMutex.Unlock()

	// Check static binding first
	if binding, exists := s.staticBindings[clientMAC]; exists {
		return binding.IP
	}

	// Check if client already has a lease
	if lease, exists := s.leases[clientMAC]; exists {
		return lease.IP
	}

	// Find available IP
	startIP := net.ParseIP(s.getConfig().IPPoolStart).To4()
	endIP := net.ParseIP(s.getConfig().IPPoolEnd).To4()

	if startIP == nil || endIP == nil {
		return ""
	}

	startInt := uint32(startIP[0])<<24 | uint32(startIP[1])<<16 | uint32(startIP[2])<<8 | uint32(startIP[3])
	endInt := uint32(endIP[0])<<24 | uint32(endIP[1])<<16 | uint32(endIP[2])<<8 | uint32(endIP[3])

	// Build set of used IPs
	usedIPs := make(map[string]bool)
	for _, lease := range s.leases {
		usedIPs[lease.IP] = true
	}

	// Find first available IP
	for ip := startInt; ip <= endInt; ip++ {
		ipBytes := []byte{
			byte(ip >> 24),
			byte(ip >> 16),
			byte(ip >> 8),
			byte(ip),
		}
		ipStr := fmt.Sprintf("%d.%d.%d.%d", ipBytes[0], ipBytes[1], ipBytes[2], ipBytes[3])

		// Skip gateway and excluded IPs
		if ipStr == s.getConfig().Gateway {
			continue
		}

		excluded := false
		for _, excl := range s.getConfig().ExcludedIPs {
			if ipStr == excl {
				excluded = true
				break
			}
		}
		if excluded {
			continue
		}

		if !usedIPs[ipStr] {
			return ipStr
		}
	}

	return ""
}

func (s *Server) verifyAssignment(clientMAC, ip string) bool {
	s.leaseMutex.RLock()
	defer s.leaseMutex.RUnlock()

	// Check static binding
	if binding, exists := s.staticBindings[clientMAC]; exists {
		return binding.IP == ip
	}

	// Check if IP is assigned to this client
	if lease, exists := s.leases[clientMAC]; exists {
		return lease.IP == ip
	}

	// IP not in lease map yet — this is a new client that just got an Offer.
	// Verify the IP is within the configured pool range and not already taken.
	startIP := net.ParseIP(s.getConfig().IPPoolStart).To4()
	endIP := net.ParseIP(s.getConfig().IPPoolEnd).To4()
	if startIP == nil || endIP == nil {
		return false
	}

	clientIP := net.ParseIP(ip).To4()
	if clientIP == nil {
		return false
	}

	clientInt := uint32(clientIP[0])<<24 | uint32(clientIP[1])<<16 | uint32(clientIP[2])<<8 | uint32(clientIP[3])
	startInt := uint32(startIP[0])<<24 | uint32(startIP[1])<<16 | uint32(startIP[2])<<8 | uint32(startIP[3])
	endInt := uint32(endIP[0])<<24 | uint32(endIP[1])<<16 | uint32(endIP[2])<<8 | uint32(endIP[3])

	if clientInt < startInt || clientInt > endInt {
		return false
	}

	// Make sure no other client already has this IP
	for mac, lease := range s.leases {
		if lease.IP == ip && mac != clientMAC {
			return false
		}
	}

	return true
}

func (s *Server) recordLease(clientMAC, ip string, data []byte) {
	lease := &db.DHCPLease{
		MAC:       clientMAC,
		IP:        ip,
		ExpiresAt: time.Now().Add(time.Duration(s.getConfig().LeaseTime) * time.Second).Unix(),
	}

	// Try to get hostname from DHCP options
	if hostname := parseHostname(data); hostname != "" {
		lease.Hostname = hostname
	}

	s.leaseMutex.Lock()
	// Check if lease already exists (e.g. from Offer phase)
	existingLease, alreadyExists := s.leases[clientMAC]
	s.leases[clientMAC] = lease
	s.usedIPs[ip] = clientMAC
	s.leaseMutex.Unlock()

	// Save to database — update if already exists to avoid duplicates
	if alreadyExists && existingLease.ID > 0 {
		lease.ID = existingLease.ID
		s.db.Save(lease)
	} else {
		s.db.Create(lease)
	}
}

func (s *Server) sendOffer(data []byte, clientMAC, offeredIP string, remoteAddr *net.UDPAddr) {
	// Build DHCP OFFER message
	response := buildDHCPMessage(MsgOffer, data, offeredIP, s.config)

	// Add options
	response = appendOption(response, OptionSubnetMask, []byte(net.ParseIP(s.getConfig().Netmask).To4()))
	response = appendOption(response, OptionRouter, []byte(net.ParseIP(s.getConfig().Gateway).To4()))

	// Add DNS servers
	if len(s.getConfig().DNSServers) > 0 {
		var dnsBytes []byte
		for _, dns := range s.getConfig().DNSServers {
			dnsBytes = append(dnsBytes, net.ParseIP(dns).To4()...)
		}
		response = appendOption(response, OptionDNS, dnsBytes)
	}

	// Add lease time
	leaseTime := make([]byte, 4)
	binary.BigEndian.PutUint32(leaseTime, uint32(s.getConfig().LeaseTime))
	response = appendOption(response, OptionLeaseTime, leaseTime)

	// Add server identifier
	response = appendOption(response, OptionServerIdentifier, []byte(s.getServerIP().To4()))

	// Add end option
	response = append(response, OptionEnd)

	// Send response — use broadcast if client has no IP yet (ciaddr == 0.0.0.0)
	targetAddr := s.getResponseAddr(data, remoteAddr)
	s.conn.WriteToUDP(response, targetAddr)
}

func (s *Server) sendACK(data []byte, clientMAC, ip string, remoteAddr *net.UDPAddr) {
	response := buildDHCPMessage(MsgACK, data, ip, s.config)

	// Add options
	response = appendOption(response, OptionSubnetMask, []byte(net.ParseIP(s.getConfig().Netmask).To4()))
	response = appendOption(response, OptionRouter, []byte(net.ParseIP(s.getConfig().Gateway).To4()))

	if len(s.getConfig().DNSServers) > 0 {
		var dnsBytes []byte
		for _, dns := range s.getConfig().DNSServers {
			dnsBytes = append(dnsBytes, net.ParseIP(dns).To4()...)
		}
		response = appendOption(response, OptionDNS, dnsBytes)
	}

	leaseTime := make([]byte, 4)
	binary.BigEndian.PutUint32(leaseTime, uint32(s.getConfig().LeaseTime))
	response = appendOption(response, OptionLeaseTime, leaseTime)

	response = appendOption(response, OptionServerIdentifier, []byte(s.getServerIP().To4()))
	response = append(response, OptionEnd)

	targetAddr := s.getResponseAddr(data, remoteAddr)
	s.conn.WriteToUDP(response, targetAddr)
}

func (s *Server) sendNAK(data []byte, remoteAddr *net.UDPAddr) {
	response := buildDHCPMessage(MsgNAK, data, "0.0.0.0", s.config)
	response = appendOption(response, OptionServerIdentifier, []byte(s.getServerIP().To4()))
	response = append(response, OptionEnd)

	// NAK must always be broadcast
	broadcastAddr := &net.UDPAddr{IP: net.IPv4bcast, Port: 68}
	s.conn.WriteToUDP(response, broadcastAddr)
}

// getResponseAddr determines where to send the DHCP response.
// Per RFC 2131: if ciaddr is 0.0.0.0, broadcast to 255.255.255.255:68.
// Otherwise unicast to ciaddr:68.
func (s *Server) getResponseAddr(data []byte, remoteAddr *net.UDPAddr) *net.UDPAddr {
	if len(data) < 28 {
		return remoteAddr
	}

	// ciaddr is at bytes 24-27
	ciaddr := net.IP(data[24:28])
	if ciaddr.Equal(net.IPv4zero) || ciaddr.IsUnspecified() {
		// Client has no IP yet — broadcast
		return &net.UDPAddr{IP: net.IPv4bcast, Port: 68}
	}

	// Client already has an IP — unicast
	return &net.UDPAddr{IP: ciaddr, Port: 68}
}

func (s *Server) loadLeases() {
	leases, err := s.db.GetActiveLeases()
	if err != nil {
		return
	}

	s.leaseMutex.Lock()
	defer s.leaseMutex.Unlock()

	for _, lease := range leases {
		s.leases[lease.MAC] = &lease
		s.usedIPs[lease.IP] = lease.MAC
	}
}

func (s *Server) loadStaticBindings() {
	bindings, err := s.db.GetStaticBindings()
	if err != nil {
		return
	}

	for _, binding := range bindings {
		s.staticBindings[binding.MAC] = binding
	}
}

func (s *Server) cleanupLeases() {
	ticker := time.NewTicker(1 * time.Minute)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			s.leaseMutex.Lock()
			now := time.Now().Unix()
			for mac, lease := range s.leases {
				if lease.ExpiresAt < now {
					delete(s.leases, mac)
					delete(s.usedIPs, lease.IP)
					log.Printf("DHCP: Lease expired for %s (%s)", mac, lease.IP)
				}
			}
			s.leaseMutex.Unlock()
		case <-s.stopChan:
			return
		}
	}
}

func (s *Server) GetLeases() []db.DHCPLease {
	s.leaseMutex.RLock()
	defer s.leaseMutex.RUnlock()

	leases := make([]db.DHCPLease, 0, len(s.leases))
	for _, lease := range s.leases {
		leases = append(leases, *lease)
	}
	return leases
}

func (s *Server) CreateStaticBinding(mac, ip, hostname, description string) error {
	binding := db.DHCPStaticBinding{
		MAC:         mac,
		IP:          ip,
		Hostname:    hostname,
		Description: description,
		Enabled:     true,
	}

	err := s.db.Create(&binding).Error
	if err == nil {
		s.staticBindings[mac] = binding
	}
	return err
}

func (s *Server) DeleteStaticBinding(id uint) error {
	// Get binding first to remove from cache
	var binding db.DHCPStaticBinding
	s.db.First(&binding, id)
	delete(s.staticBindings, binding.MAC)

	return s.db.Delete(&db.DHCPStaticBinding{}, id).Error
}

func (s *Server) GetStaticBindings() ([]db.DHCPStaticBinding, error) {
	return s.db.GetStaticBindings()
}

// Helper functions

func parseMessageType(data []byte) byte {
	// DHCP message type is option 53
	for i := 240; i < len(data)-1; i++ {
		if data[i] == OptionMessageType && i+2 < len(data) {
			return data[i+2]
		}
		if data[i] == OptionEnd {
			break
		}
	}
	return 0
}

func parseRequestedIP(data []byte) net.IP {
	// DHCP requested IP is option 50
	for i := 240; i < len(data)-1; i++ {
		if data[i] == OptionRequestedIP && i+5 < len(data) && data[i+1] == 4 {
			return net.IP(data[i+2 : i+6])
		}
		if data[i] == OptionEnd {
			break
		}
	}
	return nil
}

func parseHostname(data []byte) string {
	// DHCP hostname is option 12
	for i := 240; i < len(data)-1; i++ {
		if data[i] == OptionHostname && i+2 < len(data) {
			length := int(data[i+1])
			if i+2+length <= len(data) {
				return string(data[i+2 : i+2+length])
			}
		}
		if data[i] == OptionEnd {
			break
		}
	}
	return ""
}

func formatMAC(mac []byte) string {
	return fmt.Sprintf("%02x:%02x:%02x:%02x:%02x:%02x",
		mac[0], mac[1], mac[2], mac[3], mac[4], mac[5])
}

func buildDHCPMessage(msgType byte, request []byte, ip string, cfg *config.DHCPConfig) []byte {
	response := make([]byte, 240)

	// Copy operation, hardware type, hardware address length, hops
	response[0] = 2 // BOOTREPLY
	response[1] = request[1] // hardware type
	response[2] = request[2] // hardware address length

	// Copy transaction ID
	copy(response[4:8], request[4:8])

	// Set offered IP
	ipBytes := net.ParseIP(ip).To4()
	copy(response[16:20], ipBytes)

	// Set server IP
	serverIP := net.ParseIP(cfg.Gateway).To4()
	copy(response[128:132], serverIP)

	// Copy client MAC
	copy(response[28:34], request[28:34])

	// DHCP magic cookie
	response[236] = 99
	response[237] = 130
	response[238] = 83
	response[239] = 99

	// Add message type option
	response = append(response, OptionMessageType)
	response = append(response, 1) // length
	response = append(response, msgType)

	return response
}

func appendOption(data []byte, option byte, value []byte) []byte {
	data = append(data, option)
	data = append(data, byte(len(value)))
	data = append(data, value...)
	return data
}

// IP 地址管理工具函数
func IPInRange(ip, start, end string) bool {
	ipAddr := net.ParseIP(ip)
	startAddr := net.ParseIP(start)
	endAddr := net.ParseIP(end)

	if ipAddr == nil || startAddr == nil || endAddr == nil {
		return false
	}

	ipBytes := ipAddr.To4()
	startBytes := startAddr.To4()
	endBytes := endAddr.To4()

	if ipBytes == nil || startBytes == nil || endBytes == nil {
		return false
	}

	ipInt := uint32(ipBytes[0])<<24 | uint32(ipBytes[1])<<16 | uint32(ipBytes[2])<<8 | uint32(ipBytes[3])
	startInt := uint32(startBytes[0])<<24 | uint32(startBytes[1])<<16 | uint32(startBytes[2])<<8 | uint32(startBytes[3])
	endInt := uint32(endBytes[0])<<24 | uint32(endBytes[1])<<16 | uint32(endBytes[2])<<8 | uint32(endBytes[3])

	return ipInt >= startInt && ipInt <= endInt
}

// newBroadcastUDPConn creates a UDP listener on the given host:port with SO_BROADCAST enabled.
// This is required for DHCP because responses to clients without an IP must be sent to
// the broadcast address (255.255.255.255:68).
func newBroadcastUDPConn(host string, port int) (*net.UDPConn, error) {
	// Create a raw UDP socket so we can set SO_BROADCAST
	sock, err := unix.Socket(unix.AF_INET, unix.SOCK_DGRAM, unix.IPPROTO_UDP)
	if err != nil {
		return nil, fmt.Errorf("failed to create socket: %v", err)
	}

	// Enable SO_BROADCAST so we can send to 255.255.255.255
	if err := unix.SetsockoptInt(sock, unix.SOL_SOCKET, unix.SO_BROADCAST, 1); err != nil {
		unix.Close(sock)
		return nil, fmt.Errorf("failed to set SO_BROADCAST: %v", err)
	}

	// Bind to the address
	sa := &unix.SockaddrInet4{Port: port}
	copy(sa.Addr[:], net.ParseIP(host).To4())
	if err := unix.Bind(sock, sa); err != nil {
		unix.Close(sock)
		return nil, fmt.Errorf("failed to bind: %v", err)
	}

	// Wrap the raw socket in a net.UDPConn
	file := os.NewFile(uintptr(sock), fmt.Sprintf("udp-%s-%d", host, port))
	conn, err := net.FileConn(file)
	if err != nil {
		unix.Close(sock)
		return nil, fmt.Errorf("failed to wrap socket: %v", err)
	}

	udpConn, ok := conn.(*net.UDPConn)
	if !ok {
		conn.Close()
		return nil, fmt.Errorf("not a UDP connection")
	}

	return udpConn, nil
}

// EvictClient removes a client's lease, forcing them to get a new IP
func (s *Server) EvictClient(mac string) error {
	s.leaseMutex.Lock()
	defer s.leaseMutex.Unlock()

	lease, exists := s.leases[mac]
	if !exists {
		return fmt.Errorf("lease not found for MAC %s", mac)
	}

	// Remove from in-memory lease map
	delete(s.leases, mac)
	delete(s.usedIPs, lease.IP)

	// Delete from database
	if err := s.db.Where("mac = ?", mac).Delete(&db.DHCPLease{}).Error; err != nil {
		return fmt.Errorf("failed to delete lease from database: %v", err)
	}

	log.Printf("DHCP: Evicted client %s (%s)", mac, lease.IP)
	return nil
}