refactor: move GeoLite to its own package (#1647)

This commit is contained in:
Alessandro (Ale) Segala
2026-08-05 13:12:19 -07:00
committed by GitHub
parent 1c9233c236
commit 22e3909c6c
19 changed files with 1376 additions and 339 deletions

View File

@@ -117,6 +117,9 @@ func Bootstrap(ctx context.Context) error {
return fmt.Errorf("failed to register scheduled jobs: %w", err)
}
// Refresh the GeoLite database (this is cached per each replica)
services = append(services, svc.geoLiteModule.Run)
// The scheduler must wait on the actor host being ready, since jobs invoke actors
services = append(services, actorsReady.Await(scheduler.Run))
}

View File

@@ -14,10 +14,6 @@ func registerScheduledJobs(ctx context.Context, db *gorm.DB, svc *services, sche
if err != nil {
return fmt.Errorf("failed to register LDAP jobs in scheduler: %w", err)
}
err = scheduler.RegisterGeoLiteUpdateJobs(ctx, svc.geoLiteService)
if err != nil {
return fmt.Errorf("failed to register GeoLite DB update service: %w", err)
}
err = scheduler.RegisterDbCleanupJobs(ctx, db)
if err != nil {
return fmt.Errorf("failed to register DB cleanup jobs in scheduler: %w", err)

View File

@@ -13,6 +13,7 @@ import (
"github.com/pocket-id/pocket-id/backend/internal/devicelogin"
"github.com/pocket-id/pocket-id/backend/internal/email"
"github.com/pocket-id/pocket-id/backend/internal/emailverification"
"github.com/pocket-id/pocket-id/backend/internal/geolite"
"github.com/pocket-id/pocket-id/backend/internal/job"
"github.com/pocket-id/pocket-id/backend/internal/oidc"
"github.com/pocket-id/pocket-id/backend/internal/onetimeaccess"
@@ -27,7 +28,7 @@ type services struct {
appConfigService *appconfig.AppConfigService
appImagesService *service.AppImagesService
emailModule *email.Module
geoLiteService *service.GeoLiteService
geoLiteModule *geolite.Module
auditLogService *service.AuditLogService
jwtService *service.JwtService
scimService *service.ScimService
@@ -79,8 +80,17 @@ func initServices(
return nil, fmt.Errorf("failed to create email module: %w", err)
}
svc.geoLiteService = service.NewGeoLiteService(httpClient)
svc.auditLogService = service.NewAuditLogService(db, svc.emailModule, svc.geoLiteService, svc.appConfigService)
svc.geoLiteModule, err = geolite.New(ctx, geolite.Dependencies{
HTTPClient: httpClient,
DBPath: common.EnvConfig.GeoLiteDBPath,
DownloadURL: common.EnvConfig.GeoLiteDBUrl,
LicenseKey: common.EnvConfig.MaxMindLicenseKey,
})
if err != nil {
return nil, fmt.Errorf("failed to create GeoLite module: %w", err)
}
svc.auditLogService = service.NewAuditLogService(db, svc.emailModule, svc.geoLiteModule, svc.appConfigService)
svc.jwtService, err = service.NewJwtService(ctx, db, instanceID)
if err != nil {
return nil, fmt.Errorf("failed to create JWT service: %w", err)
@@ -104,7 +114,7 @@ func initServices(
Signer: svc.jwtService,
Reauth: svc.webauthnModule,
AuditLog: svc.auditLogService,
IPLocator: svc.geoLiteService,
IPLocator: svc.geoLiteModule,
AppConfig: svc.appConfigService,
})
if err != nil {

View File

@@ -28,7 +28,7 @@ type AuditLogger interface {
}
type IPLocationResolver interface {
GetLocationByIP(ipAddress string) (country, city string, err error)
GetLocationByIP(ctx context.Context, ipAddress string) (country string, city string, err error)
}
type AppConfigProvider interface {

View File

@@ -116,7 +116,7 @@ func (s *Service) Inspect(ctx context.Context, code string) (VerificationInfo, e
return VerificationInfo{}, err
}
country, city, err := s.ipLocator.GetLocationByIP(result.IPAddress)
country, city, err := s.ipLocator.GetLocationByIP(ctx, result.IPAddress)
if err != nil {
slog.WarnContext(ctx, "Failed to get device login request IP location", slog.String("ip", result.IPAddress), slog.Any("error", err))
}

View File

@@ -87,7 +87,7 @@ type fakeIPLocationResolver struct {
err error
}
func (f *fakeIPLocationResolver) GetLocationByIP(string) (string, string, error) {
func (f *fakeIPLocationResolver) GetLocationByIP(context.Context, string) (string, string, error) {
return f.country, f.city, f.err
}

View File

@@ -0,0 +1,173 @@
package geolite
import (
"archive/tar"
"bufio"
"compress/gzip"
"context"
"errors"
"fmt"
"io"
"net/http"
"os"
"path"
"path/filepath"
"strings"
"github.com/oschwald/maxminddb-golang/v2"
)
const (
// databaseFileName is the name of the database inside the archive published by MaxMind
databaseFileName = "GeoLite2-City.mmdb"
// maxDatabaseSize is the largest (decompressed) database we accept
maxDatabaseSize = 300 << 20 // 300 MB
)
// downloadDatabase downloads the GeoLite2 City database and puts it at targetPath
// The database is streamed to a temporary file next to the target and moved into place only once it has been verified, atomically
func downloadDatabase(ctx context.Context, httpClient *http.Client, downloadURL string, licenseKey string, targetPath string) error {
// When downloadURL contains a "%s" placeholder, it is replaced with the license key
if strings.Contains(downloadURL, "%s") {
downloadURL = fmt.Sprintf(downloadURL, licenseKey)
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil)
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
res, err := httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to download database: %w", err)
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return fmt.Errorf("failed to download database, received HTTP %d", res.StatusCode)
}
err = writeDatabase(res.Body, targetPath)
if err != nil {
return err
}
return nil
}
// writeDatabase extracts the database from a downloaded body and puts it at targetPath
func writeDatabase(body io.Reader, targetPath string) error {
baseDir := filepath.Dir(targetPath)
err := os.MkdirAll(baseDir, 0700)
if err != nil {
return fmt.Errorf("failed to create the database directory: %w", err)
}
tmpFile, err := os.CreateTemp(baseDir, "geolite.*.mmdb.tmp")
if err != nil {
return fmt.Errorf("failed to create temporary database file: %w", err)
}
tmpName := tmpFile.Name()
// Remove the temporary file unless it has been moved into place
moved := false
defer func() {
tmpFile.Close()
if !moved {
os.Remove(tmpName)
}
}()
err = extractDatabase(body, tmpFile)
if err != nil {
return fmt.Errorf("failed to extract database: %w", err)
}
err = tmpFile.Close()
if err != nil {
return fmt.Errorf("failed to write database file: %w", err)
}
// Make sure the database isn't corrupted before it replaces the one currently in place
db, err := maxminddb.Open(tmpName)
if err != nil {
return fmt.Errorf("failed to open downloaded database: %w", err)
}
_ = db.Close()
err = os.Rename(tmpName, targetPath)
if err != nil {
return fmt.Errorf("failed to replace database file: %w", err)
}
moved = true
return nil
}
// extractDatabase copies the raw MaxMind DB file out of a downloaded body and into dst
// The body is either the gzipped tarball published by MaxMind, or the database file itself
func extractDatabase(body io.Reader, dst io.Writer) error {
// A buffered reader lets the gzip magic number be checked without consuming it
reader := bufio.NewReader(body)
magic, err := reader.Peek(2)
if err != nil {
return fmt.Errorf("failed to read magic number: %w", err)
}
// If the body doesn't start with the gzip magic number, assume it's a plain database file
// Gosec returns false positive for "G602: slice index out of range"
//nolint:gosec
if magic[0] != 0x1f || magic[1] != 0x8b {
return copyDatabase(dst, reader)
}
gzr, err := gzip.NewReader(reader)
if err != nil {
return fmt.Errorf("failed to create gzip reader: %w", err)
}
defer gzr.Close()
tarReader := tar.NewReader(gzr)
for {
header, err := tarReader.Next()
if errors.Is(err, io.EOF) {
break
} else if err != nil {
return fmt.Errorf("failed to read tar archive: %w", err)
}
// The archive contains the database in a versioned folder, alongside other files such as the license
if header.Typeflag != tar.TypeReg || path.Base(header.Name) != databaseFileName {
continue
}
err = copyDatabase(dst, tarReader)
if err != nil {
return err
}
return nil
}
return errors.New(databaseFileName + " not found in archive")
}
// copyDatabase streams the database into dst, refusing anything larger than maxDatabaseSize
func copyDatabase(dst io.Writer, src io.Reader) error {
return copyDatabaseWithLimit(dst, src, maxDatabaseSize)
}
// copyDatabaseWithLimit is copyDatabase with an explicit limit, which lets tests exercise the limit without streaming hundreds of megabytes
func copyDatabaseWithLimit(dst io.Writer, src io.Reader, limit int64) error {
// Copy one byte more than the limit, so content that is exactly at the limit can be told apart from content that exceeds it
written, err := io.Copy(dst, io.LimitReader(src, limit+1))
if err != nil {
return fmt.Errorf("failed to read database: %w", err)
}
if written > limit {
return errors.New("database size exceeds maximum allowed limit")
}
return nil
}

View File

@@ -0,0 +1,171 @@
package geolite
import (
"bytes"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"testing"
"github.com/stretchr/testify/require"
)
func TestExtractDatabase(t *testing.T) {
database := readTestDatabase(t)
t.Run("gzipped tarball", func(t *testing.T) {
archive := buildTarGzForTest(t, map[string][]byte{
"GeoLite2-City_20260101/COPYRIGHT.txt": []byte("copyright"),
"GeoLite2-City_20260101/LICENSE.txt": []byte("license"),
"GeoLite2-City_20260101/" + databaseFileName: database,
})
dst := &bytes.Buffer{}
err := extractDatabase(bytes.NewReader(archive), dst)
require.NoError(t, err)
require.Equal(t, database, dst.Bytes())
})
t.Run("plain database file", func(t *testing.T) {
// A custom GEOLITE_DB_URL may serve the database uncompressed
dst := &bytes.Buffer{}
err := extractDatabase(bytes.NewReader(database), dst)
require.NoError(t, err)
require.Equal(t, database, dst.Bytes())
})
t.Run("tarball without the database", func(t *testing.T) {
archive := buildTarGzForTest(t, map[string][]byte{
"GeoLite2-City_20260101/COPYRIGHT.txt": []byte("copyright"),
})
err := extractDatabase(bytes.NewReader(archive), io.Discard)
require.Error(t, err)
require.ErrorContains(t, err, "not found in archive")
})
t.Run("truncated gzip stream", func(t *testing.T) {
archive := buildTarGzForTest(t, map[string][]byte{"GeoLite2-City_20260101/" + databaseFileName: database})
err := extractDatabase(bytes.NewReader(archive[:len(archive)/2]), io.Discard)
require.Error(t, err)
})
t.Run("empty body", func(t *testing.T) {
err := extractDatabase(strings.NewReader(""), io.Discard)
require.Error(t, err)
require.ErrorContains(t, err, "failed to read magic number")
})
}
func TestCopyDatabase(t *testing.T) {
t.Run("exactly at the limit", func(t *testing.T) {
dst := &bytes.Buffer{}
err := copyDatabaseWithLimit(dst, bytes.NewReader(bytes.Repeat([]byte{0x01}, 1024)), 1024)
require.NoError(t, err)
require.Equal(t, 1024, dst.Len())
})
t.Run("over the limit", func(t *testing.T) {
err := copyDatabaseWithLimit(io.Discard, endlessReader{}, 1024)
require.Error(t, err)
require.ErrorContains(t, err, "exceeds maximum allowed limit")
})
}
func TestDownloadDatabase(t *testing.T) {
database := readTestDatabase(t)
archive := buildTarGzForTest(t, map[string][]byte{"GeoLite2-City_20260101/" + databaseFileName: database})
t.Run("success", func(t *testing.T) {
httpClient, transport := newDownloadClientForTest(archive)
targetPath := filepath.Join(t.TempDir(), "GeoLite2-City.mmdb")
err := downloadDatabase(t.Context(), httpClient, testDownloadURL, "", targetPath)
require.NoError(t, err)
require.Equal(t, int32(1), transport.requests.Load())
written, err := os.ReadFile(targetPath)
require.NoError(t, err)
require.Equal(t, database, written)
})
t.Run("creates the target directory", func(t *testing.T) {
httpClient, _ := newDownloadClientForTest(archive)
targetPath := filepath.Join(t.TempDir(), "nested", "data", "GeoLite2-City.mmdb")
err := downloadDatabase(t.Context(), httpClient, testDownloadURL, "", targetPath)
require.NoError(t, err)
require.FileExists(t, targetPath)
})
t.Run("license key placeholder", func(t *testing.T) {
// The default MaxMind URL carries the license key, which is filled in at download time
var requestedURL string
httpClient := &http.Client{
Transport: roundTripperFunc(func(req *http.Request) (*http.Response, error) {
requestedURL = req.URL.String()
return &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(bytes.NewReader(archive)),
Header: make(http.Header),
}, nil
}),
}
err := downloadDatabase(t.Context(), httpClient, "https://example.com/download?license_key=%s", "secret-key", filepath.Join(t.TempDir(), "db.mmdb"))
require.NoError(t, err)
require.Equal(t, "https://example.com/download?license_key=secret-key", requestedURL)
})
t.Run("non-200 response", func(t *testing.T) {
httpClient, transport := newDownloadClientForTest(nil)
transport.statusCode = http.StatusUnauthorized
targetPath := filepath.Join(t.TempDir(), "GeoLite2-City.mmdb")
err := downloadDatabase(t.Context(), httpClient, testDownloadURL, "", targetPath)
require.Error(t, err)
require.ErrorContains(t, err, "received HTTP 401")
require.NoFileExists(t, targetPath)
})
t.Run("corrupted database leaves the existing file in place", func(t *testing.T) {
corrupted := buildTarGzForTest(t, map[string][]byte{
"GeoLite2-City_20260101/" + databaseFileName: []byte("not a database"),
})
httpClient, _ := newDownloadClientForTest(corrupted)
dir := t.TempDir()
targetPath := filepath.Join(dir, "GeoLite2-City.mmdb")
require.NoError(t, os.WriteFile(targetPath, database, 0600))
err := downloadDatabase(t.Context(), httpClient, testDownloadURL, "", targetPath)
require.Error(t, err)
require.ErrorContains(t, err, "failed to open downloaded database")
// The database that was already there is untouched, and no temporary file is left behind
existing, err := os.ReadFile(targetPath)
require.NoError(t, err)
require.Equal(t, database, existing)
entries, err := os.ReadDir(dir)
require.NoError(t, err)
require.Len(t, entries, 1)
})
}
// endlessReader returns an unbounded stream of zeroes, to exercise the download size limit
type endlessReader struct{}
func (endlessReader) Read(p []byte) (int, error) {
clear(p)
return len(p), nil
}
type roundTripperFunc func(req *http.Request) (*http.Response, error)
func (f roundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
return f(req)
}

View File

@@ -0,0 +1,74 @@
package geolite
import (
"context"
"errors"
"log/slog"
"net/http"
"github.com/pocket-id/pocket-id/backend/internal/common"
)
type Dependencies struct {
HTTPClient *http.Client
// DBPath is where the GeoLite2 City database is kept
// It is local to each replica: the database is a cache of a public artifact rather than state, so it doesn't need to be shared, though pointing several replicas at the same mount works too
DBPath string
// DownloadURL is the URL the GeoLite2 City database is downloaded from
// When it contains a "%s" placeholder, it is replaced with LicenseKey
DownloadURL string
// LicenseKey is the MaxMind license key
LicenseKey string
}
type Module struct {
service *Service
refresher *refresher
}
func New(ctx context.Context, deps Dependencies) (*Module, error) {
if deps.DBPath == "" {
return nil, errors.New("the GeoLite database path is empty")
}
log := slog.With(slog.String("scope", "geolite"))
service := newService(log, deps.DBPath)
// Map the database that is already on disk, if any, so lookups work before the first refresh
// A database that can't be read isn't fatal: lookups just return no location until a readable one shows up
err := service.load(ctx)
if err != nil {
log.WarnContext(ctx, "Failed to load the GeoLite2 City database", slog.String("path", deps.DBPath), slog.Any("error", err))
}
disabled := deps.LicenseKey == "" && deps.DownloadURL == common.MaxMindGeoLiteCityUrl
if disabled {
// Warn the user, and disable the periodic refresh
// The database can still be supplied by hand at DBPath, which is what air-gapped deployments do
log.Warn("MAXMIND_LICENSE_KEY environment variable is empty: the GeoLite2 City database won't be updated")
}
return &Module{
service: service,
refresher: &refresher{
log: log,
service: service,
httpClient: deps.HTTPClient,
downloadURL: deps.DownloadURL,
licenseKey: deps.LicenseKey,
disabled: disabled,
},
}, nil
}
// GetLocationByIP returns the country and city of the given IP address
func (m *Module) GetLocationByIP(ctx context.Context, ipAddress string) (country string, city string, err error) {
return m.service.GetLocationByIP(ctx, ipAddress)
}
// Run keeps the GeoLite2 City database up-to-date until the context is canceled
// It satisfies servicerunner.Service
func (m *Module) Run(ctx context.Context) error {
return m.refresher.Run(ctx)
}

View File

@@ -0,0 +1,218 @@
package geolite
import (
"context"
"errors"
"log/slog"
"math/rand/v2"
"net/http"
"os"
"path/filepath"
"time"
"github.com/fsnotify/fsnotify"
)
const (
// databaseMaxAge is how old the database on disk is allowed to get before it's downloaded again
databaseMaxAge = 14 * 24 * time.Hour
// refreshJitter is subtracted or added at random when scheduling a refresh
// On a shared mount it keeps several replicas from waking up together and racing to download the same file
refreshJitter = 30 * time.Minute
// refreshRetryInterval is how long the refresher waits before trying again after a failed download
refreshRetryInterval = time.Hour
// downloadTimeout bounds a single download of the database
downloadTimeout = 10 * time.Minute
// watcherDebounce is how long the watcher waits for the file to settle before reloading it
watcherDebounce = time.Second
)
// refresher keeps the database at dbPath up-to-date and reloads the service when the file changes
type refresher struct {
log *slog.Logger
service *Service
httpClient *http.Client
downloadURL string
licenseKey string
// disabled stops the periodic download, for deployments that have no way to reach the download URL
// The file is still watched, since it may be supplied by hand
disabled bool
// watching is closed once the watcher is established, if set
// It's a hook for tests, which need to know when a change to the file is guaranteed to be noticed
watching chan struct{}
}
// Run watches the database file and periodically refreshes it, until the context is canceled
// It satisfies servicerunner.Service
func (r *refresher) Run(ctx context.Context) error {
watcherDone := make(chan struct{})
go func() {
defer close(watcherDone)
r.watch(ctx)
}()
if !r.disabled {
r.refreshPeriodically(ctx)
} else {
<-ctx.Done()
}
<-watcherDone
return nil
}
// refreshPeriodically downloads the database whenever the one on disk has aged past databaseMaxAge
func (r *refresher) refreshPeriodically(ctx context.Context) {
for {
delay := r.timeUntilRefresh()
r.log.DebugContext(ctx, "Scheduled the next GeoLite2 City database refresh", slog.Duration("in", delay))
timer := time.NewTimer(delay)
select {
case <-ctx.Done():
timer.Stop()
return
case <-timer.C:
}
err := r.refresh(ctx)
if err != nil {
// The next attempt is scheduled by timeUntilRefresh, which still sees a missing or stale file and comes back after refreshRetryInterval
r.log.ErrorContext(ctx, "Failed to refresh the GeoLite2 City database, will try again later", slog.Any("error", err))
timer = time.NewTimer(refreshRetryInterval)
select {
case <-ctx.Done():
timer.Stop()
return
case <-timer.C:
}
}
}
}
// timeUntilRefresh reports how long to wait before the database on disk needs downloading again
//
// The age of the file is the only input, which is what makes this work across replicas without any coordination: on a shared mount, whichever replica gets there first refreshes the file and the others find it fresh and go back to sleep, and on separate disks each replica refreshes its own copy.
func (r *refresher) timeUntilRefresh() time.Duration {
info, err := os.Stat(r.service.dbPath)
if err != nil {
// Treat a database that is missing, or that can't be read, as due for a download
return 0
}
remaining := databaseMaxAge - time.Since(info.ModTime())
if remaining <= 0 {
return 0
}
return remaining + jitter()
}
// jitter returns a random offset within refreshJitter, so replicas sharing a mount don't wake up in lockstep
func jitter() time.Duration {
// #nosec G404 -- not used for anything security related
return time.Duration(rand.Int64N(int64(2*refreshJitter))) - refreshJitter
}
// refresh downloads the database and loads it
func (r *refresher) refresh(parentCtx context.Context) error {
r.log.InfoContext(parentCtx, "Refreshing the GeoLite2 City database")
ctx, cancel := context.WithTimeout(parentCtx, downloadTimeout)
defer cancel()
err := downloadDatabase(ctx, r.httpClient, r.downloadURL, r.licenseKey, r.service.dbPath)
if err != nil {
return err
}
r.log.InfoContext(parentCtx, "GeoLite2 City database successfully refreshed")
// The watcher would pick the new file up too, but loading it here makes the refresh complete on its own
return r.service.load(parentCtx)
}
// watch reloads the database whenever the file changes, so a database replaced by hand takes effect without a restart
func (r *refresher) watch(ctx context.Context) {
watcher, err := fsnotify.NewWatcher()
if err != nil {
r.log.ErrorContext(ctx, "Failed to create the GeoLite2 City database watcher, changes to the file will need a restart", slog.Any("error", err))
return
}
defer watcher.Close()
// The directory is watched rather than the file itself: a database is put in place by renaming over it, which leaves a watch on the old file pointing at an inode nothing writes to again
baseDir := filepath.Dir(r.service.dbPath)
err = os.MkdirAll(baseDir, 0700)
if err != nil {
r.log.ErrorContext(ctx, "Failed to create the GeoLite2 City database directory, changes to the file will need a restart", slog.Any("error", err))
return
}
err = watcher.Add(baseDir)
if err != nil {
r.log.ErrorContext(ctx, "Failed to watch the GeoLite2 City database directory, changes to the file will need a restart", slog.Any("error", err))
return
}
// Load once the watch is in place, to pick up a database that showed up while the watcher was being set up
// Without this, a file put there in that window would go unnoticed until the next refresh, and there is no next refresh when refreshes are disabled
err = r.service.load(ctx)
if err != nil {
r.log.ErrorContext(ctx, "Failed to load the GeoLite2 City database", slog.Any("error", err))
}
if r.watching != nil {
close(r.watching)
}
// The timer debounces the burst of events a write produces, so the database is loaded once the file has settled
reload := time.NewTimer(watcherDebounce)
if !reload.Stop() {
<-reload.C
}
defer reload.Stop()
for {
select {
case <-ctx.Done():
return
case event, ok := <-watcher.Events:
if !ok {
return
}
if filepath.Base(event.Name) != filepath.Base(r.service.dbPath) {
continue
}
if !event.Has(fsnotify.Create | fsnotify.Write | fsnotify.Rename | fsnotify.Remove) {
continue
}
r.log.DebugContext(ctx, "GeoLite2 City database change detected", slog.String("path", event.Name))
reload.Stop()
select {
case <-reload.C:
default:
}
reload.Reset(watcherDebounce)
case <-reload.C:
err := r.service.load(ctx)
if err != nil {
r.log.ErrorContext(ctx, "Failed to load the GeoLite2 City database after it changed", slog.Any("error", err))
}
case err, ok := <-watcher.Errors:
if !ok {
return
}
if !errors.Is(err, context.Canceled) {
r.log.ErrorContext(ctx, "GeoLite2 City database watcher error", slog.Any("error", err))
}
}
}
}

View File

@@ -0,0 +1,208 @@
package geolite
import (
"context"
"net/http"
"os"
"path/filepath"
"testing"
"time"
"github.com/stretchr/testify/require"
)
// newRefresherForTest returns a refresher wired to a service backed by a database file inside dir
func newRefresherForTest(t *testing.T, dir string, httpClient *http.Client) *refresher {
t.Helper()
svc := newServiceAtPathForTest(t, filepath.Join(dir, "GeoLite2-City.mmdb"))
return &refresher{
log: testLogger(),
service: svc,
httpClient: httpClient,
downloadURL: testDownloadURL,
watching: make(chan struct{}),
}
}
// runRefresherForTest starts the refresher and waits until its watcher is established, so a change made afterwards is guaranteed to be noticed
// The refresher is stopped when the test ends
func runRefresherForTest(t *testing.T, r *refresher) context.Context {
t.Helper()
ctx, cancel := context.WithCancel(t.Context())
done := make(chan error, 1)
go func() {
done <- r.Run(ctx)
}()
t.Cleanup(func() {
cancel()
require.NoError(t, <-done)
})
select {
case <-r.watching:
case <-time.After(30 * time.Second):
t.Fatal("timed out waiting for the database watcher to start")
}
return ctx
}
func TestRefresherTimeUntilRefresh(t *testing.T) {
t.Run("missing database", func(t *testing.T) {
r := newRefresherForTest(t, t.TempDir(), nil)
require.Zero(t, r.timeUntilRefresh())
})
t.Run("fresh database", func(t *testing.T) {
dir := t.TempDir()
r := newRefresherForTest(t, dir, nil)
writeDatabaseFileForTest(t, r.service.dbPath, readTestDatabase(t))
// The wait is the remaining lifetime of the file, give or take the jitter
delay := r.timeUntilRefresh()
require.InDelta(t, databaseMaxAge, delay, float64(refreshJitter+time.Minute))
require.Positive(t, delay)
})
t.Run("stale database", func(t *testing.T) {
dir := t.TempDir()
r := newRefresherForTest(t, dir, nil)
writeDatabaseFileForTest(t, r.service.dbPath, readTestDatabase(t))
aged := time.Now().Add(-databaseMaxAge - time.Hour)
err := os.Chtimes(r.service.dbPath, aged, aged)
require.NoError(t, err)
require.Zero(t, r.timeUntilRefresh())
})
}
func TestRefresherRefresh(t *testing.T) {
database := readTestDatabase(t)
archive := buildTarGzForTest(t, map[string][]byte{"GeoLite2-City_20260101/" + databaseFileName: database})
httpClient, transport := newDownloadClientForTest(archive)
r := newRefresherForTest(t, t.TempDir(), httpClient)
// Nothing to look up against before the first refresh
country, _, err := r.service.GetLocationByIP(t.Context(), "81.2.69.142")
require.NoError(t, err)
require.Empty(t, country)
err = r.refresh(t.Context())
require.NoError(t, err)
require.Equal(t, int32(1), transport.requests.Load())
// The database is on disk, and the service is serving from it without waiting for the watcher
require.FileExists(t, r.service.dbPath)
country, city, err := r.service.GetLocationByIP(t.Context(), "81.2.69.142")
require.NoError(t, err)
require.Equal(t, "United Kingdom", country)
require.Equal(t, "London", city)
// The refresh it just performed pushes the next one out by the full lifetime of the database
require.Positive(t, r.timeUntilRefresh())
}
func TestRefresherRefreshFailure(t *testing.T) {
httpClient, transport := newDownloadClientForTest(nil)
transport.statusCode = http.StatusInternalServerError
r := newRefresherForTest(t, t.TempDir(), httpClient)
err := r.refresh(t.Context())
require.Error(t, err)
require.ErrorContains(t, err, "received HTTP 500")
// No file is left behind, so the next attempt is still due right away
require.NoFileExists(t, r.service.dbPath)
require.Zero(t, r.timeUntilRefresh())
}
func TestRefresherSharedDirectorySkipsDownload(t *testing.T) {
// Replicas pointed at the same mount coordinate through the file itself: once one of them has refreshed it, the others find it fresh and don't download it again
database := readTestDatabase(t)
archive := buildTarGzForTest(t, map[string][]byte{"GeoLite2-City_20260101/" + databaseFileName: database})
httpClient, transport := newDownloadClientForTest(archive)
dir := t.TempDir()
first := newRefresherForTest(t, dir, httpClient)
second := newRefresherForTest(t, dir, httpClient)
require.Zero(t, first.timeUntilRefresh())
require.Zero(t, second.timeUntilRefresh())
err := first.refresh(t.Context())
require.NoError(t, err)
require.Equal(t, int32(1), transport.requests.Load())
// The second replica sees the file the first one wrote and goes back to sleep instead of downloading it again
require.Positive(t, second.timeUntilRefresh())
}
func TestRefresherWatchesForReplacedDatabase(t *testing.T) {
// Supplying a database by hand is how air-gapped deployments work, and it takes effect without a restart
r := newRefresherForTest(t, t.TempDir(), nil)
r.disabled = true
ctx := runRefresherForTest(t, r)
country, _, err := r.service.GetLocationByIP(ctx, "81.2.69.142")
require.NoError(t, err)
require.Empty(t, country)
writeDatabaseFileForTest(t, r.service.dbPath, readTestDatabase(t))
require.Eventually(t, func() bool {
country, _, err := r.service.GetLocationByIP(ctx, "81.2.69.142")
return err == nil && country == "United Kingdom"
}, 30*time.Second, 100*time.Millisecond, "the database put in place by hand was never picked up")
// Removing it stops lookups from resolving, rather than serving from a file that is gone
err = os.Remove(r.service.dbPath)
require.NoError(t, err)
require.Eventually(t, func() bool {
country, _, err := r.service.GetLocationByIP(ctx, "81.2.69.142")
return err == nil && country == ""
}, 30*time.Second, 100*time.Millisecond, "the removed database was still being served")
}
func TestRefresherRunRefreshesOnStart(t *testing.T) {
database := readTestDatabase(t)
archive := buildTarGzForTest(t, map[string][]byte{"GeoLite2-City_20260101/" + databaseFileName: database})
httpClient, transport := newDownloadClientForTest(archive)
r := newRefresherForTest(t, t.TempDir(), httpClient)
ctx := runRefresherForTest(t, r)
// A missing database is due right away, so the refresher downloads one as soon as it starts
require.Eventually(t, func() bool {
country, _, err := r.service.GetLocationByIP(ctx, "81.2.69.142")
return err == nil && country == "United Kingdom"
}, 30*time.Second, 100*time.Millisecond, "the database was never downloaded")
// It doesn't download again once the database on disk is fresh
require.Never(t, func() bool {
return transport.requests.Load() > 1
}, 3*time.Second, 250*time.Millisecond, "the database was downloaded again while it was still fresh")
}
func TestRefresherDisabledDoesNotDownload(t *testing.T) {
// Without a way to reach the download URL there's nothing to refresh, but the file is still watched
httpClient, transport := newDownloadClientForTest(nil)
r := newRefresherForTest(t, t.TempDir(), httpClient)
r.disabled = true
runRefresherForTest(t, r)
require.Never(t, func() bool {
return transport.requests.Load() > 0
}, 3*time.Second, 250*time.Millisecond, "the database was downloaded even though refreshes are disabled")
}

View File

@@ -0,0 +1,176 @@
package geolite
import (
"context"
"fmt"
"log/slog"
"net"
"net/netip"
"os"
"sync"
"github.com/oschwald/maxminddb-golang/v2"
"github.com/pocket-id/pocket-id/backend/internal/utils"
)
// The GeoLite2 City database is kept on disk and memory-mapped (the format is optimized for random access)
//
// The database file is considered cache, not state: it is a copy of a public artifact that any replica can rebuild on its own, so nothing is lost when a node goes away, and every replica keeps its own without needing to replicate anything
// It is also the supported way to supply a database by hand, which is what air-gapped deployments do: the file is watched, so replacing it takes effect without a restart
// internalNetworkCountry is reported for addresses that aren't routable on the public Internet
const internalNetworkCountry = "Internal Network"
// Service resolves IP addresses to locations, against a memory-mapped GeoLite2 City database
type Service struct {
log *slog.Logger
dbPath string
// mu guards the fields below
// A lookup holds it for reading throughout, so a reload can't unmap the database from under it
mu sync.RWMutex
// db is the database currently mapped, or nil when there is no readable database at dbPath
db *maxminddb.Reader
// dbModTime and dbSize identify the file that was mapped, so a reload of an unchanged file is skipped
dbModTime int64
dbSize int64
}
func newService(log *slog.Logger, dbPath string) *Service {
return &Service{
log: log,
dbPath: dbPath,
}
}
// GetLocationByIP returns the country and city of the given IP address
// Both are empty when the address isn't in the database, or when no database is available
func (s *Service) GetLocationByIP(_ context.Context, ipAddress string) (country string, city string, err error) {
if ipAddress == "" {
return "", "", nil
}
// Check the IP address against known private IP ranges, which can be short-circuited
ip := net.ParseIP(ipAddress)
if ip != nil {
switch {
case utils.IsLocalIPv6(ip):
return internalNetworkCountry, "LAN", nil
case utils.IsTailscaleIP(ip):
return internalNetworkCountry, "Tailscale", nil
case utils.IsPrivateIP(ip):
return internalNetworkCountry, "LAN", nil
case utils.IsLocalhostIP(ip):
return internalNetworkCountry, "localhost", nil
}
}
addr, err := netip.ParseAddr(ipAddress)
if err != nil {
return "", "", fmt.Errorf("failed to parse IP address: %w", err)
}
// The read lock is held for the whole lookup, including decoding, because the record is decoded straight out of the mapped file
s.mu.RLock()
defer s.mu.RUnlock()
if s.db == nil {
// No database is available
return "", "", nil
}
result := s.db.Lookup(addr)
if !result.Found() {
return "", "", nil
}
var record geoLiteRecord
err = result.Decode(&record)
if err != nil {
return "", "", fmt.Errorf("failed to decode database record: %w", err)
}
return record.Country.Names["en"], record.City.Names["en"], nil
}
// geoLiteRecord is the subset of a GeoLite2 City record that Pocket ID uses
type geoLiteRecord struct {
City struct {
Names map[string]string `maxminddb:"names"`
} `maxminddb:"city"`
Country struct {
Names map[string]string `maxminddb:"names"`
} `maxminddb:"country"`
}
// load maps the database at dbPath, replacing the one currently mapped
// A file that is already mapped is left alone, so reloading after an unrelated change is free
func (s *Service) load(ctx context.Context) error {
info, err := os.Stat(s.dbPath)
if err != nil {
if os.IsNotExist(err) {
// There's no database to map yet: lookups return no location until one shows up
s.unload()
return nil
}
return fmt.Errorf("failed to stat the GeoLite2 City database: %w", err)
}
modTime, size := info.ModTime().UnixNano(), info.Size()
s.mu.RLock()
unchanged := s.db != nil && s.dbModTime == modTime && s.dbSize == size
s.mu.RUnlock()
if unchanged {
return nil
}
db, err := maxminddb.Open(s.dbPath)
if err != nil {
return fmt.Errorf("failed to open the GeoLite2 City database: %w", err)
}
s.mu.Lock()
old := s.db
s.db = db
s.dbModTime = modTime
s.dbSize = size
s.mu.Unlock()
// The previous database is unmapped only once no lookup can still be reading it, which the write lock above guarantees
closeDatabase(old)
s.log.InfoContext(ctx, "Loaded the GeoLite2 City database",
slog.String("path", s.dbPath),
slog.Time("modTime", info.ModTime()),
slog.Int64("size", size),
)
return nil
}
// unload drops the database currently mapped, so lookups stop returning locations from a file that is no longer there
func (s *Service) unload() {
s.mu.Lock()
old := s.db
s.db = nil
s.dbModTime = 0
s.dbSize = 0
s.mu.Unlock()
closeDatabase(old)
}
// closeDatabase unmaps a database
// The caller must have already made it unreachable to lookups, since unmapping a database that is still being read would crash the process
func closeDatabase(db *maxminddb.Reader) {
if db == nil {
return
}
// Unmapping only fails if the mapping is already gone, which is not something the caller can act on
_ = db.Close()
}

View File

@@ -0,0 +1,183 @@
package geolite
import (
"os"
"sync"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestServiceGetLocationByIPPrivateRanges(t *testing.T) {
// Private addresses are short-circuited, so the service resolves them even with no database around
svc, _ := newServiceForTest(t, nil)
tests := []struct {
name string
ipAddress string
country string
city string
}{
{name: "empty address", ipAddress: ""},
{name: "private LAN IPv4", ipAddress: "192.168.1.20", country: internalNetworkCountry, city: "LAN"},
{name: "private LAN IPv4 in the 10/8 range", ipAddress: "10.4.5.6", country: internalNetworkCountry, city: "LAN"},
{name: "Tailscale IPv4", ipAddress: "100.101.102.103", country: internalNetworkCountry, city: "Tailscale"},
{name: "IPv6 unique local address", ipAddress: "fd00::1", country: internalNetworkCountry, city: "LAN"},
{name: "IPv4 loopback", ipAddress: "127.0.0.1", country: internalNetworkCountry, city: "LAN"},
{name: "IPv6 loopback", ipAddress: "::1", country: internalNetworkCountry, city: "LAN"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
country, city, err := svc.GetLocationByIP(t.Context(), tt.ipAddress)
require.NoError(t, err)
require.Equal(t, tt.country, country)
require.Equal(t, tt.city, city)
})
}
}
func TestServiceGetLocationByIPInvalidAddress(t *testing.T) {
svc, _ := newServiceForTest(t, nil)
_, _, err := svc.GetLocationByIP(t.Context(), "not-an-ip")
require.Error(t, err)
require.ErrorContains(t, err, "failed to parse IP address")
}
func TestServiceGetLocationByIP(t *testing.T) {
svc, _ := newServiceForTest(t, readTestDatabase(t))
tests := []struct {
name string
ipAddress string
country string
city string
}{
{name: "public IPv4 with country and city", ipAddress: "81.2.69.142", country: "United Kingdom", city: "London"},
{name: "public IPv4 in another country", ipAddress: "216.160.83.56", country: "United States", city: "Milton"},
{name: "public IPv4 with country only", ipAddress: "67.43.156.1", country: "Bhutan"},
{name: "public IPv6", ipAddress: "2001:218::1", country: "Japan"},
{name: "public address not in the database", ipAddress: "8.8.8.8"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
country, city, err := svc.GetLocationByIP(t.Context(), tt.ipAddress)
require.NoError(t, err)
require.Equal(t, tt.country, country)
require.Equal(t, tt.city, city)
})
}
}
func TestServiceGetLocationByIPWithoutDatabase(t *testing.T) {
// Air-gapped deployments that haven't supplied a database yet get no location, rather than an error on every audit log entry
svc, _ := newServiceForTest(t, nil)
country, city, err := svc.GetLocationByIP(t.Context(), "81.2.69.142")
require.NoError(t, err)
require.Empty(t, country)
require.Empty(t, city)
}
func TestServiceLoadMissingDatabase(t *testing.T) {
svc, dbPath := newServiceForTest(t, readTestDatabase(t))
country, _, err := svc.GetLocationByIP(t.Context(), "81.2.69.142")
require.NoError(t, err)
require.Equal(t, "United Kingdom", country)
// A database that goes away stops being used, rather than being served from a mapping of a file that no longer exists
require.NoError(t, os.Remove(dbPath))
require.NoError(t, svc.load(t.Context()))
country, _, err = svc.GetLocationByIP(t.Context(), "81.2.69.142")
require.NoError(t, err)
require.Empty(t, country)
}
func TestServiceLoadInvalidDatabase(t *testing.T) {
svc, dbPath := newServiceForTest(t, readTestDatabase(t))
// A corrupted file fails to load, and the database already mapped keeps serving lookups
writeDatabaseFileForTest(t, dbPath, []byte("not a database"))
err := svc.load(t.Context())
require.Error(t, err)
require.ErrorContains(t, err, "failed to open the GeoLite2 City database")
country, _, err := svc.GetLocationByIP(t.Context(), "81.2.69.142")
require.NoError(t, err)
require.Equal(t, "United Kingdom", country)
}
func TestServiceLoadUnchangedDatabase(t *testing.T) {
// Reloading a file that hasn't changed keeps the current mapping, so an unrelated event in the watched directory costs nothing
svc, _ := newServiceForTest(t, readTestDatabase(t))
svc.mu.RLock()
before := svc.db
svc.mu.RUnlock()
require.NoError(t, svc.load(t.Context()))
svc.mu.RLock()
after := svc.db
svc.mu.RUnlock()
require.Same(t, before, after)
}
func TestServiceConcurrentLookupsDuringReload(t *testing.T) {
// Lookups read straight out of the mapped file, so a reload must not unmap a database that a lookup is still reading
database := readTestDatabase(t)
svc, dbPath := newServiceForTest(t, database)
const (
lookers = 16
reloads = 25
)
stop := make(chan struct{})
errs := make([]error, lookers)
var wg sync.WaitGroup
wg.Add(lookers)
for i := range lookers {
go func() {
defer wg.Done()
for {
select {
case <-stop:
return
default:
}
_, _, err := svc.GetLocationByIP(t.Context(), "81.2.69.142")
if err != nil {
errs[i] = err
return
}
}
}()
}
for i := range reloads {
writeDatabaseFileForTest(t, dbPath, database)
// Force a distinct modification time, so every load really does remap the file instead of finding it unchanged
modTime := time.Now().Add(-time.Duration(i) * time.Second)
require.NoError(t, os.Chtimes(dbPath, modTime, modTime))
require.NoError(t, svc.load(t.Context()))
}
close(stop)
wg.Wait()
for _, err := range errs {
require.NoError(t, err)
}
}

Binary file not shown.

After

Width:  |  Height:  |  Size: 21 KiB

View File

@@ -0,0 +1,6 @@
# GeoLite test data
`GeoLite2-City-Test.mmdb` is the sample database published by MaxMind in the
[maxmind/MaxMind-DB](https://github.com/maxmind/MaxMind-DB) repository, under `test-data/`.
It contains a handful of synthetic networks only, so it is a few kilobytes rather than the tens of megabytes of the real database.

View File

@@ -0,0 +1,138 @@
package geolite
import (
"archive/tar"
"bytes"
"compress/gzip"
"io"
"log/slog"
"net/http"
"os"
"path/filepath"
"sync/atomic"
"testing"
"github.com/stretchr/testify/require"
testutils "github.com/pocket-id/pocket-id/backend/internal/utils/testing"
)
// testDownloadURL is the URL the mock HTTP client serves the database from
const testDownloadURL = "https://example.com/geolite/GeoLite2-City.tar.gz"
// testDatabasePath is the sample database published by MaxMind, see testdata/README.md
const testDatabasePath = "testdata/GeoLite2-City-Test.mmdb"
// readTestDatabase returns the raw sample GeoLite2 City database
func readTestDatabase(t *testing.T) []byte {
t.Helper()
data, err := os.ReadFile(testDatabasePath)
require.NoError(t, err)
return data
}
// testLogger returns a logger that discards everything, so tests don't spam the output
func testLogger() *slog.Logger {
return slog.New(slog.DiscardHandler)
}
// newServiceForTest returns a Service backed by a database file inside a temporary directory, along with the path of that file
// When data is nil no database is written, so the service starts with nothing to look up against
func newServiceForTest(t *testing.T, data []byte) (*Service, string) {
t.Helper()
dbPath := filepath.Join(t.TempDir(), "GeoLite2-City.mmdb")
if data != nil {
writeDatabaseFileForTest(t, dbPath, data)
}
return newServiceAtPathForTest(t, dbPath), dbPath
}
// newServiceAtPathForTest returns a Service backed by the database file at dbPath, loaded and unmapped when the test ends
// Unmapping is required on Windows
func newServiceAtPathForTest(t *testing.T, dbPath string) *Service {
t.Helper()
svc := newService(testLogger(), dbPath)
t.Cleanup(svc.unload)
err := svc.load(t.Context())
require.NoError(t, err)
return svc
}
// writeDatabaseFileForTest puts a database at path the same way the refresher does: written elsewhere, then moved into place
func writeDatabaseFileForTest(t *testing.T, path string, data []byte) {
t.Helper()
tmpPath := path + ".tmp"
err := os.WriteFile(tmpPath, data, 0600)
require.NoError(t, err)
err = os.Rename(tmpPath, path)
require.NoError(t, err)
}
// countingRoundTripper serves a fixed response for testDownloadURL and counts how many requests it has received
type countingRoundTripper struct {
body []byte
statusCode int
requests atomic.Int32
}
func (rt *countingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
if req.URL.String() != testDownloadURL {
return testutils.NewMockResponse(http.StatusNotFound, ""), nil
}
rt.requests.Add(1)
statusCode := rt.statusCode
if statusCode == 0 {
statusCode = http.StatusOK
}
return &http.Response{
StatusCode: statusCode,
Body: io.NopCloser(bytes.NewReader(rt.body)),
Header: make(http.Header),
ContentLength: int64(len(rt.body)),
}, nil
}
// newDownloadClientForTest returns an HTTP client that serves body at testDownloadURL, along with the transport that counts the requests it receives
func newDownloadClientForTest(body []byte) (*http.Client, *countingRoundTripper) {
rt := &countingRoundTripper{body: body}
return &http.Client{Transport: rt}, rt
}
// buildTarGzForTest returns a gzipped tarball holding the given files, mirroring the archive MaxMind publishes
func buildTarGzForTest(t *testing.T, files map[string][]byte) []byte {
t.Helper()
buf := &bytes.Buffer{}
gzw := gzip.NewWriter(buf)
tw := tar.NewWriter(gzw)
for name, content := range files {
err := tw.WriteHeader(&tar.Header{
Name: name,
Mode: 0o644,
Size: int64(len(content)),
Typeflag: tar.TypeReg,
})
require.NoError(t, err)
_, err = tw.Write(content)
require.NoError(t, err)
}
require.NoError(t, tw.Close())
require.NoError(t, gzw.Close())
return buf.Bytes()
}

View File

@@ -1,31 +0,0 @@
package job
import (
"context"
"time"
"github.com/go-co-op/gocron/v2"
"github.com/pocket-id/pocket-id/backend/internal/service"
)
type GeoLiteUpdateJobs struct {
geoLiteService *service.GeoLiteService
}
func (s *Scheduler) RegisterGeoLiteUpdateJobs(ctx context.Context, geoLiteService *service.GeoLiteService) error {
// Check if the service needs periodic updating
if geoLiteService.DisableUpdater() {
// Nothing to do
return nil
}
jobs := &GeoLiteUpdateJobs{geoLiteService: geoLiteService}
// Run every 24 hours (and right away)
return s.RegisterJob(ctx, "UpdateGeoLiteDB", gocron.DurationJob(24*time.Hour), jobs.updateGoeLiteDB, service.RegisterJobOpts{RunImmediately: true})
}
func (j *GeoLiteUpdateJobs) updateGoeLiteDB(ctx context.Context) error {
return j.geoLiteService.UpdateDatabase(ctx)
}

View File

@@ -17,28 +17,32 @@ type NewLoginEmailSender interface {
SendNewLogin(ctx context.Context, dbConfig *appconfig.AppConfigModel, userFullName, userEmail, ipAddress, country, city, device string, dateTime time.Time) error
}
type IPLocationResolver interface {
GetLocationByIP(ctx context.Context, ipAddress string) (country string, city string, err error)
}
type AuditLogService struct {
db *gorm.DB
emailSender NewLoginEmailSender
geoliteService *GeoLiteService
ipLocator IPLocationResolver
appConfigService *appconfig.AppConfigService
}
func NewAuditLogService(db *gorm.DB, emailSender NewLoginEmailSender, geoliteService *GeoLiteService, appConfigService *appconfig.AppConfigService) *AuditLogService {
func NewAuditLogService(db *gorm.DB, emailSender NewLoginEmailSender, ipLocator IPLocationResolver, appConfigService *appconfig.AppConfigService) *AuditLogService {
return &AuditLogService{
db: db,
emailSender: emailSender,
geoliteService: geoliteService,
ipLocator: ipLocator,
appConfigService: appConfigService,
}
}
// Create creates a new audit log entry in the database
func (s *AuditLogService) Create(ctx context.Context, event model.AuditLogEvent, ipAddress, userAgent, userID string, data model.AuditLogData, tx *gorm.DB) (model.AuditLog, bool) {
country, city, err := s.geoliteService.GetLocationByIP(ipAddress)
country, city, err := s.ipLocator.GetLocationByIP(ctx, ipAddress)
if err != nil {
// Log the error but don't interrupt the operation
slog.Warn("Failed to get IP location", slog.String("ip", ipAddress), slog.Any("error", err))
slog.WarnContext(ctx, "Failed to get IP location", slog.String("ip", ipAddress), slog.Any("error", err))
}
auditLog := model.AuditLog{

View File

@@ -1,292 +0,0 @@
package service
import (
"archive/tar"
"bytes"
"compress/gzip"
"context"
"errors"
"fmt"
"io"
"log/slog"
"net"
"net/http"
"net/netip"
"os"
"path/filepath"
"strings"
"sync"
"time"
"github.com/oschwald/maxminddb-golang/v2"
"github.com/pocket-id/pocket-id/backend/internal/utils"
"github.com/pocket-id/pocket-id/backend/internal/common"
)
const maxTotalSize = 300 * 1024 * 1024 // 300 MB limit for total decompressed size
type GeoLiteService struct {
httpClient *http.Client
disableUpdater bool
mutex sync.RWMutex
}
// NewGeoLiteService initializes a new GeoLiteService instance and starts a goroutine to update the GeoLite2 City database.
func NewGeoLiteService(httpClient *http.Client) *GeoLiteService {
service := &GeoLiteService{
httpClient: httpClient,
}
if common.EnvConfig.MaxMindLicenseKey == "" && common.EnvConfig.GeoLiteDBUrl == common.MaxMindGeoLiteCityUrl {
// Warn the user, and disable the periodic updater
slog.Warn("MAXMIND_LICENSE_KEY environment variable is empty: the GeoLite2 City database won't be updated")
service.disableUpdater = true
}
return service
}
func (s *GeoLiteService) DisableUpdater() bool {
return s.disableUpdater
}
// GetLocationByIP returns the country and city of the given IP address.
func (s *GeoLiteService) GetLocationByIP(ipAddress string) (country, city string, err error) {
if ipAddress == "" {
return "", "", nil
}
// Check the IP address against known private IP ranges
if ip := net.ParseIP(ipAddress); ip != nil {
if utils.IsLocalIPv6(ip) {
return "Internal Network", "LAN", nil
}
if utils.IsTailscaleIP(ip) {
return "Internal Network", "Tailscale", nil
}
if utils.IsPrivateIP(ip) {
return "Internal Network", "LAN", nil
}
if utils.IsLocalhostIP(ip) {
return "Internal Network", "localhost", nil
}
}
addr, err := netip.ParseAddr(ipAddress)
if err != nil {
return "", "", fmt.Errorf("failed to parse IP address: %w", err)
}
// Race condition between reading and writing the database.
s.mutex.RLock()
defer s.mutex.RUnlock()
db, err := maxminddb.Open(common.EnvConfig.GeoLiteDBPath)
if err != nil {
return "", "", err
}
defer db.Close()
var record struct {
City struct {
Names map[string]string `maxminddb:"names"`
} `maxminddb:"city"`
Country struct {
Names map[string]string `maxminddb:"names"`
} `maxminddb:"country"`
}
err = db.Lookup(addr).Decode(&record)
if err != nil {
return "", "", err
}
return record.Country.Names["en"], record.City.Names["en"], nil
}
// UpdateDatabase checks the age of the database and updates it if it's older than 14 days.
func (s *GeoLiteService) UpdateDatabase(parentCtx context.Context) error {
if s.isDatabaseUpToDate() {
slog.Info("GeoLite2 City database is up-to-date")
return nil
}
slog.Info("Updating GeoLite2 City database")
downloadUrl := common.EnvConfig.GeoLiteDBUrl
if strings.Contains(downloadUrl, "%s") {
downloadUrl = fmt.Sprintf(downloadUrl, common.EnvConfig.MaxMindLicenseKey)
}
ctx, cancel := context.WithTimeout(parentCtx, 10*time.Minute)
defer cancel()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadUrl, nil)
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
resp, err := s.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to download database: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("failed to download database, received HTTP %d", resp.StatusCode)
}
// Extract the database file directly to the target path
err = s.extractDatabase(resp.Body)
if err != nil {
return fmt.Errorf("failed to extract database: %w", err)
}
slog.Info("GeoLite2 City database successfully updated.")
return nil
}
// isDatabaseUpToDate checks if the database file is older than 14 days.
func (s *GeoLiteService) isDatabaseUpToDate() bool {
info, err := os.Stat(common.EnvConfig.GeoLiteDBPath)
if err != nil {
// If the file doesn't exist, treat it as not up-to-date
return false
}
return time.Since(info.ModTime()) < 14*24*time.Hour
}
// extractDatabase extracts the database file from the tar.gz archive directly to the target location.
func (s *GeoLiteService) extractDatabase(reader io.Reader) error {
// Check for gzip magic number
buf := make([]byte, 2)
_, err := io.ReadFull(reader, buf)
if err != nil {
return fmt.Errorf("failed to read magic number: %w", err)
}
// Check if the file starts with the gzip magic number
// Gosec returns false positive for "G602: slice index out of range"
//nolint:gosec
isGzip := buf[0] == 0x1f && buf[1] == 0x8b
if !isGzip {
// If not gzip, assume it's a regular database file
return s.writeDatabaseFile(io.MultiReader(bytes.NewReader(buf), reader))
}
gzr, err := gzip.NewReader(io.MultiReader(bytes.NewReader(buf), reader))
if err != nil {
return fmt.Errorf("failed to create gzip reader: %w", err)
}
defer gzr.Close()
tarReader := tar.NewReader(gzr)
var totalSize int64
// Iterate over the files in the tar archive
for {
header, err := tarReader.Next()
if errors.Is(err, io.EOF) {
break
} else if err != nil {
return fmt.Errorf("failed to read tar archive: %w", err)
}
// Check if the file is the GeoLite2-City.mmdb file
if header.Typeflag == tar.TypeReg && filepath.Base(header.Name) == "GeoLite2-City.mmdb" {
totalSize += header.Size
if totalSize > maxTotalSize {
return errors.New("total decompressed size exceeds maximum allowed limit")
}
// extract to a temporary file to avoid having a corrupted db in case of write failure.
baseDir := filepath.Dir(common.EnvConfig.GeoLiteDBPath)
tmpFile, err := os.CreateTemp(baseDir, "geolite.*.mmdb.tmp")
if err != nil {
return fmt.Errorf("failed to create temporary database file: %w", err)
}
tempName := tmpFile.Name()
// Write the file contents directly to the target location
if _, err := io.Copy(tmpFile, tarReader); err != nil { //nolint:gosec
// if fails to write, then cleanup and throw an error
tmpFile.Close()
os.Remove(tempName)
return fmt.Errorf("failed to write database file: %w", err)
}
tmpFile.Close()
// ensure the database is not corrupted
db, err := maxminddb.Open(tempName)
if err != nil {
// if fails to write, then cleanup and throw an error
os.Remove(tempName)
return fmt.Errorf("failed to open downloaded database file: %w", err)
}
db.Close()
// ensure we lock the structure before we overwrite the database
// to prevent race conditions between reading and writing the mmdb.
s.mutex.Lock()
// replace the old file with the new file
err = os.Rename(tempName, common.EnvConfig.GeoLiteDBPath)
s.mutex.Unlock()
if err != nil {
// if cannot overwrite via rename, then cleanup and throw an error
os.Remove(tempName)
return fmt.Errorf("failed to replace database file: %w", err)
}
return nil
}
}
return errors.New("GeoLite2-City.mmdb not found in archive")
}
func (s *GeoLiteService) writeDatabaseFile(reader io.Reader) error {
baseDir := filepath.Dir(common.EnvConfig.GeoLiteDBPath)
tmpFile, err := os.CreateTemp(baseDir, "geolite.*.mmdb.tmp")
if err != nil {
return fmt.Errorf("failed to create temporary database file: %w", err)
}
defer tmpFile.Close()
// Limit the amount we read to maxTotalSize.
// We read one extra byte to detect if the source is larger than the limit.
limitReader := io.LimitReader(reader, maxTotalSize+1)
// Write the file contents directly to the temporary file
written, err := io.Copy(tmpFile, limitReader)
if err != nil {
os.Remove(tmpFile.Name())
return fmt.Errorf("failed to write database file: %w", err)
}
if written > maxTotalSize {
os.Remove(tmpFile.Name())
return errors.New("total database size exceeds maximum allowed limit")
}
// Validate the downloaded database file
if db, err := maxminddb.Open(tmpFile.Name()); err == nil {
db.Close()
} else {
os.Remove(tmpFile.Name())
return fmt.Errorf("failed to open downloaded database file: %w", err)
}
// Ensure atomic replacement of the old database file
s.mutex.Lock()
err = os.Rename(tmpFile.Name(), common.EnvConfig.GeoLiteDBPath)
s.mutex.Unlock()
if err != nil {
os.Remove(tmpFile.Name())
return fmt.Errorf("failed to replace database file: %w", err)
}
return nil
}