From bc9f481717861f2063b006498b04f43a95e58e22 Mon Sep 17 00:00:00 2001
From: skidoodle <contact@albert.lol>
Date: Sun, 29 Sep 2024 17:58:21 +0200
Subject: [PATCH] Refactor IP data lookup with caching and enable gzip
 compression

---
 readme.md | 18 ++++++++++++
 server.go | 84 +++++++++++++++++++++++++++++++++----------------------
 2 files changed, 69 insertions(+), 33 deletions(-)

diff --git a/readme.md b/readme.md
index ffd5ade..1617ecf 100644
--- a/readme.md
+++ b/readme.md
@@ -1,7 +1,9 @@
 # ipinfo
+
 `ipinfo` is a powerful and efficient IP information service written in Go. It fetches GeoIP data to provide detailed information about an IP address, including geographical location, ASN, and related network details. The service automatically updates its GeoIP databases to ensure accuracy and reliability.
 
 ## Features
+
 - **IP Geolocation**: Provides city, region, country, continent, and coordinates for any IP address.
 - **ASN Information**: Includes autonomous system number and organization.
 - **Hostname Lookup**: Retrieves the hostname associated with the IP address.
@@ -9,7 +11,9 @@
 - **JSONP Support**: Allows JSONP responses for cross-domain requests.
 
 ## Example Endpoints
+
 ### Get information about an IP address
+
 ```sh
 $ curl https://ip.albert.lol/9.9.9.9
 {
@@ -25,14 +29,18 @@ $ curl https://ip.albert.lol/9.9.9.9
   "loc": "37.8767,-122.2676"
 }
 ```
+
 ### Get specific information (e.g., city) about an IP address
+
 ```sh
 $ curl https://ip.albert.lol/9.9.9.9/city
 {
   "city": "Berkeley"
 }
 ```
+
 ### Use JSONP callback function
+
 ```sh
 $ curl https://test.albert.lol/9.9.9.9?callback=Quad9
 /**/ typeof Quad9 === 'function' && Quad9({
@@ -48,6 +56,7 @@ $ curl https://test.albert.lol/9.9.9.9?callback=Quad9
   "loc": "37.8767,-122.2676"
 });
 ```
+
 ```html
 <script>
 let Quad9 = function(data) {
@@ -58,7 +67,9 @@ let Quad9 = function(data) {
 ```
 
 ## Running Locally
+
 ### With Docker
+
 ```sh
 git clone https://github.com/skidoodle/ipinfo
 cd ipinfo
@@ -71,7 +82,9 @@ docker run \
 -e GEOIPUPDATE_DB_DIR=<> \
 ipinfo:main
 ```
+
 ### Without Docker
+
 ```sh
 git clone https://github.com/skidoodle/ipinfo
 cd ipinfo
@@ -79,7 +92,9 @@ go run .
 ```
 
 ## Deploying
+
 ### Docker Compose
+
 ```yaml
 services:
   ipinfo:
@@ -100,7 +115,9 @@ services:
       retries: 3
       start_period: 40s
 ```
+
 ### Docker Run
+
 ```sh
 docker run \
   -d \
@@ -115,4 +132,5 @@ docker run \
 ```
 
 ## LICENSE
+
 [GPL-3.0](https://github.com/skidoodle/ipinfo/blob/main/license)
diff --git a/server.go b/server.go
index 239604e..581fa60 100644
--- a/server.go
+++ b/server.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"compress/gzip"
 	"encoding/json"
 	"fmt"
 	"log"
@@ -11,7 +12,6 @@ import (
 
 var invalidIPBytes = []byte("Please provide a valid IP address.")
 
-// Struct to hold IP data
 type dataStruct struct {
 	IP        *string `json:"ip"`
 	Hostname  *string `json:"hostname"`
@@ -36,9 +36,27 @@ func startServer() {
 	log.Fatal(http.ListenAndServe(":3000", nil))
 }
 
+type gzipResponseWriter struct {
+	http.ResponseWriter
+	Writer *gzip.Writer
+}
+
+func (w gzipResponseWriter) Write(b []byte) (int, error) {
+	return w.Writer.Write(b)
+}
+
 func handler(w http.ResponseWriter, r *http.Request) {
 	requestedThings := strings.Split(r.URL.Path, "/")
 
+	// Enable gzip compression if requested
+	if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
+		w.Header().Set("Content-Encoding", "gzip")
+		gz := gzip.NewWriter(w)
+		defer gz.Close()
+		w = &gzipResponseWriter{Writer: gz, ResponseWriter: w}
+	}
+
+	// Extract IP and field
 	var IPAddress, field string
 	if len(requestedThings) > 1 && net.ParseIP(requestedThings[1]) != nil {
 		IPAddress = requestedThings[1]
@@ -46,19 +64,23 @@ func handler(w http.ResponseWriter, r *http.Request) {
 			field = requestedThings[2]
 		}
 	} else if len(requestedThings) > 1 {
-		field = requestedThings[1]
+		IPAddress = requestedThings[1] // This might be an invalid IP
 	}
 
-	if IPAddress == "" || IPAddress == "self" {
+	// Check if the IP is the client's IP
+	if IPAddress == "" {
 		IPAddress = getRealIP(r)
 	}
+
+	// Validate the IP address
 	ip := net.ParseIP(IPAddress)
 	if ip == nil {
-		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
-		w.Write(invalidIPBytes)
+		// Send 400 Bad Request with invalid IP message
+		http.Error(w, string(invalidIPBytes), http.StatusBadRequest)
 		return
 	}
 
+	// Check if the IP is a bogon IP
 	if isBogon(ip) {
 		w.Header().Set("Content-Type", "application/json; charset=utf-8")
 		bogonData := bogonDataStruct{
@@ -69,13 +91,15 @@ func handler(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// Look up IP data
 	data := lookupIPData(ip)
 	if data == nil {
-		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
-		w.Write(invalidIPBytes)
+		// Send 400 Bad Request with invalid IP message
+		http.Error(w, string(invalidIPBytes), http.StatusBadRequest)
 		return
 	}
 
+	// Handle specific field requests
 	if field != "" {
 		value := getField(data, field)
 		if value != nil {
@@ -83,12 +107,14 @@ func handler(w http.ResponseWriter, r *http.Request) {
 			json.NewEncoder(w).Encode(map[string]*string{field: value})
 			return
 		} else {
+			// Handle invalid field request
 			w.Header().Set("Content-Type", "application/json; charset=utf-8")
 			json.NewEncoder(w).Encode(map[string]*string{field: nil})
 			return
 		}
 	}
 
+	// If no specific field is requested, return the whole data
 	w.Header().Set("Content-Type", "application/json; charset=utf-8")
 	callback := r.URL.Query().Get("callback")
 	enableJSONP := callback != "" && len(callback) < 2000 && callbackJSONP.MatchString(callback)
@@ -106,30 +132,22 @@ func handler(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
-// Get specific field from dataStruct
-func getField(data *dataStruct, field string) *string {
-	switch field {
-	case "ip":
-		return data.IP
-	case "hostname":
-		return data.Hostname
-	case "asn":
-		return data.ASN
-	case "org":
-		return data.Org
-	case "city":
-		return data.City
-	case "region":
-		return data.Region
-	case "country":
-		return data.Country
-	case "continent":
-		return data.Continent
-	case "timezone":
-		return data.Timezone
-	case "loc":
-		return data.Loc
-	default:
-		return nil
-	}
+var fieldMap = map[string]func(*dataStruct) *string{
+	"ip":        func(d *dataStruct) *string { return d.IP },
+	"hostname":  func(d *dataStruct) *string { return d.Hostname },
+	"asn":       func(d *dataStruct) *string { return d.ASN },
+	"org":       func(d *dataStruct) *string { return d.Org },
+	"city":      func(d *dataStruct) *string { return d.City },
+	"region":    func(d *dataStruct) *string { return d.Region },
+	"country":   func(d *dataStruct) *string { return d.Country },
+	"continent": func(d *dataStruct) *string { return d.Continent },
+	"timezone":  func(d *dataStruct) *string { return d.Timezone },
+	"loc":       func(d *dataStruct) *string { return d.Loc },
+}
+
+func getField(data *dataStruct, field string) *string {
+	if f, ok := fieldMap[field]; ok {
+		return f(data)
+	}
+	return nil
 }