mirror of
https://github.com/pocket-id/pocket-id.git
synced 2026-09-05 18:51:27 +02:00
refactor: move GeoLite to its own package (#1647)
This commit is contained in:
committed by
GitHub
parent
1c9233c236
commit
22e3909c6c
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
173
backend/internal/geolite/download.go
Normal file
173
backend/internal/geolite/download.go
Normal 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
|
||||
}
|
||||
171
backend/internal/geolite/download_test.go
Normal file
171
backend/internal/geolite/download_test.go
Normal 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)
|
||||
}
|
||||
74
backend/internal/geolite/module.go
Normal file
74
backend/internal/geolite/module.go
Normal 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)
|
||||
}
|
||||
218
backend/internal/geolite/refresh.go
Normal file
218
backend/internal/geolite/refresh.go
Normal 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))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
208
backend/internal/geolite/refresh_test.go
Normal file
208
backend/internal/geolite/refresh_test.go
Normal 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")
|
||||
}
|
||||
176
backend/internal/geolite/service.go
Normal file
176
backend/internal/geolite/service.go
Normal 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()
|
||||
}
|
||||
183
backend/internal/geolite/service_test.go
Normal file
183
backend/internal/geolite/service_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
BIN
backend/internal/geolite/testdata/GeoLite2-City-Test.mmdb
vendored
Normal file
BIN
backend/internal/geolite/testdata/GeoLite2-City-Test.mmdb
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 21 KiB |
6
backend/internal/geolite/testdata/README.md
vendored
Normal file
6
backend/internal/geolite/testdata/README.md
vendored
Normal 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.
|
||||
138
backend/internal/geolite/testutils_test.go
Normal file
138
backend/internal/geolite/testutils_test.go
Normal 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()
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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{
|
||||
|
||||
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user