From 796b0b383996209fc1e6f6a96ecd036bcdfac591 Mon Sep 17 00:00:00 2001
From: Guanran Wang <guanran928@outlook.com>
Date: Sat, 19 Oct 2024 16:04:28 +0800
Subject: [PATCH] feat: show ip country locally

---
 assets/index.html |  2 +-
 go.mod            |  7 +++++
 go.sum            | 14 ++++++++++
 main.go           | 67 ++++++++++++++++++++++++++++++++++++++---------
 4 files changed, 76 insertions(+), 14 deletions(-)
 create mode 100644 go.sum

diff --git a/assets/index.html b/assets/index.html
index b05ae4f..2f41b55 100644
--- a/assets/index.html
+++ b/assets/index.html
@@ -27,7 +27,7 @@
           <tr>
             <td class="table-cell">local</td>
             <td class="table-cell">{{.sourceIP}}</td>
-            <td class="table-cell">N/A</td>
+            <td class="table-cell">{{.sourceCountry}}</td>
           </tr>
           <tr>
             <td class="table-cell">ip.sb</td>
diff --git a/go.mod b/go.mod
index a2f3241..e98e720 100644
--- a/go.mod
+++ b/go.mod
@@ -1,3 +1,10 @@
 module ip-checker
 
 go 1.22.5
+
+require github.com/oschwald/geoip2-golang v1.11.0
+
+require (
+	github.com/oschwald/maxminddb-golang v1.13.0 // indirect
+	golang.org/x/sys v0.20.0 // indirect
+)
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..060edde
--- /dev/null
+++ b/go.sum
@@ -0,0 +1,14 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/oschwald/geoip2-golang v1.11.0 h1:hNENhCn1Uyzhf9PTmquXENiWS6AlxAEnBII6r8krA3w=
+github.com/oschwald/geoip2-golang v1.11.0/go.mod h1:P9zG+54KPEFOliZ29i7SeYZ/GM6tfEL+rgSn03hYuUo=
+github.com/oschwald/maxminddb-golang v1.13.0 h1:R8xBorY71s84yO06NgTmQvqvTvlS/bnYZrrWX1MElnU=
+github.com/oschwald/maxminddb-golang v1.13.0/go.mod h1:BU0z8BfFVhi1LQaonTwwGQlsHUEu9pWNdMfmq4ztm0o=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+golang.org/x/sys v0.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y=
+golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/main.go b/main.go
index a7305a7..e59037e 100644
--- a/main.go
+++ b/main.go
@@ -4,27 +4,47 @@ import (
 	"fmt"
 	"html/template"
 	"log"
+	"net"
 	"net/http"
 	"os"
 	"strings"
+	"sync"
+
+	"github.com/oschwald/geoip2-golang"
+)
+
+var (
+	db      *geoip2.Reader
+	tmpl    *template.Template
+	config  ServerConfig
+	ipCache sync.Map
 )
 
 type ServerConfig struct {
-	listen string
+	listen    string
+	countrydb string
 }
 
 func main() {
-	config := ServerConfig{
-		listen: getEnvOr("IP_CHECKER_LISTEN", ":8080"),
+	config = ServerConfig{
+		listen:    getEnvOr("IP_CHECKER_LISTEN", ":8080"),
+		countrydb: os.Getenv("IP_CHECKER_COUNTRY_DB"),
 	}
 
-	tmpl := template.Must(template.ParseFiles("assets/index.html"))
+	if config.countrydb != "" {
+		var err error
+		db, err = geoip2.Open(config.countrydb)
+		if err != nil {
+			log.Fatalf("Error opening GeoIP database: %v", err)
+		}
+		defer db.Close()
+	}
+
+	tmpl = template.Must(template.ParseFiles("assets/index.html"))
 
 	http.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("./assets/static"))))
 
-	http.HandleFunc("/{$}", func(w http.ResponseWriter, r *http.Request) {
-		handleRequest(w, r, tmpl)
-	})
+	http.HandleFunc("/{$}", handleRequest)
 
 	http.HandleFunc("/robots.txt", func(w http.ResponseWriter, r *http.Request) {
 		fmt.Fprintf(w, `User-Agent: *
@@ -38,13 +58,33 @@ Disallow: /harm/to/self
 	log.Fatal(http.ListenAndServe(config.listen, nil))
 }
 
-func handleRequest(w http.ResponseWriter, r *http.Request, tmpl *template.Template) {
+func getIPCountry(ip string) string {
+	if config.countrydb == "" {
+		return "unknown"
+	}
+
+	if country, found := ipCache.Load(ip); found {
+		return country.(string)
+	}
+
+	record, err := db.Country(net.ParseIP(ip))
+	if err != nil {
+		log.Print(err)
+		return "an error occurred"
+	}
+
+	country := record.Country.Names["en"]
+	ipCache.Store(ip, country) // Cache the result
+	return country
+}
+
+func handleRequest(w http.ResponseWriter, r *http.Request) {
 	sourceIP := func() string {
-		if r.Header.Get("CF-Connecting-IP") != "" {
-			return r.Header.Get("CF-Connecting-IP")
+		if ip := r.Header.Get("CF-Connecting-IP"); ip != "" {
+			return ip
 		}
-		if r.Header.Get("X-Forwarded-For") != "" {
-			return r.Header.Get("X-Forwarded-For")
+    if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
+			return ip
 		}
 		return strings.Split(r.RemoteAddr, ":")[0]
 	}()
@@ -67,7 +107,8 @@ func handleRequest(w http.ResponseWriter, r *http.Request, tmpl *template.Templa
 		w.Header().Set("Content-Security-Policy", strings.Join(csp, "; "))
 
 		data := map[string]string{
-			"sourceIP": sourceIP,
+			"sourceIP":      sourceIP,
+			"sourceCountry": getIPCountry(sourceIP),
 		}
 
 		err := tmpl.Execute(w, data)