geoipapi/geoipapi.go

// Package main implements a GeoIP API service that provides MaxMind GeoIP data over HTTP.
//
// This service is designed to run as a small daemon within Kubernetes clusters
// to serve geolocation data to other services. It exposes HTTP endpoints for
// retrieving country codes and full GeoIP data for given IP addresses.
//
// The service supports OpenTelemetry tracing and automatic MaxMind database
// discovery in standard system paths.
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/oschwald/geoip2-golang"
	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"

	"go.ntppool.org/common/logger"
	"go.ntppool.org/common/tracing"
	"go.ntppool.org/common/version"
)

// geoType represents the type of MaxMind GeoIP database being accessed.
// Each type corresponds to different levels of geographical detail.
type geoType uint8

const (
	countryDB geoType = iota // Country-level database (GeoIP2-Country, GeoLite2-Country)
	cityDB                   // City-level database with detailed location data (GeoIP2-City, GeoLite2-City)
	asnDB                    // ASN/ISP database for network provider information (GeoIP2-ISP)
)

// dbFiles maps each geoType to the possible MaxMind database filenames.
// The system searches for these files in order of preference.
var dbFiles map[geoType][]string

// init initializes the mapping between database types and their corresponding
// MaxMind database filenames, supporting both commercial and free editions.
func init() {
	dbFiles = map[geoType][]string{
		countryDB: {"GeoIP2-Country.mmdb", "GeoLite2-Country.mmdb"},
		asnDB:     {"GeoIP2-ISP.mmdb"},
		cityDB:    {"GeoIP2-City.mmdb", "GeoLite2-City.mmdb"},
	}
}

// main is the entry point for the GeoIP API service.
//
// When run without arguments, it starts an HTTP server on port 8009 that provides
// GeoIP lookup endpoints. When run with IP addresses as arguments, it operates
// in CLI mode and outputs country codes for each provided IP.
//
// The service automatically sets up OpenTelemetry tracing and searches for
// MaxMind databases in standard system locations.
func main() {
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill, syscall.SIGTERM)
	defer cancel()

	log := logger.FromContext(ctx)

	tpShutdown, err := tracing.SetupSDK(ctx, &tracing.TracerConfig{
		ServiceName: "geoipapi",
	})
	if err != nil {
		log.ErrorContext(ctx, "could not setup tracing", "err", err)
		os.Exit(2)
	}
	defer tpShutdown(context.Background()) // todo: with timeout

	rdr, err := open(cityDB)
	if err != nil {
		log.ErrorContext(ctx, "could not open geodb", "err", err)
		os.Exit(3)
	}

	if len(os.Args) > 1 {
		for _, str := range os.Args[1:] {
			ip := net.ParseIP(str)
			city, err := rdr.City(ip)
			if err != nil {
				fmt.Printf("error looking up %q\n: %s", ip, err)
				continue
			}
			fmt.Printf("%s: %s\n", ip, city.Country.IsoCode)
		}
		os.Exit(0)
	}

	err = setupHTTP(ctx)
	if err != nil {
		log.ErrorContext(ctx, "http server", "err", err)
	}
}

// setupHTTP configures and starts the HTTP server with all routes and middleware.
//
// The server listens on port 8009 and provides three endpoints:
//   - /api/country?ip=X.X.X.X - returns ISO country code
//   - /api/json?ip=X.X.X.X - returns full GeoIP data as JSON
//   - /healthz - health check with actual database lookup
//
// The server includes OpenTelemetry tracing (excluding health checks),
// version headers, and graceful shutdown support.
func setupHTTP(ctx context.Context) error {
	mux := http.NewServeMux()
	mux.HandleFunc("/api/country", handleCountry)
	mux.HandleFunc("/api/json", handleJSON)
	mux.HandleFunc("/healthz", handleHealth)

	versionHandler := func(next http.Handler) http.Handler {
		vinfo := version.VersionInfo()
		v := "geoipapi/" + vinfo.Version + "+" + vinfo.GitRevShort
		return http.HandlerFunc(
			func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Server", v)
				span := trace.SpanFromContext(r.Context())
				w.Header().Set("Traceparent", span.SpanContext().TraceID().String())
				next.ServeHTTP(w, r)
			})
	}

	srv := &http.Server{
		Addr:         ":8009",
		BaseContext:  func(_ net.Listener) context.Context { return ctx },
		ReadTimeout:  time.Second,
		WriteTimeout: 10 * time.Second,
		Handler: otelhttp.NewHandler(
			versionHandler(mux),
			"geoipapi",
			otelhttp.WithFilter(func(r *http.Request) bool {
				return r.URL.Path != "/healthz"
			}),
		),
	}
	srvErr := make(chan error, 1)
	go func() {
		srvErr <- srv.ListenAndServe()
	}()

	select {
	case err := <-srvErr:
		// Error when starting HTTP server.
		return err
	case <-ctx.Done():
	}

	return srv.Shutdown(context.Background())
}

// getCityIP retrieves comprehensive GeoIP city data for the given IP address.
//
// This function opens a MaxMind city database and performs a lookup to get
// detailed location information including country, city, coordinates, and
// administrative divisions. It logs warnings for lookup failures but returns
// a generic error message to avoid exposing internal details.
func getCityIP(ctx context.Context, ip net.IP) (*geoip2.City, error) {
	rdr, err := open(cityDB)
	if err != nil {
		return nil, err
	}
	city, err := rdr.City(ip)
	if err != nil {
		logger.FromContext(ctx).WarnContext(ctx, "error looking up city", "ip", ip, "err", err)
		return nil, fmt.Errorf("db lookup error")
	}
	return city, nil
}

// getCity extracts an IP address from an HTTP request and retrieves its GeoIP data.
//
// It parses the 'ip' query parameter, validates it as a valid IP address,
// and adds tracing attributes before delegating to getCityIP for the database lookup.
// Returns an error if the IP parameter is missing or invalid.
func getCity(req *http.Request) (*geoip2.City, error) {
	ctx := req.Context()
	span := trace.SpanFromContext(ctx)
	req.ParseForm()
	ipStr := req.FormValue("ip")
	span.SetAttributes(attribute.String("ip", ipStr))
	ip := net.ParseIP(ipStr)
	if ip == nil {
		return nil, fmt.Errorf("missing IP address")
	}
	return getCityIP(ctx, ip)
}

// handleJSON handles the /api/json endpoint, returning comprehensive GeoIP data as JSON.
//
// This endpoint provides the complete geoip2.City structure with all available
// location information including country, subdivisions, city, postal code,
// coordinates, and timezone data. The response is the raw MaxMind data structure
// serialized to JSON.
func handleJSON(w http.ResponseWriter, req *http.Request) {
	ctx := req.Context()
	span := trace.SpanFromContext(ctx)
	span.SetName("/api/json")
	city, err := getCity(req)
	if err != nil {
		logger.FromContext(ctx).ErrorContext(ctx, "getCity error ", "err", err)
		http.Error(w, "data error", 500)
		return
	}

	b, err := json.Marshal(&city)
	if err != nil {
		logger.FromContext(ctx).ErrorContext(ctx, "error marshaling JSON", "err", err)
		http.Error(w, "internal error", 500)
		return
	}
	w.WriteHeader(200)
	w.Write(b)
}

// handleCountry handles the /api/country endpoint, returning only the ISO country code.
//
// This endpoint provides a lightweight response containing just the two-letter
// ISO 3166-1 alpha-2 country code in lowercase format (e.g., "us", "gb", "ca").
// This is ideal for applications that only need basic country-level geolocation.
func handleCountry(w http.ResponseWriter, req *http.Request) {
	ctx := req.Context()
	span := trace.SpanFromContext(ctx)
	span.SetName("/api/country")

	city, err := getCity(req)
	if err != nil {
		logger.FromContext(ctx).ErrorContext(ctx, "getCity error ", "err", err)
		http.Error(w, "data error", 500)
		return
	}

	w.WriteHeader(200)
	w.Write([]byte(strings.ToLower(city.Country.IsoCode)))
}

// handleHealth handles the /healthz endpoint for Kubernetes-style health checks.
//
// Unlike a simple "OK" response, this endpoint performs an actual GeoIP lookup
// against a known IP address (199.43.0.43) to verify that the MaxMind database
// is accessible and functional. This provides a more meaningful health check
// that can detect database corruption or missing files.
func handleHealth(w http.ResponseWriter, req *http.Request) {
	ctx := req.Context()
	span := trace.SpanFromContext(ctx)
	span.SetAttributes(attribute.Bool("app.drop_sample", true))
	span.SetName("/healthz")

	ip := net.ParseIP("199.43.0.43")
	city, err := getCityIP(ctx, ip)
	if err != nil {
		logger.FromContext(ctx).WarnContext(ctx, "health check getCity error ", "ip", ip, "err", err)
		http.Error(w, "data error", 500)
		return
	}

	w.WriteHeader(200)
	w.Write([]byte(strings.ToLower(city.Country.IsoCode)))
}

// open opens a MaxMind database of the specified type and returns a reader.
//
// It searches through the standard system database directories and looks for
// the appropriate database files based on the geoType. Returns an error if
// no suitable database file is found in any of the searched locations.
func open(t geoType) (*geoip2.Reader, error) {
	dir := findDB()

	var fileName string

	found := false
	for _, f := range dbFiles[t] {
		fileName = filepath.Join(dir, f)
		if _, err := os.Stat(fileName); err == nil {
			found = true
			break
		}
	}
	if !found {
		return nil, fmt.Errorf("could not find '%s' in '%s'", dbFiles[t], dir)
	}

	rdr, err := geoip2.Open(fileName)

	return rdr, err
}

// findDB searches for MaxMind database directories in standard system locations.
//
// It checks common installation paths used by package managers and manual installs:
//   - /usr/share/GeoIP/ (Linux distributions)
//   - /usr/local/share/GeoIP/ (FreeBSD, source installs)
//   - /opt/local/share/GeoIP/ (MacPorts)
//   - /opt/homebrew/var/GeoIP/ (Homebrew on Apple Silicon)
//
// Returns the first existing directory found, or empty string if none exist.
func findDB() string {
	dirs := []string{
		"/usr/share/GeoIP/",        // Linux default
		"/usr/share/local/GeoIP/",  // source install?
		"/usr/local/share/GeoIP/",  // FreeBSD
		"/opt/local/share/GeoIP/",  // MacPorts
		"/opt/homebrew/var/GeoIP/", // Homebrew

	}
	for _, dir := range dirs {
		if _, err := os.Stat(dir); err != nil {
			if os.IsExist(err) {
				log.Println(err)
			}
			continue
		}
		return dir
	}
	return ""
}