kdhcp/iptools/pool.go

package iptools

import (
	"bytes"
	"fmt"
	"net/netip"
	"sync"
	"time"

	"git.wntrmute.dev/kyle/goutils/log"
	"git.wntrmute.dev/kyle/kdhcp/dhcp"
	"git.wntrmute.dev/kyle/kdhcp/leases"
)

const DefaultExpiry = 168 * time.Hour

type Pool struct {
	Name       string                      `yaml:"name" json:"name"`
	Expiry     time.Duration               `yaml:"expiry" json:"expiry"`
	Available  []*leases.Info              `yaml:"available" json:"available"`
	Active     map[netip.Addr]*leases.Info `yaml:"active" json:"active"`
	Limbo      map[netip.Addr]*leases.Info `yaml:"limbo" json:"limbo"`             // leases that are currently being offered
	NoHostName bool                        `yaml:"no_hostname" json:"no_hostname"` // don't set the hostname
	mtx        *sync.Mutex
}

func (p *Pool) checkMaps() {
	if p.Active == nil {
		p.Active = map[netip.Addr]*leases.Info{}
	}

	if p.Limbo == nil {
		p.Limbo = map[netip.Addr]*leases.Info{}
	}
}

func (p *Pool) lock() {
	if p.mtx == nil {
		p.mtx = &sync.Mutex{}
	}

	p.mtx.Lock()
}

func (p *Pool) unlock() {
	if p.mtx == nil {
		panic("pool: attempted to unlock uninitialized mutex")
	}

	p.mtx.Unlock()
}

func NewPool(name string, r *Range) (*Pool, error) {
	if r.Expiry == 0 {
		r.Expiry = DefaultExpiry
	}

	p := &Pool{
		Name:       name,
		Expiry:     r.Expiry,
		NoHostName: r.NoHostName,
		Available:  enumerateRange(name, r, true),
		Limbo:      map[netip.Addr]*leases.Info{},
	}

	return p, nil
}

func (p *Pool) sort() {
	p.Available = leases.Sort(p.Available)
}

func (p *Pool) Sort() {
	p.lock()
	defer p.unlock()

	p.sort()
}

func (p *Pool) IsAddressAvailable() bool {
	return len(p.Available) > 0
}

// Peek returns the first available address from the pool and moves it
// from available to limbo. When the client is given the address, Accept
// should be called on the address to move it from limbo to active.
func (p *Pool) Peek(req *dhcp.Packet, t time.Time, waitFor time.Duration) *leases.Info {
	p.lock()
	defer p.unlock()

	li := p.peek(req, t, waitFor)
	return li
}

func (p *Pool) peek(req *dhcp.Packet, t time.Time, waitFor time.Duration) *leases.Info {
	p.checkMaps()

	for _, li := range p.Active {
		if bytes.Equal(req.HardwareAddress, li.HardwareAddress) {
			log.Debugf("returning existing lease to %x of %s", req.HardwareAddress, li.Addr)
			return li
		}
	}

	for _, li := range p.Limbo {
		if bytes.Equal(req.HardwareAddress, li.HardwareAddress) {
			log.Debugf("returning existing offer to %x of %s", req.HardwareAddress, li.Addr)
			return li
		}
	}

	if len(p.Available) == 0 {
		return nil
	}

	lease := p.Available[0]
	p.Available = p.Available[1:]

	lease.HardwareAddress = req.HardwareAddress
	lease.ResetExpiry(t, waitFor)
	p.Limbo[lease.Addr] = lease

	return lease
}

func (p *Pool) accept(addr netip.Addr) error {
	p.checkMaps()

	if active, ok := p.Active[addr]; ok {
		return fmt.Errorf("limbo address %s is already active: %#v", addr, *active)
	}

	p.Active[addr] = p.Limbo[addr]
	delete(p.Limbo, addr)
	return nil
}

// Accept will move an address out limbo and mark it as active.
func (p *Pool) Accept(addr netip.Addr, t time.Time) error {
	p.lock()
	defer p.unlock()

	p.update(t)
	p.accept(addr)

	return nil
}

// Update checks expirations on leases, moving any expired leases back
// to the available pool and removing any limbo leases that are expired.
func (p *Pool) Update(t time.Time) bool {
	p.lock()
	defer p.unlock()

	return p.update(t)
}

func (p *Pool) update(t time.Time) bool {
	p.checkMaps()

	var updated bool

	for k, v := range p.Active {
		if v.IsExpired(t) {
			updated = true
			log.Infof("expiring active address %s for %x", v.Addr, v.HardwareAddress)
			delete(p.Active, k)
			v = v.Reset()
			p.Available = append(p.Available, v)
		}
	}

	for k, v := range p.Limbo {
		if v.IsExpired(t) {
			updated = true
			log.Infof("expiring limbo address %s for %x", v.Addr, v.HardwareAddress)
			delete(p.Limbo, k)
			v = v.Reset()
			p.Available = append(p.Available, v)
		}
	}

	p.sort()
	return updated
}

func (p *Pool) Renew(lease *leases.Info, t time.Time) {
	p.checkMaps()

	p.Active[lease.Addr].ResetExpiry(t, p.Expiry)
}