demo/DHCP/internal/service/dhcp.go

package service

import (
	"context"
	"dhcp/api/request"
	"dhcp/code"
	"dhcp/internal/backend"
	"dhcp/internal/handler/proxy"
	"dhcp/internal/ip"
	"dhcp/internal/server"
	"dhcp/result"
	"fmt"
	"github.com/go-logr/logr"
	"github.com/insomniacslk/dhcp/dhcpv4/server4"
	"net"
	"net/netip"
)

var DHCPServiceMap = map[string]*DHCPService{}
var DHCPServiceInfoMap = map[string]request.DHCPService{}

type DhcpConfig struct {
	Mode              string `yaml:"mode"`                 // 模式
	BindAddr          string `yaml:"bind_addr"`            // 绑定地址
	BindInterface     string `yaml:"bind_interface"`       //	绑定接口
	IpForPacket       string `yaml:"ip_for_packet"`        // 包的ip
	SyslogIP          string `yaml:"syslog_ip"`            // syslog ip
	TftpIP            string `yaml:"tftp_ip"`              // tftp ip
	HttpIpxeBinaryURL string `yaml:"http_ipxe_binary_url"` // http ipxe二进制url
}

func StartDHCPService(ctx context.Context, log logr.Logger, req request.DHCPService) result.Result {
	dhcp := DhcpConfig{
		Mode:          req.Mode,
		BindAddr:      "0.0.0.0:67",
		BindInterface: req.BindInterface,
		IpForPacket:   ip.DetectPublicIPv4(""),
		SyslogIP:      ip.DetectPublicIPv4(""),
		TftpIP:        ip.DetectPublicIPv4(":69"),
	}
	var dhcpService = &DHCPService{}
	dhcpService.StartDHCP(ctx, dhcp, log)
	DHCPServiceMap[req.BindInterface] = dhcpService
	req.Status = true
	DHCPServiceInfoMap[req.BindInterface] = req
	return result.Result{
		Status: true,
		Code:   code.SUCCESS,
		Msg:    "启动成功",
	}
}

func StopDHCPService(ctx context.Context, log logr.Logger, req request.DHCPService) result.Result {
	if dhcpService, ok := DHCPServiceMap[req.BindInterface]; ok {
		err := dhcpService.server.Close()
		if err != nil {
			return result.Result{
				Status: false,
				Code:   code.SystemException,
				Msg:    "关闭失败",
			}
		}
		delete(DHCPServiceMap, req.BindInterface)
		DHCPServiceInfo := DHCPServiceInfoMap[req.BindInterface]
		DHCPServiceInfo.Status = false
		DHCPServiceInfo.Info = "自主关闭"
		DHCPServiceInfoMap[req.BindInterface] = DHCPServiceInfo
	}
	return result.Result{
		Status: true,
		Code:   code.SUCCESS,
		Msg:    "关闭成功",
	}
}

func DeleteDHCPService(ctx context.Context, log logr.Logger, req request.DHCPService) result.Result {
	if dhcpService, ok := DHCPServiceMap[req.BindInterface]; ok {
		_ = dhcpService.server.Close()
		delete(DHCPServiceMap, req.BindInterface)
		delete(DHCPServiceInfoMap, req.BindInterface)
	} else {
		return result.Result{
			Status: false,
			Code:   code.SystemException,
			Msg:    "删除失败",
		}
	}
	return result.Result{
		Status: true,
		Code:   code.SUCCESS,
		Msg:    "删除成功",
	}
}

type DHCPService struct {
	server *server.DHCP
}

func (dhcp *DHCPService) StartDHCP(ctx context.Context, d DhcpConfig, log logr.Logger) {
	dh, err := dhcpHandler(ctx, &d, log)
	if err != nil {
		log.Error(err, "failed to create dhcpServer listener")
		panic(fmt.Errorf("failed to create dhcpServer listener: %w", err))
	}
	log.Info("starting dhcpServer server", "bind_addr", d.BindAddr)
	go func(ctx context.Context) {
		bindAddr, err := netip.ParseAddrPort(d.BindAddr)
		if err != nil {
			panic(fmt.Errorf("invalid tftp address for dhcpServer server: %w", err))
		}
		conn, err := server4.NewIPv4UDPConn(d.BindInterface, net.UDPAddrFromAddrPort(bindAddr))
		if err != nil {
			log.Info("failed to create dhcpServer listener", "error", err)
			panic(err)
		}
		defer conn.Close()
		dhcp.server = &server.DHCP{Logger: log, Conn: conn, Handlers: []server.Handler{dh}}
		err = dhcp.server.Serve(ctx)

		if err != nil {
			delete(DHCPServiceMap, d.BindInterface)
			DHCPServiceInfo := DHCPServiceInfoMap[d.BindInterface]
			DHCPServiceInfo.Status = false
			DHCPServiceInfo.Info = err.Error()
			DHCPServiceInfoMap[d.BindInterface] = DHCPServiceInfo
		}
	}(ctx)
}

func (dhcp *DHCPService) StopDHCP(ctx context.Context) {
	dhcp.server.Close()
}

func dhcpHandler(ctx context.Context, c *DhcpConfig, log logr.Logger) (server.Handler, error) {
	// 1. create the handler
	// 2. create the macAndIPXE
	// 3. add the macAndIPXE to the handler
	pktIP, err := netip.ParseAddr(c.IpForPacket)
	if err != nil {
		return nil, fmt.Errorf("invalid bind address: %w", err)
	}
	br := backend.NewBackend()
	dh := &proxy.Handler{
		Backend: br,
		IPAddr:  pktIP,
		Log:     log,
	}
	return dh, nil
}