Fix country code persistence

3 files changed, 12 insertions, 11 deletions

diff --git a/cmd/server/hits.go b/cmd/server/hits.go
index 9420665..0c4ac3f 100644
--- a/cmd/server/hits.go
+++ b/cmd/server/hits.go
@@ -46,20 +46,21 @@ func handleHitCounter(db sws.CounterStore, mmdbPath string) http.HandlerFunc {
 		}
 		hit.SiteID = site.ID
 		hit.Addr = r.RemoteAddr
-
-		host, _, err := net.SplitHostPort(addr)
-		if err == nil {
+		if strings.Contains(r.RemoteAddr, ":") {
+			hit.Addr, _, err = net.SplitHostPort(r.RemoteAddr)
+		}
+		if err == nil && hit.Addr != "" {
 			var cc *string
-			if v, ok := cache.Get(host); ok {
+			if v, ok := cache.Get(hit.Addr); ok {
 				cc = v.(*string)
 			} else if mmdbPath != "" {
-				cc, _ = fetchCountryCode(mmdbPath, host)
-				if cc != nil {
-					debug("geoip", host, "=>", *cc)
+				if cc, err = fetchCountryCode(mmdbPath, hit.Addr); err != nil {
+					log("geoip lookup failed:", err)
 				}
-				cache.Add(host, cc)
+				cache.Add(hit.Addr, cc)
 			}
 			hit.CountryCode = cc
+			debug("geolocated:", hit.Addr, "to", *hit.CountryCode)
 		}
 
 		if err := db.SaveHit(hit); err != nil {
diff --git a/cmd/server/main.go b/cmd/server/main.go
index 1486f6d..6a4f089 100644
--- a/cmd/server/main.go
+++ b/cmd/server/main.go
@@ -38,7 +38,7 @@ func init() {
 	flag.StringVar(&addr, "listen", "localhost:5000", "listen address")
 	flag.StringVar(&dsn, "dsn", "file:sws.db?cache=shared", "database password")
 	flag.StringVar(&domain, "domain", "stats.userspace.com.au", "stats domain")
-	flag.StringVar(&logFile, "l", "", "log to file")
+	flag.StringVar(&logFile, "log", "", "log to file")
 	flag.StringVar(&override, "override", "", "override path")
 	flag.StringVar(&maxmind, "maxmind", "", "maxmind country DB path")
 	flag.BoolVar(&noMigrate, "no-migrate", false, "disable migrations")
diff --git a/store/sqlite3.go b/store/sqlite3.go
index 1102f9d..a74d7dc 100644
--- a/store/sqlite3.go
+++ b/store/sqlite3.go
@@ -174,9 +174,9 @@ on conflict(hash) do update set last_seen_at = :last_seen_at`,
 
 	"saveHit": `insert into hits (
 site_id, addr, scheme, host, path, query, title, referrer, user_agent_hash,
-view_port, no_script, created_at)
+view_port, country_code, no_script, created_at)
 values (:site_id, :addr, :scheme, :host, :path, :query, :title, :referrer,
-:user_agent_hash, :view_port, :no_script, :created_at)`,
+:user_agent_hash, :view_port, :country_code, :no_script, :created_at)`,
 
 	"hits": `select h.*,
 ua.hash as "ua.hash", ua.name as "ua.name", ua.last_seen_at as "ua.last_seen_at"