f-93: Merges onelease behavior in onerange

Signed-off-by: Aleix Ramírez <[email protected]>

appliances/VRouter/DHCP4v2/dhcpcore-onelease/plugins/onerange/plugin.go

@@ -14,6 +14,8 @@ import (
 	"errors"
 	"fmt"
 	"net"
+	"regexp"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -24,6 +26,7 @@ import (
 	"github.com/coredhcp/coredhcp/plugins/allocators"
 	"github.com/coredhcp/coredhcp/plugins/allocators/bitmap"
 	"github.com/insomniacslk/dhcp/dhcpv4"
+	"github.com/spf13/pflag"
 )
 
 var log = logger.GetLogger("plugins/onerange")
@@ -46,11 +49,15 @@ type PluginState struct {
 	// Rough lock for the whole plugin, we'll get better performance once we use leasestorage
 	sync.Mutex
 	// Recordsv4 holds a MAC -> IP address and lease time mapping
-	Recordsv4   map[string]*Record
-	LeaseTime   time.Duration
-	excludedIPs []net.IP
-	leasedb     *sql.DB
-	allocator   allocators.Allocator
+	Recordsv4    map[string]*Record
+	LeaseTime    time.Duration
+	ExcludedIPs  []net.IP
+	leasedb      *sql.DB
+	allocator    allocators.Allocator
+	enableMAC2IP bool
+	MACPrefix    [2]byte
+	rangeStartIP net.IP
+	rangeEndIP   net.IP
 }
 
 // Handler4 handles DHCPv4 packets for the range plugin
@@ -62,11 +69,47 @@ func (p *PluginState) Handler4(req, resp *dhcpv4.DHCPv4) (*dhcpv4.DHCPv4, bool)
 	if !ok {
 		// Allocating new address since there isn't one allocated
 		log.Printf("MAC address %s is new, leasing new IPv4 address", req.ClientHWAddr.String())
-		ip, err := p.allocator.Allocate(net.IPNet{})
+		// Check if the MAC address should be mapped to a specific IP address
+		ipToAllocate := net.IPNet{}
+		macPrefixMatches := false
+		if p.enableMAC2IP {
+			macAddress := req.ClientHWAddr
+			ipFromMAC, ok, err := p.checkMACPrefix(macAddress)
+			if err != nil {
+				log.Errorf("MAC2IP lease failed for mac %v: %v", macAddress.String(), err)
+				// return nack to the client
+				resp.Options.Update(dhcpv4.OptMessageType(dhcpv4.MessageTypeNak))
+				return resp, true
+			}
+			if ok {
+				macPrefixMatches = true
+				//propose the least 4 bytes of the mac address for allocating an IP address
+				ipToAllocate = net.IPNet{IP: ipFromMAC}
+				log.Infof("MAC %s matches the prefix %x, trying to allocate IP %s...", macAddress.String(), p.MACPrefix, ipToAllocate.IP.String())
+			} else {
+				log.Infof("MAC %s does not match the prefix %x, providing conventional lease", macAddress.String(), p.MACPrefix)
+			}
+		}
+		// The allocator will try to allocate the given IP address, but if it exits, it will return a different one (an available one)
+		ip, err := p.allocator.Allocate(ipToAllocate)
 		if err != nil {
 			log.Errorf("Could not allocate IP for MAC %s: %v", req.ClientHWAddr.String(), err)
-			return nil, true
+			resp.Options.Update(dhcpv4.OptMessageType(dhcpv4.MessageTypeNak))
+			// return nack to the client
+			resp.Options.Update(dhcpv4.OptMessageType(dhcpv4.MessageTypeNak))
+			return resp, true
 		}
+
+		// if the MAC address is mapped to an IP address, check if the allocated IP address matches the requested one, if not, revert the allocation and return
+		if p.enableMAC2IP && macPrefixMatches && !ip.IP.Equal(ipToAllocate.IP) {
+			log.Warnf("Allocated IP %s for MAC %s does not match the requested IP %s", ip.IP.String(), req.ClientHWAddr.String(), ipToAllocate.IP.String())
+			//revert the allocation of the undesired IP
+			p.allocator.Free(ip)
+			// return nack to the client
+			resp.Options.Update(dhcpv4.OptMessageType(dhcpv4.MessageTypeNak))
+			return resp, true
+		}
+
 		rec := Record{
 			IP:       ip.IP.To4(),
 			expires:  int(time.Now().Add(p.LeaseTime).Unix()),
@@ -93,7 +136,42 @@ func (p *PluginState) Handler4(req, resp *dhcpv4.DHCPv4) (*dhcpv4.DHCPv4, bool)
 	resp.YourIPAddr = record.IP
 	resp.Options.Update(dhcpv4.OptIPAddressLeaseTime(p.LeaseTime.Round(time.Second)))
 	log.Printf("found IP address %s for MAC %s", record.IP, req.ClientHWAddr.String())
-	return resp, false
+	return resp, true
+}
+
+func (p *PluginState) checkIPInRange(ip net.IP) bool {
+	return binary.BigEndian.Uint32(ip.To4()) >= binary.BigEndian.Uint32(p.rangeStartIP.To4()) &&
+		binary.BigEndian.Uint32(ip.To4()) <= binary.BigEndian.Uint32(p.rangeEndIP.To4())
+}
+
+func (p *PluginState) checkMACPrefix(mac net.HardwareAddr) (net.IP, bool, error) {
+	// verify that the MAC address is valid
+	if _, err := net.ParseMAC(mac.String()); err != nil {
+		return nil, false, fmt.Errorf("invalid MAC address: %v", mac)
+	}
+
+	// verify that the two first bytes equal to macPrefix, if not, return
+	if mac[0] != p.MACPrefix[0] || mac[1] != p.MACPrefix[1] {
+		return nil, false, nil
+	}
+
+	// retrieve the IP address from the MAC address
+	// the IP address is the last 4 bytes of the MAC address
+	ip := net.IPv4(mac[2], mac[3], mac[4], mac[5])
+
+	// check if the ip is in the excluded list
+	for _, excluded := range p.ExcludedIPs {
+		if ip.Equal(excluded) {
+			return nil, false, fmt.Errorf("excluded IP %v", ip)
+		}
+	}
+
+	// check if the ip is in the lease range
+	if !p.checkIPInRange(ip) {
+		return nil, false, fmt.Errorf("IP %v is not in the range", ip)
+	}
+
+	return ip, true, nil
 }
 
 func setupRange(args ...string) (handler.Handler4, error) {
@@ -103,25 +181,19 @@ func setupRange(args ...string) (handler.Handler4, error) {
 	)
 
 	if len(args) < 4 {
-		return nil, fmt.Errorf("invalid number of arguments, want: 4 (file name, start IP, end IP, lease time), got: %d", len(args))
+		return nil, fmt.Errorf("invalid number of arguments, want at least: 4 (file name, start IP, end IP, lease time), got: %d", len(args))
 	}
 	filename := args[0]
 	if filename == "" {
 		return nil, errors.New("file name cannot be empty")
 	}
-	ipRangeStart := net.ParseIP(args[1])
-	if ipRangeStart.To4() == nil {
-		return nil, fmt.Errorf("invalid IPv4 address: %v", args[1])
-	}
-	ipRangeEnd := net.ParseIP(args[2])
-	if ipRangeEnd.To4() == nil {
-		return nil, fmt.Errorf("invalid IPv4 address: %v", args[2])
-	}
-	if binary.BigEndian.Uint32(ipRangeStart.To4()) >= binary.BigEndian.Uint32(ipRangeEnd.To4()) {
-		return nil, errors.New("start of IP range has to be lower than the end of an IP range")
+
+	p.rangeStartIP, p.rangeEndIP, err = parseIPRange(args[1], args[2])
+	if err != nil {
+		return nil, fmt.Errorf("invalid IP range: %v", err)
 	}
 
-	p.allocator, err = bitmap.NewIPv4Allocator(ipRangeStart, ipRangeEnd)
+	p.allocator, err = bitmap.NewIPv4Allocator(p.rangeStartIP, p.rangeEndIP)
 	if err != nil {
 		return nil, fmt.Errorf("could not create an allocator: %w", err)
 	}
@@ -131,28 +203,35 @@ func setupRange(args ...string) (handler.Handler4, error) {
 		return nil, fmt.Errorf("invalid lease duration: %v", args[3])
 	}
 
-	// parse a list of excluded IPs as fifth argument
-	if len(args) > 4 {
-		excludedIPs := args[4]
-		for _, ip := range strings.Split(excludedIPs, ",") {
-			excluded := net.ParseIP(strings.TrimSpace(ip))
-			if excluded.To4() == nil {
-				return nil, fmt.Errorf("invalid excluded IP address: %v", ip)
-			}
-			p.excludedIPs = append(p.excludedIPs, excluded)
+	optionalArgs := args[4:]
+
+	var excludedIPs string
+	var macPrefix string
+
+	pflag.StringVar(&excludedIPs, "excluded-ips", "", "Comma-separated list of excluded IP addresses")
+	pflag.BoolVar(&p.enableMAC2IP, "mac2ip", false, "Enables MAC to IP address mapping")
+	pflag.StringVar(&macPrefix, "mac-prefix", "02:00", "2-byte MAC prefix for MAC to IP address mapping. Defaults to [02:00]")
+
+	pflag.CommandLine.Parse(optionalArgs)
+
+	if p.enableMAC2IP && macPrefix != "" {
+		p.MACPrefix, err = parseMACPrefix(macPrefix)
+		if err != nil {
+			return nil, fmt.Errorf("invalid MAC prefix: %v", macPrefix)
 		}
 	}
 
-	// check that the excluded IPs belongs to the range and pre-allocate them
-	// preallocated IPs will not be stored in the lease database, but they will be kept at the allocator level
-	// who is the responsible of managing the IP availability
-	for _, excluded := range p.excludedIPs {
-		if binary.BigEndian.Uint32(excluded.To4()) < binary.BigEndian.Uint32(ipRangeStart.To4()) ||
-			binary.BigEndian.Uint32(excluded.To4()) > binary.BigEndian.Uint32(ipRangeEnd.To4()) {
-			return nil, fmt.Errorf("excluded IP %v is not in the range %v-%v", excluded, ipRangeStart, ipRangeEnd)
+	if excludedIPs != "" {
+		p.ExcludedIPs, err = parseExcludedIPs(excludedIPs)
+		if err != nil {
+			return nil, fmt.Errorf("invalid excluded IPs: %v", excludedIPs)
 		}
-		if _, err := p.allocator.Allocate(net.IPNet{IP: excluded}); err != nil {
-			return nil, fmt.Errorf("could not pre-allocate excluded IP %v: %w", excluded, err)
+		//check if excluded IPs are in the range and pre-allocate them
+		for _, excluded := range p.ExcludedIPs {
+			p.checkIPInRange(excluded)
+			if _, err := p.allocator.Allocate(net.IPNet{IP: excluded}); err != nil {
+				return nil, fmt.Errorf("could not pre-allocate excluded IP %v: %w", excluded, err)
+			}
 		}
 	}
 
@@ -178,3 +257,51 @@ func setupRange(args ...string) (handler.Handler4, error) {
 
 	return p.Handler4, nil
 }
+
+func parseIPRange(startIP, endIP string) (net.IP, net.IP, error) {
+	ipRangeStart := net.ParseIP(startIP)
+	if ipRangeStart.To4() == nil {
+		return nil, nil, fmt.Errorf("invalid IPv4 address: %v", startIP)
+	}
+	ipRangeEnd := net.ParseIP(endIP)
+	if ipRangeEnd.To4() == nil {
+		return nil, nil, fmt.Errorf("invalid IPv4 address: %v", endIP)
+	}
+	if binary.BigEndian.Uint32(ipRangeStart.To4()) >= binary.BigEndian.Uint32(ipRangeEnd.To4()) {
+		return nil, nil, errors.New("start of IP range has to be lower than the end of an IP range")
+	}
+	return ipRangeStart, ipRangeEnd, nil
+}
+
+func parseExcludedIPs(ipList string) ([]net.IP, error) {
+	excludedIPs := []net.IP{}
+	for _, ip := range strings.Split(ipList, ",") {
+		excluded := net.ParseIP(strings.TrimSpace(ip))
+		if excluded.To4() == nil {
+			return nil, fmt.Errorf("invalid excluded IP address: %v", ip)
+		}
+		excludedIPs = append(excludedIPs, excluded)
+	}
+	return excludedIPs, nil
+}
+
+func parseMACPrefix(prefix string) ([2]byte, error) {
+	regex := `^[0-9A-Fa-f]{2}:[0-9A-Fa-f]{2}$`
+	matched, err := regexp.MatchString(regex, prefix)
+	if err != nil {
+		return [2]byte{}, fmt.Errorf("error matching regex: %v", err)
+	}
+	if !matched {
+		return [2]byte{}, fmt.Errorf("invalid MAC prefix format: %s", prefix)
+	}
+	parts := strings.Split(prefix, ":")
+	macByte0, err := strconv.ParseUint(parts[0], 16, 8)
+	if err != nil {
+		return [2]byte{}, fmt.Errorf("invalid MAC prefix byte [0]: %v", err)
+	}
+	macByte1, err := strconv.ParseUint(parts[1], 16, 8)
+	if err != nil {
+		return [2]byte{}, fmt.Errorf("invalid MAC prefix byte [1]: %v", err)
+	}
+	return [2]byte{byte(macByte0), byte(macByte1)}, nil
+}