kdhcp/iptools/pool.go

package iptools

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

	"git.wntrmute.dev/kyle/goutils/assert"
	"git.wntrmute.dev/kyle/goutils/log"
)

const DefaultExpiry = 168 * time.Hour

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

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]*LeaseInfo{},
	}

	return p, nil
}

func (p *Pool) sort() {
	p.Available = SortLeases(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(t time.Time, waitFor time.Duration) *LeaseInfo {
	p.lock()
	defer p.unlock()

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

func (p *Pool) peek(t time.Time, waitFor time.Duration) *LeaseInfo {
	if len(p.Available) == 0 {
		return nil
	}

	lease := p.Available[0]
	p.Available = p.Available[1:]
	lease.ResetExpiry(t, waitFor)

	p.Limbo[lease.Addr] = lease

	return lease
}

func (p *Pool) accept(addr netip.Addr) error {
	assert.Bool(p.Active[addr] == nil, fmt.Sprintf("limbo address %s is already active: %#v", addr, *p.Active[addr]))

	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 {
	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 *LeaseInfo, t time.Time) {
	p.Active[lease.Addr].ResetExpiry(t, p.Expiry)
}