diff options
Diffstat (limited to 'pkg')
-rw-r--r-- | pkg/atlas/config.go | 7 | ||||
-rw-r--r-- | pkg/atlas/server.go | 20 | ||||
-rw-r--r-- | pkg/atlas/util.go | 53 |
3 files changed, 80 insertions, 0 deletions
diff --git a/pkg/atlas/config.go b/pkg/atlas/config.go index a0a8108..4f411bc 100644 --- a/pkg/atlas/config.go +++ b/pkg/atlas/config.go @@ -164,6 +164,13 @@ type Config struct { // of top-level names to URLs. Web string `env:"ATLAS_WEB="` + // The path to the IP2Location database, which should contain at least the + // country and region fields. The database must not be modified while atlas + // is running, but it can be replaced (and a reload can be triggered with + // SIGHUP). If not provided, geolocation-dependent features like server + // regions will not be enabled. + IP2Location string `env:"ATLAS_IP2LOCATION"` + // For sd-notify. NotifySocket string `env:"NOTIFY_SOCKET"` diff --git a/pkg/atlas/server.go b/pkg/atlas/server.go index 3978911..2e90af2 100644 --- a/pkg/atlas/server.go +++ b/pkg/atlas/server.go @@ -242,6 +242,18 @@ func NewServer(c *Config) (*Server, error) { } else { return nil, fmt.Errorf("initialize main menu promos: %w", err) } + if ip2l, err := configureIP2Location(c); err == nil { + if ip2l != nil { + s.reload = append(s.reload, func() { + if err := ip2l.Load(""); err != nil { + s.Logger.Err(err).Msg("failed to reload ip2location database") + } + }) + s.API0.LookupIP = ip2l.LookupFields + } + } else { + return nil, fmt.Errorf("initialize ip2location: %w", err) + } s.MetricsSecret = c.MetricsSecret @@ -611,6 +623,14 @@ func configureMainMenuPromos(c *Config) (func(*http.Request) api0.MainMenuPromos } } +func configureIP2Location(c *Config) (*ip2locationMgr, error) { + if c.IP2Location == "" { + return nil, nil + } + mgr := new(ip2locationMgr) + return mgr, mgr.Load(c.IP2Location) +} + // Run runs the server, shutting it down gracefully when ctx is canceled, then // waiting indefinitely for it to exit. It must only ever be called once, and // the server is useless afterwards. diff --git a/pkg/atlas/util.go b/pkg/atlas/util.go index 83feb20..d5f71e1 100644 --- a/pkg/atlas/util.go +++ b/pkg/atlas/util.go @@ -1,13 +1,66 @@ package atlas import ( + "fmt" "io" "net/http" + "net/netip" + "os" "sync" + "github.com/pg9182/ip2x/ip2location" "github.com/rs/zerolog" ) +// ip2locationMgr wraps a file-backed IP2Location database. +type ip2locationMgr struct { + file *os.File + db *ip2location.DB + mu sync.RWMutex +} + +// Load replaces the currently loaded database with the specified file. If name +// is empty, the existing database, if any, is reopened. +func (m *ip2locationMgr) Load(name string) error { + if name == "" { + m.mu.RLock() + if m.file == nil { + return fmt.Errorf("no ip2location database loaded") + } + name = m.file.Name() + m.mu.RUnlock() + } + + f, err := os.Open(name) + if err != nil { + return err + } + + db, err := ip2location.New(f) + if err != nil { + f.Close() + return err + } + + m.mu.Lock() + defer m.mu.Unlock() + + m.file.Close() + m.file = f + m.db = db + return nil +} + +// LookupFields calls (*ip2location.DB).LookupFields if a database is loaded. +func (m *ip2locationMgr) LookupFields(ip netip.Addr, mask ip2location.Field) (ip2location.Record, error) { + m.mu.RLock() + defer m.mu.RUnlock() + if m.db == nil { + return ip2location.Record{}, fmt.Errorf("no ip2location database loaded") + } + return m.db.LookupFields(ip, mask) +} + type zerologWriterLevel struct { w io.Writer // or zerolog.LevelWriter l zerolog.Level |