github.com/zooyer/miskit@v1.0.71/dhcp/main.go (about)

     1  /**
     2   * @Author: zzy
     3   * @Email: zhangzhongyuan@didiglobal.com
     4   * @Description:
     5   * @File: main.go
     6   * @Package: dhcp4
     7   * @Version: 1.0.0
     8   * @Date: 2022/11/9 23:36
     9   */
    10  
    11  package main
    12  
    13  import (
    14  	"bytes"
    15  	"encoding/binary"
    16  	"encoding/json"
    17  	"fmt"
    18  	"github.com/krolaw/dhcp4"
    19  	"log"
    20  	"math/rand"
    21  	"net"
    22  	"time"
    23  )
    24  
    25  // ExampleHandler using DHCP with a single network interface device
    26  func ExampleHandler() {
    27  	serverIP := net.IP{192, 168, 1, 1}
    28  	handler := &DHCPHandler{
    29  		ip:            serverIP,
    30  		leaseDuration: 2 * time.Hour,
    31  		start:         net.IP{192, 168, 1, 2},
    32  		leaseRange:    50,
    33  		leases:        make(map[int]lease, 10),
    34  		options: dhcp4.Options{
    35  			dhcp4.OptionSubnetMask:       []byte{255, 255, 255, 0},
    36  			dhcp4.OptionRouter:           []byte(serverIP), // Presuming Server is also your router
    37  			dhcp4.OptionDomainNameServer: []byte(serverIP), // Presuming Server is also your DNS server
    38  		},
    39  	}
    40  	log.Fatal(dhcp4.ListenAndServe(handler))
    41  	// log.Fatal(dhcp4.Serve(dhcp4.NewUDP4BoundListener("eth0",":67"), handler)) // Select interface on multi interface device - just linux for now
    42  	// log.Fatal(dhcp4.Serve(dhcp4.NewUDP4FilterListener("en0",":67"), handler)) // Work around for other OSes
    43  }
    44  
    45  type lease struct {
    46  	nic    string    // Client's CHAddr
    47  	expiry time.Time // When the lease expires
    48  }
    49  
    50  type DHCPHandler struct {
    51  	ip            net.IP        // Server IP to use
    52  	options       dhcp4.Options // Options to send to DHCP Clients
    53  	start         net.IP        // Start of IP range to distribute
    54  	leaseRange    int           // Number of IPs to distribute (starting from start)
    55  	leaseDuration time.Duration // Lease period
    56  	leases        map[int]lease // Map to keep track of leases
    57  }
    58  
    59  func (h *DHCPHandler) ServeDHCP(p dhcp4.Packet, msgType dhcp4.MessageType, options dhcp4.Options) (d dhcp4.Packet) {
    60  	var opt = make(map[string]interface{})
    61  	for key, val := range options {
    62  		switch key {
    63  		case dhcp4.OptionDHCPMessageType:
    64  			u, err := binary.ReadUvarint(bytes.NewReader(val))
    65  			if err != nil {
    66  				opt[key.String()+"Error"] = err.Error()
    67  				continue
    68  			}
    69  			opt[key.String()] = u
    70  		case dhcp4.OptionParameterRequestList:
    71  			fallthrough
    72  		case dhcp4.OptionRequestedIPAddress:
    73  			opt[key.String()] = net.IP{val[0], val[1], val[2], val[3]}.String()
    74  		default:
    75  			opt[key.String()] = string(val)
    76  		}
    77  	}
    78  	data, err := json.Marshal(opt)
    79  	if err != nil {
    80  		panic(err)
    81  	}
    82  
    83  	fmt.Println(string(data))
    84  
    85  	switch msgType {
    86  
    87  	case dhcp4.Discover:
    88  		free, nic := -1, p.CHAddr().String()
    89  		for i, v := range h.leases { // Find previous lease
    90  			if v.nic == nic {
    91  				free = i
    92  				goto reply
    93  			}
    94  		}
    95  		if free = h.freeLease(); free == -1 {
    96  			return
    97  		}
    98  	reply:
    99  		return dhcp4.ReplyPacket(p, dhcp4.Offer, h.ip, dhcp4.IPAdd(h.start, free), h.leaseDuration,
   100  			h.options.SelectOrderOrAll(options[dhcp4.OptionParameterRequestList]))
   101  
   102  	case dhcp4.Request:
   103  		if server, ok := options[dhcp4.OptionServerIdentifier]; ok && !net.IP(server).Equal(h.ip) {
   104  			return nil // Message not for this dhcp4 server
   105  		}
   106  		reqIP := net.IP(options[dhcp4.OptionRequestedIPAddress])
   107  		if reqIP == nil {
   108  			reqIP = net.IP(p.CIAddr())
   109  		}
   110  
   111  		if len(reqIP) == 4 && !reqIP.Equal(net.IPv4zero) {
   112  			if leaseNum := dhcp4.IPRange(h.start, reqIP) - 1; leaseNum >= 0 && leaseNum < h.leaseRange {
   113  				if l, exists := h.leases[leaseNum]; !exists || l.nic == p.CHAddr().String() {
   114  					h.leases[leaseNum] = lease{nic: p.CHAddr().String(), expiry: time.Now().Add(h.leaseDuration)}
   115  					return dhcp4.ReplyPacket(p, dhcp4.ACK, h.ip, reqIP, h.leaseDuration,
   116  						h.options.SelectOrderOrAll(options[dhcp4.OptionParameterRequestList]))
   117  				}
   118  			}
   119  		}
   120  		return dhcp4.ReplyPacket(p, dhcp4.NAK, h.ip, nil, 0, nil)
   121  
   122  	case dhcp4.Release, dhcp4.Decline:
   123  		nic := p.CHAddr().String()
   124  		for i, v := range h.leases {
   125  			if v.nic == nic {
   126  				delete(h.leases, i)
   127  				break
   128  			}
   129  		}
   130  	}
   131  	return nil
   132  }
   133  
   134  func (h *DHCPHandler) freeLease() int {
   135  	now := time.Now()
   136  	b := rand.Intn(h.leaseRange) // Try random first
   137  	for _, v := range [][]int{[]int{b, h.leaseRange}, []int{0, b}} {
   138  		for i := v[0]; i < v[1]; i++ {
   139  			if l, ok := h.leases[i]; !ok || l.expiry.Before(now) {
   140  				return i
   141  			}
   142  		}
   143  	}
   144  	return -1
   145  }
   146  
   147  func main() {
   148  	ExampleHandler()
   149  }