golib/net/ippolicy/ipprefix.go

package ippolicy

import (
	"context"
	"fmt"
	"net/netip"
	"os"
	"sync/atomic"
	"time"

	"github.com/therootcompany/golib/net/gitshallow"
	"github.com/therootcompany/golib/net/ipcohort"
)

const DefaultPrefixSetRefreshInterval = time.Hour + 57*time.Minute + 13*time.Second

type PrefixSet struct {
	ctx    context.Context
	cancel context.CancelFunc
	repo   *gitshallow.Repo
	files  []string
	cohort atomic.Pointer[ipcohort.Cohort]
}

func EmptyPrefixSet() *PrefixSet {
	ps := &PrefixSet{}
	ps.cohort.Store(&ipcohort.Cohort{})
	return ps
}

func NewPrefixSet(ctx context.Context, repoURL, dataPath string, files []string, interval time.Duration) (*PrefixSet, error) {
	if err := os.MkdirAll(dataPath, 0o755); err != nil {
		return nil, fmt.Errorf("ippolicy: create data dir: %w", err)
	}

	if interval <= 0 {
		interval = DefaultPrefixSetRefreshInterval
	}
	setCtx, cancel := context.WithCancel(ctx)
	repo := gitshallow.New(repoURL, dataPath, 0, "")

	ps := &PrefixSet{
		ctx:    setCtx,
		cancel: cancel,
		repo:   repo,
		files:  files,
	}
	ps.cohort.Store(&ipcohort.Cohort{})

	go ps.refreshLoop(interval)

	return ps, nil
}

func (ps *PrefixSet) Close() error {
	if ps != nil && ps.cancel != nil {
		ps.cancel()
	}
	return nil
}

func (ps *PrefixSet) Contains(addr netip.Addr) bool {
	cohort := ps.cohort.Load()
	if cohort == nil {
		cohort = &ipcohort.Cohort{}
		ps.cohort.CompareAndSwap(nil, cohort)
	}
	return cohort.ContainsAddr(addr)
}

func (ps *PrefixSet) reload(ctx context.Context) error {
	updated, err := ps.repo.Fetch(ctx)
	if err != nil {
		return err
	}

	paths := make([]string, len(ps.files))
	for i, f := range ps.files {
		paths[i] = ps.repo.FilePath(f)
	}

	if cachedCohortValid(updated, ps.cohort.Load(), paths) {
		return nil
	}

	cohort, err := ipcohort.LoadFiles(paths...)
	if err != nil {
		return fmt.Errorf("load files: %w", err)
	}

	ps.cohort.Store(cohort)

	log().Info("prefix set loaded", "entries", commaify(cohort.Size()))
	return nil
}

// cachedCohortValid reports whether the current cohort can be reused
// without reloading from disk: the repo was not updated, the cohort is
// non-nil and non-empty, and all data files are present.
func cachedCohortValid(updated bool, cohort *ipcohort.Cohort, paths []string) bool {
	return !updated && cohort != nil && cohort.Size() > 0 && filesPresent(paths)
}

func filesPresent(paths []string) bool {
	for _, path := range paths {
		if _, err := os.Stat(path); err != nil {
			return false
		}
	}
	return true
}

func (ps *PrefixSet) refreshLoop(interval time.Duration) {
	if err := ps.reload(ps.ctx); err != nil {
		log().Warn("prefix set initial load (will retry)", "err", err)
	}

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ps.ctx.Done():
			return
		case <-ticker.C:
			if err := ps.reload(ps.ctx); err != nil {
				log().Warn("prefix set reload failed", "err", err)
			}
		}
	}
}