ip-checker/api/main.go

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/oschwald/geoip2-golang"
)

var (
	asnDB  *geoip2.Reader
	cityDB *geoip2.Reader
	config serverConfig
	cache  sync.Map
)

const (
	devMode  = "dev"
	prodMode = "prod"
)

type serverConfig struct {
	ASNDB  string
	CityDB string
	Listen string
	Mode   string
}

type body struct {
	IP           net.IP `json:"ip"`
	Country      string `json:"country"`
	Region       string `json:"region"`
	City         string `json:"city"`
	ASN          uint   `json:"asn"`
	Organization string `json:"organization"`
}

type bodyError struct {
	Error string `json:"error"`
}

func main() {
	config = serverConfig{
		ASNDB:  os.Getenv("IP_CHECKER_ASN_DB"),
		CityDB: os.Getenv("IP_CHECKER_CITY_DB"),
		Listen: getEnvOr("IP_CHECKER_LISTEN", ":8080"),
		Mode:   os.Getenv("IP_CHECKER_MODE"),
	}

	if config.CityDB != "" {
		var err error
		cityDB, err = geoip2.Open(config.CityDB)
		if err != nil {
			log.Fatalf("Error opening GeoIP City database: %v", err)
		}
		defer cityDB.Close()
	} else {
		log.Fatalln("'IP_CHECKER_CITY_DB' not set")
	}
	if config.ASNDB != "" {
		var err error
		asnDB, err = geoip2.Open(config.ASNDB)
		if err != nil {
			log.Fatalf("Error opening GeoIP ASN database: %v", err)
		}
		defer asnDB.Close()
	} else {
		log.Fatalln("'IP_CHECKER_ASN_DB' not set")
	}

	http.HandleFunc("/api/v1/ip", handleRequest)
	log.Printf("Starting server on %s", config.Listen)
	log.Fatal(http.ListenAndServe(config.Listen, nil))
}

func generateJSON(ip string) (body, error) {
	var data body

	parsedIP := net.ParseIP(ip)

	// check for reserved ip
	if parsedIP == nil {
		return body{}, errors.New("invalid IP")
	} else if parsedIP.IsPrivate() {
		return body{}, errors.New("private IP")
	} else if parsedIP.IsLoopback() {
		return body{}, errors.New("loopback IP")
	} else if parsedIP.IsUnspecified() {
		return body{}, errors.New("unspecified IP")
	}

	// cache
	if cached, found := cache.Load(ip); found && config.Mode == prodMode {
		return cached.(body), nil
	}

	// query
	cityRecord, err := cityDB.City(parsedIP)
	if err != nil {
		log.Printf("GeoIP City lookup failed for IP: %s, error: %v", ip, err)
		return body{}, errors.New("City record lookup failed")
	}

	ASNRecord, err := asnDB.ASN(parsedIP)
	if err != nil {
		log.Printf("GeoIP ASN lookup failed for IP: %s, error: %v", ip, err)
		return body{}, errors.New("ASN record lookup failed")
	}

	data.IP = parsedIP
	data.Country = cityRecord.Country.Names["en"]
	data.Region = cityRecord.Subdivisions[0].Names["en"]
	data.City = cityRecord.City.Names["en"]
	data.ASN = ASNRecord.AutonomousSystemNumber
	data.Organization = ASNRecord.AutonomousSystemOrganization

	if config.Mode == prodMode {
		cache.Store(ip, data)
	}
	return data, nil
}

func handleRequest(w http.ResponseWriter, r *http.Request) {
	sourceIP := func() string {
		if q := r.URL.Query().Get("q"); q != "" {
			return q
		}

		if config.Mode == "dev" {
			return "223.5.5.5"
		}

		if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
			return strings.Split(ip, ",")[0]
		}

		host, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			errorText := fmt.Sprintf("Unable to parse remote address: %v", err)
			log.Printf(errorText)
			return errorText
		}

		return host
	}()

	bodyData, err := generateJSON(sourceIP)
	if err != nil {
		errorBody := bodyError{
			Error: err.Error(),
		}
		errorData, _ := json.MarshalIndent(errorBody, "", "  ")
		http.Error(w, string(errorData)+"\n", http.StatusInternalServerError)
		return
	}
	bodyJson, _ := json.MarshalIndent(bodyData, "", "  ")

	fmt.Fprintf(w, "%s\n", bodyJson)
}

func getEnvOr(key string, defaultValue string) string {
	if value, exists := os.LookupEnv(key); exists {
		return value
	}
	return defaultValue
}