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) } } } }