pkg/atlas: Add ATLAS_IP2LOCATION config option

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