diff --git a/net/ipgate/domains.go b/net/ipgate/domains.go index d6e8465..865b3e7 100644 --- a/net/ipgate/domains.go +++ b/net/ipgate/domains.go @@ -3,6 +3,7 @@ package ipgate import ( "context" "net/netip" + "sync" "sync/atomic" "time" @@ -10,13 +11,21 @@ import ( "github.com/therootcompany/golib/net/ipcohort" ) -const domainSetRefreshInterval = 5 * time.Minute +const DefaultDomainSetRefreshInterval = 5 * time.Minute -type DomainSet struct { +type domainSources struct { staticPrefixes []string domains []string - resolved atomic.Pointer[map[string][]string] - cohort atomic.Pointer[ipcohort.Cohort] +} + +type DomainSet struct { + ctx context.Context + cancel context.CancelFunc + sources atomic.Pointer[domainSources] + resolved atomic.Pointer[map[string][]string] + cohort atomic.Pointer[ipcohort.Cohort] + start sync.Once + interval time.Duration } func EmptyDomainSet() *DomainSet { @@ -28,11 +37,16 @@ func EmptyDomainSet() *DomainSet { // NewDomainSet creates a DomainSet from pre-parsed inputs. // staticPrefixes is a list of CIDRs or bare IPs. // domains is a list of hostnames to resolve periodically. -func NewDomainSet(ctx context.Context, staticPrefixes []string, domains []string) *DomainSet { - ds := &DomainSet{ - staticPrefixes: staticPrefixes, - domains: domains, +func NewDomainSet(ctx context.Context, staticPrefixes []string, domains []string, interval time.Duration) *DomainSet { + if interval <= 0 { + interval = DefaultDomainSetRefreshInterval } + setCtx, cancel := context.WithCancel(ctx) + ds := &DomainSet{ctx: setCtx, cancel: cancel, interval: interval} + ds.sources.Store(&domainSources{ + staticPrefixes: append([]string(nil), staticPrefixes...), + domains: append([]string(nil), domains...), + }) ds.cohort.Store(&ipcohort.Cohort{}) @@ -43,11 +57,36 @@ func NewDomainSet(ctx context.Context, staticPrefixes []string, domains []string log().Info("domain set loaded", "static", commaify(len(staticPrefixes)), "domains", commaify(len(domains))) - go ds.refreshLoop(ctx) + if len(domains) > 0 { + ds.startRefreshLoop() + } return ds } +func (ds *DomainSet) Replace(staticPrefixes, domains []string) { + ds.sources.Store(&domainSources{ + staticPrefixes: append([]string(nil), staticPrefixes...), + domains: append([]string(nil), domains...), + }) + if len(domains) > 0 { + ds.startRefreshLoop() + } + ds.resolveDomains(ds.ctx) + ds.rebuildCohort() +} + +func (ds *DomainSet) startRefreshLoop() { + ds.start.Do(func() { go ds.refreshLoop(ds.interval) }) +} + +func (ds *DomainSet) Close() error { + if ds != nil && ds.cancel != nil { + ds.cancel() + } + return nil +} + func (ds *DomainSet) Contains(addr netip.Addr) bool { cohort := ds.cohort.Load() if cohort == nil { @@ -57,34 +96,35 @@ func (ds *DomainSet) Contains(addr netip.Addr) bool { return cohort.ContainsAddr(addr) } -func (ds *DomainSet) refreshLoop(ctx context.Context) { - ds.resolveDomains(ctx) +func (ds *DomainSet) refreshLoop(interval time.Duration) { + ds.resolveDomains(ds.ctx) ds.rebuildCohort() - ticker := time.NewTicker(domainSetRefreshInterval) + ticker := time.NewTicker(interval) defer ticker.Stop() for { select { - case <-ctx.Done(): + case <-ds.ctx.Done(): return case <-ticker.C: - ds.resolveDomains(ctx) + ds.resolveDomains(ds.ctx) ds.rebuildCohort() } } } func (ds *DomainSet) resolveDomains(ctx context.Context) { - if len(ds.domains) == 0 { + sources := ds.sources.Load() + if sources == nil || len(sources.domains) == 0 { return } prev := *ds.resolved.Load() - next := make(map[string][]string, len(ds.domains)) + next := make(map[string][]string, len(sources.domains)) resolver := dnsresolver.New() - for _, domain := range ds.domains { + for _, domain := range sources.domains { ips, _, err := resolver.LookupIP(ctx, domain) if err != nil || len(ips) == 0 { @@ -109,7 +149,11 @@ func (ds *DomainSet) resolveDomains(ctx context.Context) { func (ds *DomainSet) rebuildCohort() { var all []string - all = append(all, ds.staticPrefixes...) + sources := ds.sources.Load() + if sources == nil { + return + } + all = append(all, sources.staticPrefixes...) resolved := *ds.resolved.Load() for _, addrs := range resolved { diff --git a/net/ipgate/ipprefix.go b/net/ipgate/ipprefix.go index fa469f3..53923e3 100644 --- a/net/ipgate/ipprefix.go +++ b/net/ipgate/ipprefix.go @@ -12,9 +12,11 @@ import ( "github.com/therootcompany/golib/net/ipcohort" ) -const prefixSetRefreshInterval = 47 * time.Minute +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] @@ -26,24 +28,37 @@ func EmptyPrefixSet() *PrefixSet { return ps } -func NewPrefixSet(ctx context.Context, repoURL, dataPath string, files []string) (*PrefixSet, error) { +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("ipgate: create data dir: %w", err) } + if interval <= 0 { + interval = DefaultPrefixSetRefreshInterval + } + setCtx, cancel := context.WithCancel(ctx) repo := gitshallow.New(repoURL, dataPath, 0, "") ps := &PrefixSet{ - repo: repo, - files: files, + ctx: setCtx, + cancel: cancel, + repo: repo, + files: files, } ps.cohort.Store(&ipcohort.Cohort{}) - go ps.refreshLoop(ctx) + 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 { @@ -95,20 +110,20 @@ func filesPresent(paths []string) bool { return true } -func (ps *PrefixSet) refreshLoop(ctx context.Context) { - if err := ps.reload(ctx); err != nil { +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(prefixSetRefreshInterval) + ticker := time.NewTicker(interval) defer ticker.Stop() for { select { - case <-ctx.Done(): + case <-ps.ctx.Done(): return case <-ticker.C: - if err := ps.reload(ctx); err != nil { + if err := ps.reload(ps.ctx); err != nil { log().Warn("prefix set reload failed", "err", err) } }