glider/service/dhcpd/pool.go

package dhcpd

import (
	"bytes"
	"errors"
	"math/rand"
	"net"
	"time"
)

// Pool is a dhcp pool.
type Pool struct {
	items []*item
}

// NewPool returns a new dhcp ip pool.
func NewPool(lease time.Duration, start, end net.IP) (*Pool, error) {
	if start == nil || end == nil {
		return nil, errors.New("start ip or end ip is wrong/nil, please check your config")
	}

	s, e := ip2num(start.To4()), ip2num(end.To4())
	if e < s {
		return nil, errors.New("start ip larger than end ip")
	}

	items := make([]*item, 0, e-s+1)
	for n := s; n <= e; n++ {
		items = append(items, &item{lease: lease, ip: num2ip(n)})
	}
	rand.Seed(time.Now().Unix())
	return &Pool{items: items}, nil
}

// AssignIP assigns an ip to mac from dhco pool.
func (p *Pool) AssignIP(mac net.HardwareAddr) (net.IP, error) {
	var ip net.IP
	for _, item := range p.items {
		if bytes.Equal(mac, item.mac) {
			return item.ip, nil
		}
	}

	idx := rand.Intn(len(p.items))
	for _, item := range p.items[idx:] {
		if ip = item.take(mac); ip != nil {
			return ip, nil
		}
	}

	for _, item := range p.items {
		if ip = item.take(mac); ip != nil {
			return ip, nil
		}
	}
	return nil, errors.New("no more ip can be assigned")
}

type item struct {
	taken bool
	ip    net.IP
	lease time.Duration
	mac   net.HardwareAddr
}

func (i *item) take(addr net.HardwareAddr) net.IP {
	if !i.taken {
		i.taken = true
		go func() {
			timer := time.NewTimer(i.lease)
			<-timer.C
			i.mac = nil
			i.taken = false
		}()
		i.mac = addr
		return i.ip
	}
	return nil
}

func ip2num(ip net.IP) uint32 {
	n := uint32(ip[0])<<24 + uint32(ip[1])<<16
	return n + uint32(ip[2])<<8 + uint32(ip[3])
}

func num2ip(n uint32) net.IP {
	return []byte{byte(n >> 24), byte(n >> 16), byte(n >> 8), byte(n)}
}