root/golib
1
0
Fork 0
mirror of https://github.com/therootcompany/golib.git synced 2026-11-04 06:55:00 +00:00
Code Issues Projects Releases Wiki Activity

feat(cmd/check-ip): reference CLI exercising ipcohort + geoip + dataset

Browse Source
This commit is contained in:
AJ ONeal 2026-09-09 11:34:21 -06:00
parent 5221e0e1c3
commit 4e85e00d5a
No known key found for this signature in database
4 changed files with 631 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
cmd/check-ip/go.mod Normal file
View File

@ -0,0 +1,21 @@
module github.com/therootcompany/golib/cmd/check-ip
go 1.26.5
require (
github.com/therootcompany/golib/net/geoip v0.5.0
github.com/therootcompany/golib/net/gitshallow v0.9.1
github.com/therootcompany/golib/net/httpcache v0.5.1
github.com/therootcompany/golib/net/ipcohort v0.9.0
github.com/therootcompany/golib/net/iplist v0.8.5
github.com/therootcompany/golib/sync/dataset v0.5.0
)
require (
github.com/oschwald/geoip2-golang v1.13.0 // indirect
github.com/oschwald/maxminddb-golang v1.13.0 // indirect
github.com/therootcompany/golib/https v0.9.0 // indirect
github.com/therootcompany/golib/io/transform/gsheet2csv v1.1.0 // indirect
golang.org/x/sync v0.20.0 // indirect
golang.org/x/sys v0.45.0 // indirect
)

32
cmd/check-ip/go.sum Normal file
View File

@ -0,0 +1,32 @@
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.13.0 h1:Q44/Ldc703pasJeP5V9+aFSZFmBN7DKHbNsSFzQATJI=
github.com/oschwald/geoip2-golang v1.13.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=
github.com/therootcompany/golib/https v0.9.0 h1:D9gh0K6ljkvF805JUbylITn8SHPMo08q5PglQGXsK8c=
github.com/therootcompany/golib/https v0.9.0/go.mod h1:ilExqR/Uh7oB36oy9mwuNhANcFh+zdCfD3uJiZ11YAY=
github.com/therootcompany/golib/io/transform/gsheet2csv v1.1.0 h1:R5zAPun5p3JPlNfcMD9es4LD9sqh6nWDNTflZBeeUJo=
github.com/therootcompany/golib/io/transform/gsheet2csv v1.1.0/go.mod h1:CsntOAhhIMdzSAYuwuRZtfkdoiER9+TfSgYOIDWN71o=
github.com/therootcompany/golib/net/geoip v0.5.0 h1:dmduwHLX3RCe2NExhOKcFZ4m1ML2C3gQKnNyaBhNE88=
github.com/therootcompany/golib/net/geoip v0.5.0/go.mod h1:BLf5g25UhvI/+/x/WpM1NTVs4ZMnjySoHv+LXfALT6M=
github.com/therootcompany/golib/net/gitshallow v0.9.1 h1:VI/Q4jViKvVuJ+lzZrD7suhSXLfNNhUcd4ZJrkpTP/Q=
github.com/therootcompany/golib/net/gitshallow v0.9.1/go.mod h1:S/3OK7wMOqJJjP79U9Qw/l3jjEyrb/zsjJ8HAPAX/Fc=
github.com/therootcompany/golib/net/httpcache v0.5.1 h1:vQM758tBwPEqCo7ojpSeL3fHqgq29LYmvD2hvDTBfKY=
github.com/therootcompany/golib/net/httpcache v0.5.1/go.mod h1:mKSQJWpGSCDBzWkTCwT7dO1SCVCsFmdkRgKr39fNS1U=
github.com/therootcompany/golib/net/ipcohort v0.9.0 h1:j3AaoW2u4XQJz1ya6fixv3pb3KgxIpZV8t0082/kS3E=
github.com/therootcompany/golib/net/ipcohort v0.9.0/go.mod h1:Ljcnt2xOryVDGIVV5Kzt5hj2LhW4lBAWo5OeeVYU9ys=
github.com/therootcompany/golib/net/iplist v0.8.5 h1:swn/hcCt6fAG3tYApmxh1OlDak2PjQdczldZWaH1Eus=
github.com/therootcompany/golib/net/iplist v0.8.5/go.mod h1:vK5Na53RcLdtp/IPM/NX9BudyKhXp+8T+elmWBVWzeU=
github.com/therootcompany/golib/sync/dataset v0.5.0 h1:4i3+7Izy/vnpK6IW5m9u4UqLA9RM1IjUsedbZo3MW6E=
github.com/therootcompany/golib/sync/dataset v0.5.0/go.mod h1:HMa25n3eqYnjRcPXny0IUGU2wFKD5/x2g9UUkt3iAmI=
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY=
golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

393
cmd/check-ip/main.go Normal file
View File

@ -0,0 +1,393 @@
// check-ip runs an HTTP API that reports whether an IP appears in the
// configured blocklist repo and enriches the response with MaxMind
// GeoLite2 City + ASN data.
package main
import (
"context"
"encoding/base64"
"encoding/csv"
"encoding/json"
"errors"
"flag"
"fmt"
"io"
"log"
"net/http"
"os"
"os/signal"
"path/filepath"
"strconv"
"strings"
"sync/atomic"
"syscall"
"time"
"github.com/therootcompany/golib/net/geoip"
"github.com/therootcompany/golib/net/gitshallow"
"github.com/therootcompany/golib/net/httpcache"
"github.com/therootcompany/golib/net/ipcohort"
"github.com/therootcompany/golib/net/iplist"
"github.com/therootcompany/golib/sync/dataset"
)
// Replaced by goreleaser / ldflags at build time.
var (
name = "check-ip"
version = "0.0.0-dev"
commit = "0000000"
date = "0001-01-01"
licenseYear = "2026"
licenseOwner = "AJ ONeal"
licenseType = "MPL-2.0"
)
const (
defaultBlocklistRepo = "https://github.com/bitwire-it/ipblocklist.git"
refreshInterval = 47 * time.Minute
)
// printVersion writes version, copyright, and license info to w.
func printVersion(w io.Writer) {
_, _ = fmt.Fprintf(w, "%s v%s %s (%s)\n", name, version, commit[:7], date)
_, _ = fmt.Fprintf(w, "Copyright (C) %s %s\n", licenseYear, licenseOwner)
_, _ = fmt.Fprintf(w, "Licensed under %s\n", licenseType)
}
// isTTYish reports whether f is a terminal device.
func isTTYish(f *os.File) bool {
fi, err := f.Stat()
if err != nil {
return false
}
m := os.ModeDevice | os.ModeCharDevice
return fi.Mode()&m == m
}
// parseFormat validates and returns the output format.
func parseFormat(s string) (string, error) {
switch s {
case "", "pretty", "tsv", "csv", "json":
return s, nil
default:
return "", fmt.Errorf("invalid format %q: want pretty, tsv, csv, json", s)
}
}
// commafy formats n with comma thousands separators (e.g. 3406727 -> "3,406,727").
func commafy(n int) string {
s := strconv.Itoa(n)
neg := ""
if n < 0 {
neg, s = "-", s[1:]
}
if len(s) <= 3 {
return neg + s
}
var b strings.Builder
head := len(s) % 3
if head > 0 {
b.WriteString(s[:head])
b.WriteByte(',')
}
for i := head; i < len(s); i += 3 {
b.WriteString(s[i : i+3])
if i+3 < len(s) {
b.WriteByte(',')
}
}
return neg + b.String()
}
// IPCheck holds the parsed CLI config and the loaded data sources used by
// the HTTP handler.
type IPCheck struct {
Bind string
GeoIPConfPath string
GeoIPURL string
RepoURL string
CacheDir string
WhitelistPath string
AsyncLoad bool
Format string
// GeoIPBasicAuth is the pre-encoded Authorization header value for
// MaxMind downloads.
GeoIPBasicAuth string
inbound *dataset.View[ipcohort.Cohort]
outbound *dataset.View[ipcohort.Cohort]
whitelist atomic.Pointer[ipcohort.Cohort]
geo *dataset.View[geoip.Databases]
}
func main() {
cfg := IPCheck{}
fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
fs.StringVar(&cfg.Bind, "serve", "", "bind address for the HTTP API, e.g. :8080")
fs.StringVar(&cfg.GeoIPConfPath, "geoip-conf", "", "path to GeoIP.conf (default: ./GeoIP.conf or ~/.config/maxmind/GeoIP.conf)")
fs.StringVar(&cfg.RepoURL, "blocklist-repo", defaultBlocklistRepo, "git URL of the blocklist repo (must match bitwire-it layout)")
fs.StringVar(&cfg.CacheDir, "cache-dir", "", "cache parent dir, holds bitwire-it/ and maxmind/ subdirs (default: ~/.cache)")
fs.StringVar(&cfg.GeoIPURL, "geoip-url", "", "GeoIP download URL (overrides geoip.DownloadBase)")
fs.StringVar(&cfg.WhitelistPath, "whitelist", "", "path to a file of IPs and/or CIDRs (one per line) that override block decisions")
fs.BoolVar(&cfg.AsyncLoad, "async-load", false, "with --serve: start the HTTP server immediately and populate blocklists+whitelist in the background (/healthz returns 503 until ready). Ignored in CLI mode.")
fs.StringVar(&cfg.Format, "format", "", "output format: pretty, tsv, csv, json (default: auto)")
fs.Usage = func() {
fmt.Fprintf(os.Stderr, "Usage: %s [flags] <ip> [ip...]\n", os.Args[0])
fmt.Fprintf(os.Stderr, " %s --serve <bind> [flags]\n", os.Args[0])
fs.PrintDefaults()
}
if len(os.Args) > 1 {
switch os.Args[1] {
case "-V", "-version", "--version", "version":
printVersion(os.Stdout)
os.Exit(0)
case "help", "-help", "--help":
printVersion(os.Stdout)
fmt.Fprintln(os.Stdout)
fs.SetOutput(os.Stdout)
fs.Usage()
os.Exit(0)
}
}
if err := fs.Parse(os.Args[1:]); err != nil {
if errors.Is(err, flag.ErrHelp) {
os.Exit(0)
}
os.Exit(1)
}
ips := fs.Args()
if cfg.Bind == "" && len(ips) == 0 {
fmt.Fprintln(os.Stderr, "error: provide at least one IP argument or --serve <bind>")
fs.Usage()
os.Exit(1)
}
if cfg.CacheDir == "" {
home, err := os.UserHomeDir()
if err != nil {
log.Fatalf("cache-dir: %v", err)
}
cfg.CacheDir = filepath.Join(home, ".cache")
}
format, err := parseFormat(cfg.Format)
if err != nil {
log.Fatalf("format: %v", err)
}
cfg.Format = format
// GeoIP config discovery: explicit --geoip-conf wins; otherwise check the
// default locations. If found, parse it and stash the basic-auth header
// value for later MaxMind downloads.
if cfg.GeoIPConfPath == "" {
for _, p := range geoip.DefaultConfPaths() {
if _, err := os.Stat(p); err == nil {
cfg.GeoIPConfPath = p
break
}
}
}
if cfg.GeoIPConfPath != "" {
data, err := os.ReadFile(cfg.GeoIPConfPath)
if err != nil {
log.Fatalf("geoip-conf: %v", err)
}
conf, err := geoip.ParseConf(string(data))
if err != nil {
log.Fatalf("geoip-conf: %v", err)
}
cfg.GeoIPBasicAuth = "Basic " + base64.StdEncoding.EncodeToString([]byte(conf.AccountID+":"+conf.LicenseKey))
}
// Allow overriding the license key via env var.
if envKey := os.Getenv("MAXMIND_LICENSE_KEY"); envKey != "" {
if cfg.GeoIPConfPath == "" {
log.Fatalf("MAXMIND_LICENSE_KEY requires --geoip-conf for AccountID")
}
cfg.GeoIPBasicAuth = "Basic " + base64.StdEncoding.EncodeToString([]byte(cfg.GeoIPBasicAuth[6:]+":"+envKey))
}
// Blocklists: git repo with inbound + outbound IP cohort files.
repo := gitshallow.New(cfg.RepoURL, filepath.Join(cfg.CacheDir, "bitwire-it"), 1, "")
repo.MaxAge = refreshInterval
blocklists := dataset.NewSet(repo)
asyncServe := cfg.AsyncLoad && cfg.Bind != ""
addCohort := func(s *dataset.Set, loader func(context.Context) (*ipcohort.Cohort, error)) *dataset.View[ipcohort.Cohort] {
if asyncServe {
return dataset.AddInitial(s, &ipcohort.Cohort{}, loader)
}
return dataset.Add(s, loader)
}
cfg.inbound = addCohort(blocklists, func(_ context.Context) (*ipcohort.Cohort, error) {
return ipcohort.LoadFiles(
repo.FilePath("tables/inbound/single_ips.txt"),
repo.FilePath("tables/inbound/networks.txt"),
)
})
cfg.outbound = addCohort(blocklists, func(_ context.Context) (*ipcohort.Cohort, error) {
return ipcohort.LoadFiles(
repo.FilePath("tables/outbound/single_ips.txt"),
repo.FilePath("tables/outbound/networks.txt"),
)
})
loadBlocklists := func() {
fmt.Fprint(os.Stderr, "Loading blocklists... ")
t := time.Now()
if err := blocklists.Load(context.Background()); err != nil {
fmt.Fprintln(os.Stderr)
log.Printf("blocklists: %v", err)
if !asyncServe {
os.Exit(1)
}
return
}
fmt.Fprintf(os.Stderr, "%s (inbound=%s, outbound=%s)\n",
time.Since(t).Round(time.Millisecond),
commafy(cfg.inbound.Value().Size()),
commafy(cfg.outbound.Value().Size()),
)
}
if !asyncServe {
loadBlocklists()
}
// GeoIP: download the City + ASN tar.gz archives via httpcache
// conditional GETs. geoip.Open extracts in-memory — no .mmdb files
// are written to disk.
if cfg.GeoIPBasicAuth == "" {
log.Fatalf("geoip-conf: not found; set --geoip-conf or place GeoIP.conf in a default location.\n"+
"GeoLite2 registration and the AccountID/LicenseKey needed for GeoIP.conf are free at:\n"+
" https://www.maxmind.com/en/geolite2/signup\n"+
"Then create a license key and write ./GeoIP.conf (or ~/.config/maxmind/GeoIP.conf):\n"+
" AccountID <your-account-id>\n"+
" LicenseKey <your-license-key>\n"+
" EditionIDs GeoLite2-City GeoLite2-ASN\n"+
"Default search paths: %v", geoip.DefaultConfPaths())
}
maxmindDir := filepath.Join(cfg.CacheDir, "maxmind")
authHeader := http.Header{"Authorization": []string{cfg.GeoIPBasicAuth}}
downloadBase := geoip.DownloadBase
if cfg.GeoIPURL != "" {
downloadBase = strings.TrimRight(cfg.GeoIPURL, "/")
}
cityCacher := httpcache.New(
downloadBase+"/GeoLite2-City/download?suffix=tar.gz",
filepath.Join(maxmindDir, geoip.TarGzName(geoip.CityEdition)))
cityCacher.Header = authHeader
cityCacher.MaxAge = 3 * 24 * time.Hour
asnCacher := httpcache.New(
downloadBase+"/GeoLite2-ASN/download?suffix=tar.gz",
filepath.Join(maxmindDir, geoip.TarGzName(geoip.ASNEdition)))
asnCacher.Header = authHeader
asnCacher.MaxAge = 3 * 24 * time.Hour
geoSet := dataset.NewSet(cityCacher, asnCacher)
cfg.geo = dataset.Add(geoSet, func(_ context.Context) (*geoip.Databases, error) {
return geoip.Open(maxmindDir)
})
fmt.Fprint(os.Stderr, "Loading geoip... ")
tGeo := time.Now()
if err := geoSet.Load(context.Background()); err != nil {
fmt.Fprintln(os.Stderr)
log.Fatalf("geoip: %v", err)
}
fmt.Fprintf(os.Stderr, "%s\n", time.Since(tGeo).Round(time.Millisecond))
defer func() { _ = cfg.geo.Value().Close() }()
// Whitelist: combined IPs + CIDRs in one file, refreshed on an interval.
// A match here overrides any block decision from the blocklists.
// Uses net/iplist.Source — the refreshed-IP-list primitive net/ippolicy
// builds on — instead of a dataset wrapper; iplist owns the refresh.
if cfg.WhitelistPath != "" {
t := time.Now()
whitelistSrc, err := iplist.NewSource(context.Background(), iplist.SourceConfig{
Source: cfg.WhitelistPath,
RefreshInterval: refreshInterval,
Optional: true,
})
if err != nil {
log.Fatalf("whitelist: %v", err)
}
store := func() {
coh, perr := ipcohort.Parse(whitelistSrc.Entries())
if perr != nil {
log.Printf("whitelist: %v", perr)
}
cfg.whitelist.Store(coh)
}
store()
whitelistSrc.Subscribe(func(e iplist.SourceEvent) {
if e.Err != nil {
log.Printf("whitelist refresh: %v", e.Err)
return
}
store()
})
fmt.Fprintf(os.Stderr, "Loading whitelist... %s (entries=%s)\n",
time.Since(t).Round(time.Millisecond),
commafy(cfg.whitelist.Load().Size()),
)
}
// Blank line separates the stderr "Loading ..." block from the real
// output (stdout results for CLI mode, or the stderr "listening on"
// log for serve mode).
fmt.Fprintln(os.Stderr)
// Auto-detect format: pretty on TTY, tsv when piped.
if cfg.Format == "" {
if isTTYish(os.Stdout) {
cfg.Format = "pretty"
} else {
cfg.Format = "tsv"
}
}
for _, ip := range ips {
res := cfg.lookup(ip)
switch cfg.Format {
case "json":
enc := json.NewEncoder(os.Stdout)
enc.SetIndent("", " ")
_ = enc.Encode(res)
case "csv":
w := csv.NewWriter(os.Stdout)
w.Write([]string{"ip", "blocked", "blocked_inbound", "blocked_outbound", "allowlisted", "city", "region", "country", "country_iso", "asn", "asn_org"})
w.Write([]string{res.IP, strconv.FormatBool(res.Blocked), strconv.FormatBool(res.BlockedInbound), strconv.FormatBool(res.BlockedOutbound), strconv.FormatBool(res.Allowlisted), res.Geo.City, res.Geo.Region, res.Geo.Country, res.Geo.CountryISO, strconv.Itoa(int(res.Geo.ASN)), res.Geo.ASNOrg})
w.Flush()
case "tsv":
fmt.Fprintf(os.Stdout, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n",
res.IP,
strconv.FormatBool(res.Blocked),
strconv.FormatBool(res.BlockedInbound),
strconv.FormatBool(res.BlockedOutbound),
strconv.FormatBool(res.Allowlisted),
res.Geo.City,
res.Geo.Region,
res.Geo.Country,
res.Geo.CountryISO,
strconv.Itoa(int(res.Geo.ASN)),
res.Geo.ASNOrg,
)
default:
cfg.writePretty(os.Stdout, res)
}
}
if cfg.Bind == "" {
return
}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()
if asyncServe {
go loadBlocklists()
}
go blocklists.Tick(ctx, refreshInterval, func(err error) {
log.Printf("blocklists refresh: %v", err)
})
go geoSet.Tick(ctx, refreshInterval, func(err error) {
log.Printf("geoip refresh: %v", err)
})
if err := cfg.serve(ctx); err != nil {
log.Fatalf("serve: %v", err)
}
}

185
cmd/check-ip/server.go Normal file
View File

@ -0,0 +1,185 @@
package main
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log"
"net"
"net/http"
"net/netip"
"strings"
"time"
"github.com/therootcompany/golib/net/geoip"
"github.com/therootcompany/golib/net/ipcohort"
"github.com/therootcompany/golib/sync/dataset"
)
// Result is the JSON verdict for a single IP.
type Result struct {
IP string `json:"ip"`
Blocked bool `json:"blocked"`
BlockedInbound bool `json:"blocked_inbound"`
BlockedOutbound bool `json:"blocked_outbound"`
Allowlisted bool `json:"allowlisted,omitzero"`
Geo geoip.Info `json:"geo,omitzero"`
}
// lookup builds a Result for ip against the currently loaded blocklists
// and GeoIP databases.
func (c *IPCheck) lookup(ip string) Result {
res := Result{IP: ip, Geo: c.geo.Value().Lookup(ip)}
addr, err := netip.ParseAddr(ip)
if err != nil {
res.Blocked = true
res.BlockedInbound = true
res.BlockedOutbound = true
return res
}
if cur := c.whitelist.Load(); cur != nil && cur.ContainsAddr(addr) {
res.Allowlisted = true
return res
}
res.BlockedInbound = c.inbound.Value().ContainsAddr(addr)
res.BlockedOutbound = c.outbound.Value().ContainsAddr(addr)
res.Blocked = res.BlockedInbound || res.BlockedOutbound
return res
}
// writePretty renders res as space-aligned plain text.
func (c *IPCheck) writePretty(w io.Writer, res Result) {
switch {
case res.Allowlisted:
fmt.Fprintf(w, "%s is ALLOWED (whitelist)\n", res.IP)
case res.BlockedInbound && res.BlockedOutbound:
fmt.Fprintf(w, "%s is BLOCKED (inbound + outbound)\n", res.IP)
case res.BlockedInbound:
fmt.Fprintf(w, "%s is BLOCKED (inbound)\n", res.IP)
case res.BlockedOutbound:
fmt.Fprintf(w, "%s is BLOCKED (outbound)\n", res.IP)
default:
fmt.Fprintf(w, "%s is allowed\n", res.IP)
}
var parts []string
if res.Geo.City != "" {
parts = append(parts, res.Geo.City)
}
if res.Geo.Region != "" {
parts = append(parts, res.Geo.Region)
}
if res.Geo.Country != "" {
parts = append(parts, fmt.Sprintf("%s (%s)", res.Geo.Country, res.Geo.CountryISO))
}
if len(parts) > 0 {
fmt.Fprintf(w, " Location: %s\n", strings.Join(parts, ", "))
}
if res.Geo.ASN != 0 {
fmt.Fprintf(w, " ASN: AS%d %s\n", res.Geo.ASN, res.Geo.ASNOrg)
}
}
func (c *IPCheck) handle(w http.ResponseWriter, r *http.Request) {
ip := strings.TrimSpace(r.URL.Query().Get("ip"))
if ip == "" {
if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
first, _, _ := strings.Cut(xff, ",")
ip = strings.TrimSpace(first)
} else if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
ip = host
} else {
ip = r.RemoteAddr
}
}
res := c.lookup(ip)
if r.URL.Query().Get("format") == "json" ||
strings.Contains(r.Header.Get("Accept"), "application/json") {
w.Header().Set("Content-Type", "application/json; charset=utf-8")
enc := json.NewEncoder(w)
enc.SetIndent("", " ")
_ = enc.Encode(res)
return
}
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
c.writePretty(w, res)
}
type dsStatus struct {
Loaded bool `json:"loaded"`
Size int `json:"size,omitzero"`
LoadedAt time.Time `json:"loaded_at,omitzero"`
}
// healthz reports per-dataset load state and an overall ready flag.
// Returns 200 when all required sets (inbound, outbound, geoip) are
// loaded, 503 while any are still empty.
func (c *IPCheck) healthz(w http.ResponseWriter, _ *http.Request) {
cohortStatus := func(v *dataset.View[ipcohort.Cohort]) dsStatus {
s := dsStatus{LoadedAt: v.LoadedAt()}
if cur := v.Value(); cur != nil {
s.Loaded, s.Size = true, cur.Size()
}
return s
}
datasets := map[string]dsStatus{
"inbound": cohortStatus(c.inbound),
"outbound": cohortStatus(c.outbound),
"geoip": {
Loaded: c.geo.Value() != nil,
LoadedAt: c.geo.LoadedAt(),
},
}
if cur := c.whitelist.Load(); cur != nil {
datasets["whitelist"] = dsStatus{Loaded: true, Size: cur.Size()}
}
ready := datasets["inbound"].Loaded && datasets["outbound"].Loaded && datasets["geoip"].Loaded
resp := struct {
Ready bool `json:"ready"`
Version string `json:"version"`
Datasets map[string]dsStatus `json:"datasets"`
}{ready, version, datasets}
w.Header().Set("Content-Type", "application/json; charset=utf-8")
if !ready {
w.WriteHeader(http.StatusServiceUnavailable)
}
enc := json.NewEncoder(w)
enc.SetIndent("", " ")
_ = enc.Encode(resp)
}
func (c *IPCheck) serve(ctx context.Context) error {
mux := http.NewServeMux()
mux.HandleFunc("GET /check", c.handle)
mux.HandleFunc("GET /healthz", c.healthz)
mux.HandleFunc("GET /{$}", c.handle)
srv := &http.Server{
Addr: c.Bind,
Handler: mux,
BaseContext: func(_ net.Listener) context.Context { return ctx },
ReadHeaderTimeout: 5 * time.Second,
ReadTimeout: 10 * time.Second,
WriteTimeout: 15 * time.Second,
IdleTimeout: 60 * time.Second,
MaxHeaderBytes: 1 << 16, // 64KB
}
go func() {
<-ctx.Done()
shutCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
_ = srv.Shutdown(shutCtx)
}()
log.Printf("listening on %s", c.Bind)
if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
return err
}
return nil
}