Merge pull request 'Fix data races, robustness & cleanup (#10-#21)' (#22) from fix/review-10-21 into main

Reviewed-on: #22
Reviewed-by: codex <codex@mail.thuanle.me>

.env.example

@@ -1,2 +1,6 @@
-API_PORT=28080
+# API_PORT is the port the service listens on INSIDE the container.
+# With docker-compose's "28080:8080" mapping, set this to 8080 so the host
+# can reach it at http://localhost:28080. Override only if you also change
+# the host-side mapping / non-Docker run.
+API_PORT=8080
 GIN_TRUSTED_PROXY_IP=127.0.0.1

README.md

@@ -31,7 +31,7 @@ The service will be available at `http://localhost:28080`
 
 ### Manual Installation
 
-1. Install Go 1.22 or later
+1. Install Go 1.25 or later
 2. Clone and build:
 ```bash
 git clone <repository-url>

go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/gin-gonic/gin v1.12.0
 	github.com/go-resty/resty/v2 v2.17.2
 	github.com/joho/godotenv v1.5.1
-	github.com/oschwald/maxminddb-golang v1.13.1
+	github.com/oschwald/maxminddb-golang/v2 v2.4.1
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/rs/zerolog v1.35.1
 	github.com/safing/mmdbmeld v0.3.0
@@ -32,7 +32,6 @@ require (
 	github.com/maxmind/mmdbwriter v1.2.0 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
-	github.com/oschwald/maxminddb-golang/v2 v2.4.1 // indirect
 	github.com/pelletier/go-toml/v2 v2.4.3 // indirect
 	github.com/quic-go/qpack v0.6.0 // indirect
 	github.com/quic-go/quic-go v0.60.0 // indirect

go.sum

@@ -55,8 +55,6 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w
 github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
 github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
-github.com/oschwald/maxminddb-golang v1.13.1 h1:G3wwjdN9JmIK2o/ermkHM+98oX5fS+k5MbwsmL4MRQE=
-github.com/oschwald/maxminddb-golang v1.13.1/go.mod h1:K4pgV9N/GcK694KSTmVSDTODk4IsCNThNdTmnaBZ/F8=
 github.com/oschwald/maxminddb-golang/v2 v2.4.1 h1:OffzqSABE3Sw354GdBThqDsKfpA4GWBqOY2P91V8tjI=
 github.com/oschwald/maxminddb-golang/v2 v2.4.1/go.mod h1:CZK8iQQMKfy6mKOifoyUmrj4vTHnMiGVaS7hDaZZxQ0=
 github.com/pelletier/go-toml/v2 v2.4.3 h1:GTRvJQutkOSftxIFD5xw9aepkYNuPWmVJpffdDPYVpY=

internal/data/data_test.go

@@ -0,0 +1,133 @@
+package data
+
+import (
+	"errors"
+	"net"
+	"os"
+	"path/filepath"
+	"sync"
+	"sync/atomic"
+	"testing"
+)
+
+// nested looks up a nested map value by successive string keys, returning "" if absent.
+func nested(m map[string]any, keys ...string) string {
+	var cur any = m
+	for _, k := range keys {
+		mm, ok := cur.(map[string]any)
+		if !ok {
+			return ""
+		}
+		cur = mm[k]
+	}
+	s, _ := cur.(string)
+	return s
+}
+
+// parseIP is a test helper that panics on invalid input (used only with literals).
+func parseIP(s string) net.IP {
+	ip := net.ParseIP(s)
+	if ip == nil {
+		panic("invalid ip in test: " + s)
+	}
+	return ip
+}
+
+// The package's data helpers use relative paths ("data/..."), so tests must
+// run with the working directory at the repository root. TestMain moves there
+// once for the whole package.
+func TestMain(m *testing.M) {
+	if root, err := filepath.Abs("../.."); err == nil {
+		_ = os.Chdir(root)
+	}
+	os.Exit(m.Run())
+}
+
+// freshDB returns a loaded IpDb, skipping the test when no mmdb is present.
+func freshDB(t *testing.T) *IpDb {
+	t.Helper()
+	db := Ins()
+	if !db.IsLoaded() {
+		t.Skip("no mmdb present, skipping test that needs real data")
+	}
+	return db
+}
+
+func TestQueryValidIP(t *testing.T) {
+	db := freshDB(t)
+
+	data, netRec, err := db.Query("8.8.8.8")
+	if err != nil {
+		t.Fatalf("Query(8.8.8.8) unexpected error: %v", err)
+	}
+	if data == nil {
+		t.Fatal("Query(8.8.8.8) returned nil data")
+	}
+	// 8.8.8.8 is Google (AS15169), US.
+	if got := nested(data, "country", "iso_code"); got != "US" {
+		t.Errorf("country.iso_code = %v, want US", got)
+	}
+	if netRec == nil {
+		t.Error("expected a non-nil network record")
+	} else if !netRec.Contains(parseIP("8.8.8.8")) {
+		t.Errorf("network %s should contain 8.8.8.8", netRec)
+	}
+}
+
+func TestQueryInvalidIP(t *testing.T) {
+	db := freshDB(t)
+
+	if _, _, err := db.Query("not-an-ip"); err == nil {
+		t.Fatal("expected error for invalid IP, got nil")
+	}
+}
+
+func TestQueryNotFound(t *testing.T) {
+	db := freshDB(t)
+
+	// 240.0.0.0/4 is reserved/unallocated and should not be in a public GeoIP DB.
+	_, _, err := db.Query("240.0.0.1")
+	if err == nil {
+		t.Fatal("expected not-found error, got nil")
+	}
+	if !errors.Is(err, ErrNotFound) {
+		t.Errorf("expected ErrNotFound sentinel, got %v", err)
+	}
+}
+
+// TestQueryReloadConcurrent is a regression test for the data race between
+// Query (reads d.r) and Reload (reassigns d.r / closes the old reader). It
+// must pass under `go test -race`.
+func TestQueryReloadConcurrent(t *testing.T) {
+	db := freshDB(t)
+
+	var reloadOK, queries int64
+	var wg sync.WaitGroup
+
+	for w := 0; w < 8; w++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for i := 0; i < 5000; i++ {
+				_, _, _ = db.Query("8.8.8.8")
+				atomic.AddInt64(&queries, 1)
+			}
+		}()
+	}
+	for r := 0; r < 4; r++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for i := 0; i < 30; i++ {
+				if db.Reload() == nil {
+					atomic.AddInt64(&reloadOK, 1)
+				}
+			}
+		}()
+	}
+	wg.Wait()
+
+	if reloadOK == 0 {
+		t.Fatalf("expected reloads to succeed, got reloads=%d queries=%d", reloadOK, queries)
+	}
+}

internal/data/ipdb.go

@@ -9,7 +9,7 @@ import (
 	"sync"
 	"thuanle.me/ip-info/configs"
 )
-import reader "github.com/oschwald/maxminddb-golang"
+import reader "github.com/oschwald/maxminddb-golang/v2"
 
 type IpDb struct {
 	r  *reader.Reader
@@ -61,7 +61,7 @@ func CleanupDataDir() error {
 			if err != nil {
 				return fmt.Errorf("failed to delete file %s: %w", filePath, err)
 			}
-			fmt.Printf("Deleted file: %s\n", filePath)
+			log.Info().Str("file", filePath).Msg("Deleted temp file")
 		}
 	}
 
@@ -83,7 +83,11 @@ func (d *IpDb) Reload() error {
 	d.dbFile = tmpFile
 	r, err := reader.Open(tmpFile)
 	if err != nil {
-		log.Err(err).Msg("Failed to open mmdb")
+		log.Err(err).Str("file", tmpFile).Msg("Failed to open mmdb")
+		// Open failed: drop the orphaned temp clone and leave d.dbFile/d.r
+		// pointing at the previously loaded DB.
+		_ = os.Remove(tmpFile)
+		d.dbFile = wilDeleteFile
 		return err
 	}
 	log.Info().Str("file", tmpFile).Msg("MMDB reloaded")
@@ -104,6 +108,8 @@ func (d *IpDb) Reload() error {
 }
 
 func (d *IpDb) IsLoaded() bool {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
 	return d.r != nil
 }
 
@@ -125,6 +131,8 @@ func cloneDBFile() (string, error) {
 	// Copy the data from srcFile to dstFile
 	_, err = io.CopyBuffer(tmpFile, srcFile, buf)
 	if err != nil {
+		// Leave nothing behind on a failed copy.
+		_ = os.Remove(tmpFile.Name())
 		return "", err
 	}
 

internal/data/query.go

@@ -1,23 +1,57 @@
 package data
 
 import (
+	"errors"
 	"fmt"
 	"net"
+	"net/netip"
 )
 
+// ErrNotFound is returned by Query when the IP parses but is not present in the DB.
+var ErrNotFound = errors.New("IP not found")
+
 func (d *IpDb) Query(ipArg string) (map[string]any, *net.IPNet, error) {
-	ip := net.ParseIP(ipArg)
-	if ip == nil {
+	// maxminddb-golang v2 requires a netip.Addr.
+	addr, err := netip.ParseAddr(ipArg)
+	if err != nil {
 		return nil, nil, fmt.Errorf("invalid IP: %s", ipArg)
 	}
-	// Get data of IP.
-	anyData := make(map[string]any)
-	recordNet, ok, err := d.r.LookupNetwork(ip, &anyData)
-	if err != nil {
+
+	// Hold the read lock: Reload() reassigns d.r and closes the old reader,
+	// so reading d.r / calling Lookup without synchronization is a data race.
+	d.mu.RLock()
+	defer d.mu.RUnlock()
+
+	if d.r == nil {
+		return nil, nil, fmt.Errorf("db not loaded")
+	}
+
+	// Look up the record for the IP.
+	result := d.r.Lookup(addr)
+	if err := result.Err(); err != nil {
 		return nil, nil, err
 	}
-	if !ok {
-		return nil, nil, fmt.Errorf("IP not found: %s", ipArg)
+	if !result.Found() {
+		return nil, nil, fmt.Errorf("%w: %s", ErrNotFound, ipArg)
+	}
+
+	// Decode the record into a generic map.
+	anyData := make(map[string]any)
+	if err := result.Decode(&anyData); err != nil {
+		return nil, nil, err
+	}
+
+	// Convert the matched prefix back to *net.IPNet for callers.
+	var recordNet *net.IPNet
+	if prefix := result.Prefix(); prefix.IsValid() {
+		recordNet = prefixToIPNet(prefix)
 	}
 	return anyData, recordNet, nil
 }
+
+// prefixToIPNet converts a netip.Prefix to a *net.IPNet.
+func prefixToIPNet(prefix netip.Prefix) *net.IPNet {
+	addr := prefix.Addr().AsSlice()
+	mask := net.CIDRMask(prefix.Bits(), len(addr)*8)
+	return &net.IPNet{IP: addr, Mask: mask}
+}

internal/services/api/handler_json.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"errors"
 	"github.com/gin-gonic/gin"
 	"github.com/rs/zerolog/log"
 	"net"
@@ -10,34 +11,63 @@ import (
 
 func HandleJson(c *gin.Context) {
 	ip := c.ClientIP()
-	HandleIpInfo(c, ip)
+	// ClientIP can yield an unparseable value from malformed proxy headers;
+	// reject it explicitly instead of letting Query fail into a 500.
+	if net.ParseIP(ip) == nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid ip", "ip": ip})
+		return
+	}
+	// selfLookup=true preserves the "default value when lookup internal ip"
+	// contract (commit 9d07639): a valid but DB-absent client IP (loopback /
+	// private addresses) returns 200 {"ip": ...} rather than a 404, since the
+	// caller did not ask for a specific lookup.
+	HandleIpInfo(c, ip, true)
 }
 
 func HandleOtherIp(c *gin.Context) {
 	ip := c.Param("ip")
 	if net.ParseIP(ip) == nil {
-		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid ip"})
+		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid ip", "ip": ip})
 		return
 	}
-	HandleIpInfo(c, ip)
+	// The caller explicitly requested this IP, so a not-found is a real 404.
+	HandleIpInfo(c, ip, false)
 }
 
-func HandleIpInfo(c *gin.Context, ip string) {
+func HandleIpInfo(c *gin.Context, ip string, selfLookup bool) {
 	if !data.Ins().IsLoaded() {
 		log.Error().Msg("DB is not loaded")
-		c.String(http.StatusInternalServerError, "Try again later")
+		c.JSON(http.StatusServiceUnavailable, gin.H{
+			"error": "service unavailable, db still loading",
+			"ip":    ip,
+		})
 		return
 	}
 
 	ipData, _, err := data.Ins().Query(ip)
 	if err != nil {
+		// Distinguish "not found" (valid IP, absent from DB) from other errors.
+		if isNotFound(err) {
+			if selfLookup {
+				// Graceful default for the self-IP (/json) path: the client's
+				// address is valid but simply absent from the public GeoIP DB.
+				c.JSON(http.StatusOK, gin.H{"ip": ip})
+				return
+			}
+			c.JSON(http.StatusNotFound, gin.H{"error": "not found", "ip": ip})
+			return
+		}
 		log.Err(err).Msg("Failed to query IP")
-		c.JSON(http.StatusOK, gin.H{
-			"ip": ip,
-		})
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "internal error", "ip": ip})
 		return
 	}
 
 	ipData["ip"] = ip
 	c.JSON(http.StatusOK, ipData)
 }
+
+// isNotFound reports whether the query error means the IP was not present in
+// the DB (as opposed to an invalid IP or an internal lookup failure).
+func isNotFound(err error) bool {
+	return errors.Is(err, data.ErrNotFound)
+}

internal/services/api/handler_metrics.go

@@ -8,6 +8,6 @@ import (
 
 func HandleMetrics(c *gin.Context) {
 	c.JSON(http.StatusOK, gin.H{
-		"updated_at": db_updater.DbUpdatedAt.UnixMilli(),
+		"updated_at": db_updater.DbUpdatedAtMillis(),
 	})
 }

internal/services/api/main.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"errors"
 	"github.com/gin-gonic/gin"
 	"github.com/rs/zerolog/log"
@@ -8,6 +9,7 @@ import (
 	"os"
 	"thuanle.me/ip-info/configs"
 	"thuanle.me/ip-info/configs/key"
+	"time"
 )
 
 var srv *http.Server
@@ -48,6 +50,12 @@ func StartApiService() {
 }
 
 func Shutdown() {
-	_ = srv.Close()
+	// Graceful shutdown: let in-flight requests finish (bounded by a timeout)
+	// instead of abruptly closing connections.
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	if err := srv.Shutdown(ctx); err != nil {
+		log.Err(err).Msg("Error shutting down API service")
+	}
 	log.Info().Msg("API service stopped")
 }

internal/services/db_updater/downloader.go

@@ -7,6 +7,7 @@ import (
 	"os"
 	"path"
 	"thuanle.me/ip-info/configs"
+	"time"
 )
 
 func download(url string) bool {
@@ -16,7 +17,10 @@ func download(url string) bool {
 
 	etag, _ := readFile(filenameEtag)
 
-	client := resty.New()
+	client := resty.New().
+		SetTimeout(2 * time.Minute).
+		SetRetryCount(2).
+		SetRetryWaitTime(5 * time.Second)
 	resp, err := client.R().
 		SetHeader("If-None-Match", etag).
 		Get(url)

internal/services/db_updater/main.go

@@ -3,12 +3,24 @@ package db_updater
 import (
 	"github.com/robfig/cron/v3"
 	"github.com/rs/zerolog/log"
+	"os"
+	"path"
+	"sync/atomic"
 	"thuanle.me/ip-info/configs"
 	"thuanle.me/ip-info/internal/data"
-	"time"
+	"thuanle.me/ip-info/pkg/osx"
 )
 
-var DbUpdatedAt time.Time
+// DbUpdatedAt stores the DB last-updated time as unix milliseconds.
+// Accessed concurrently by the cron goroutine (writer) and the /metrics
+// handler (reader), so it must be synchronized — hence atomic.Int64.
+var DbUpdatedAt atomic.Int64
+
+// DbUpdatedAtMillis returns the last DB update time in unix milliseconds,
+// or 0 if the DB has never been updated.
+func DbUpdatedAtMillis() int64 {
+	return DbUpdatedAt.Load()
+}
 
 func StartUpdateDbService() {
 	c := cron.New()
@@ -25,14 +37,108 @@ func fetchDbs() {
 
 	if newFlag {
 		log.Info().Msg("New DB downloaded. Recreating mmdb")
-		_ = mergeMmdb()
+		// mmdbmeld.WriteMMDB truncates the canonical mmdb before writing. If
+		// the merge fails the file would be left empty/corrupt — and because
+		// download() already persisted the etag, the next daily run would get
+		// 304 Not Modified and never retry, leaving the service broken until a
+		// restart (which then fails to open the empty file). Back the file up
+		// first so we can restore it on failure.
+		backup, backupErr := backupMmdb()
+		if backupErr != nil {
+			// We could not back up the current canonical mmdb. Do NOT proceed:
+			// mergeMmdb() truncates the canonical file before writing, and with
+			// no backup a later merge/reload failure would leave it corrupt
+			// (unrecoverable on restart). Keep serving the current DB and let
+			// the next run retry.
+			log.Err(backupErr).Msg("Failed to back up mmdb, skipping rebuild")
+			return
+		}
+		if err := mergeMmdb(); err != nil {
+			log.Err(err).Msg("Failed to merge mmdb")
+			rollBackMmdb(backup)
+			return
+		}
 
 		err := data.Ins().Reload()
 		if err != nil {
 			log.Err(err).Msg("Failed to reload mmdb")
+			// Reload failed on the newly merged file: restore the previous
+			// canonical mmdb so we keep serving known-good data.
+			rollBackMmdb(backup)
 			return
 		}
 
-		DbUpdatedAt = time.Now()
+		// Both merge and reload succeeded: the new canonical file is live, the
+		// backup is no longer needed.
+		cleanupMmdbBackup(backup)
+
+		// Stamp the update time from the freshly written canonical file.
+		stampDbUpdatedAt()
+		return
+	}
+
+	// No new download: on startup this is the path that loads an existing DB.
+	// Ensure the timestamp reflects the file we are actually serving.
+	if DbUpdatedAt.Load() == 0 {
+		stampDbUpdatedAt()
+	}
+}
+
+// stampDbUpdatedAt records the mtime of the canonical mmdb (unix ms) as the
+// last DB update time. This covers both initial load (pre-existing files) and
+// post-download reloads, so /metrics never reports 0 while serving real data.
+func stampDbUpdatedAt() {
+	info, err := os.Stat(configs.MmdbDbFile)
+	if err != nil {
+		log.Err(err).Str("file", configs.MmdbDbFile).Msg("Failed to stat mmdb")
+		return
+	}
+	DbUpdatedAt.Store(info.ModTime().UnixMilli())
+}
+
+// backupMmdb copies the canonical mmdb to a sibling backup file. It returns
+// ("", nil) when there is no canonical file yet (first-ever build), in which
+// case there is nothing to restore on failure.
+func backupMmdb() (string, error) {
+	if _, err := os.Stat(configs.MmdbDbFile); os.IsNotExist(err) {
+		return "", nil
+	}
+	backup := configs.MmdbDbFile + ".bak"
+	if err := osx.Copy(configs.MmdbDbFile, backup); err != nil {
+		return "", err
+	}
+	return backup, nil
+}
+
+// rollBackMmdb restores the canonical mmdb from backup (if any) and deletes
+// the etag files so the next cron run re-downloads and retries the merge.
+// backup is "" on a first-ever build (no canonical file existed to back up),
+// in which case there is nothing to restore.
+func rollBackMmdb(backup string) {
+	if backup != "" {
+		if err := osx.Copy(backup, configs.MmdbDbFile); err != nil {
+			log.Err(err).Msg("Failed to restore mmdb backup")
+		} else {
+			log.Info().Str("file", configs.MmdbDbFile).Msg("Restored mmdb from backup")
+		}
+		_ = os.Remove(backup)
+	}
+	// Drop etags so a failed merge is retried on the next run instead of being
+	// masked by a 304 Not Modified.
+	for _, url := range configs.GeoDbSourcePaths {
+		etag := configs.GeoDbFolder + path.Base(url) + ".etag"
+		if err := os.Remove(etag); err != nil && !os.IsNotExist(err) {
+			log.Err(err).Str("file", etag).Msg("Failed to remove etag")
+		}
+	}
+}
+
+// cleanupMmdbBackup removes a successful merge's backup file.
+func cleanupMmdbBackup(backup string) {
+	if backup == "" {
+		return
+	}
+	if err := os.Remove(backup); err != nil && !os.IsNotExist(err) {
+		log.Err(err).Str("file", backup).Msg("Failed to remove mmdb backup")
 	}
 }

pkg/osx/file.go

@@ -6,7 +6,7 @@ import (
 )
 
 // Copy copies a file from source to destination
-func Copy(from, to string) error {
+func Copy(from, to string) (err error) {
 	// Open the source file for reading
 	srcFile, err := os.Open(from)
 	if err != nil {
@@ -19,7 +19,16 @@ func Copy(from, to string) error {
 	if err != nil {
 		return err
 	}
-	defer dstFile.Close()
+	// Closing may flush buffered writes, so its error must be surfaced — a
+	// successful copy with a failed close means the destination is incomplete.
+	// Defer the close and only let it overwrite the return value when the copy
+	// itself succeeded; otherwise keep the first error.
+	defer func() {
+		cerr := dstFile.Close()
+		if err == nil {
+			err = cerr
+		}
+	}()
 
 	// Use a buffer to copy the file in chunks
 	buf := make([]byte, 1024*1024) // 1 MB buffer