geoipapi/geoipapi.go

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"
)

type geoType uint8

const (
	countryDB geoType = iota
	cityDB
	asnDB
)

var dbFiles map[geoType][]string

func init() {
	dbFiles = map[geoType][]string{
		countryDB: {"GeoIP2-Country.mmdb", "GeoLite2-Country.mmdb"},
		asnDB:     {"GeoIP2-ISP.mmdb"},
		cityDB:    {"GeoIP2-City.mmdb", "GeoLite2-City.mmdb"},
	}
}

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)
	}
}

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())
}

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
}

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)
}

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)
}

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)))
}

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)))
}

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
}

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 ""
}