From e8cea8d07de5caf16f6d0e2e8ec3724b03de1d6a Mon Sep 17 00:00:00 2001
From: pg9182 <96569817+pg9182@users.noreply.github.com>
Date: Mon, 12 Dec 2022 16:17:41 -0500
Subject: pkg/atlas: Support region overrides for IP prefixes

---
 pkg/atlas/config.go |  3 +++
 pkg/atlas/server.go | 46 +++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/pkg/atlas/config.go b/pkg/atlas/config.go
index cc48a9c..4300d04 100644
--- a/pkg/atlas/config.go
+++ b/pkg/atlas/config.go
@@ -110,6 +110,9 @@ type Config struct {
 	// "none", region maps are disabled. Options: none, default.
 	API0_RegionMap string `env:"ATLAS_API0_REGION_MAP?=default"`
 
+	// Region mapping overrides. Comma-separated list of prefix=region.
+	API0_RegionMap_Override []string `env:"ATLAS_API0_REGION_MAP_OVERRIDE"`
+
 	// The time after registration for a gameserver to complete verification by.
 	API0_ServerList_VerifyTime time.Duration `env:"ATLAS_API0_SERVERLIST_VERIFY_TIME=10s"`
 
diff --git a/pkg/atlas/server.go b/pkg/atlas/server.go
index fe52441..e780b87 100644
--- a/pkg/atlas/server.go
+++ b/pkg/atlas/server.go
@@ -693,15 +693,55 @@ func configureIP2Location(c *Config) (*ip2xMgr, error) {
 	return mgr, mgr.Load(c.IP2Location)
 }
 
-func configureRegionMap(c *Config) (func(netip.Addr, ip2x.Record) (string, error), error) {
+func configureRegionMap(c *Config) (fn func(netip.Addr, ip2x.Record) (string, error), err error) {
 	switch m := c.API0_RegionMap; m {
 	case "", "none":
-		return nil, nil
+		fn = nil
 	case "default":
-		return regionmap.GetRegion, nil
+		fn = regionmap.GetRegion
 	default:
 		return nil, fmt.Errorf("unknown region map type %q", m)
 	}
+	if len(c.API0_RegionMap_Override) != 0 {
+		type regionMapOverride struct {
+			Prefix netip.Prefix
+			Region string
+		}
+		var mos []regionMapOverride
+		for _, x := range c.API0_RegionMap_Override {
+			a, r, ok := strings.Cut(x, "=")
+			if !ok {
+				return nil, fmt.Errorf("parse region override %q: missing equals sign", x)
+			}
+			if strings.ContainsRune(a, '/') {
+				if pfx, err := netip.ParsePrefix(a); err == nil {
+					mos = append(mos, regionMapOverride{pfx, r})
+				} else {
+					return nil, fmt.Errorf("parse region override %q: invalid prefix: %w", x, err)
+				}
+			} else {
+				if x, err := netip.ParseAddr(a); err == nil {
+					if pfx, err := x.Prefix(x.BitLen()); err == nil {
+						mos = append(mos, regionMapOverride{pfx, r})
+					} else {
+						panic(err)
+					}
+				} else {
+					return nil, fmt.Errorf("parse region override %q: invalid prefix: %w", x, err)
+				}
+			}
+		}
+		next := fn
+		fn = func(a netip.Addr, r ip2x.Record) (string, error) {
+			for _, mo := range mos {
+				if mo.Prefix.Contains(a) {
+					return mo.Region, nil
+				}
+			}
+			return next(a, r)
+		}
+	}
+	return
 }
 
 // Run runs the server, shutting it down gracefully when ctx is canceled, then
-- 
cgit v1.2.3