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

Ties gitshallow, httpcache, geoip, ipcohort, and sync/dataset together
into a single tool that resolves an IP against the configured blocklist
and GeoIP databases, with optional --serve mode for production use.

Flags: --git, --blocklist, --whitelist, --geoip-conf, --serve,
--async-load, --format pretty|json, plus -V/--version and help.
Includes /healthz, http.Server timeouts, and stage-timing logs to stderr.
This commit is contained in:
AJ ONeal 2026-04-30 16:49:08 -06:00
parent 8bed31e00e
commit e2e2a9694a
No known key found for this signature in database
3 changed files with 495 additions and 0 deletions

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

@ -0,0 +1,11 @@
module github.com/therootcompany/golib/cmd/check-ip
go 1.26.0
require (
github.com/therootcompany/golib/net/geoip v0.0.0
github.com/therootcompany/golib/net/gitshallow v0.0.0
github.com/therootcompany/golib/net/httpcache v0.0.0
github.com/therootcompany/golib/net/ipcohort v0.0.0
github.com/therootcompany/golib/sync/dataset v0.0.0
)

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

@ -0,0 +1,299 @@
// 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"
"errors"
"flag"
"fmt"
"log"
"net/http"
"os"
"os/signal"
"path/filepath"
"strconv"
"strings"
"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/sync/dataset"
)
const (
defaultBlocklistRepo = "https://github.com/bitwire-it/ipblocklist.git"
refreshInterval = 47 * time.Minute
version = "dev"
)
// 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
RepoURL string
CacheDir string
WhitelistPath string
AsyncLoad bool
// GeoIPBasicAuth is the pre-encoded Authorization header value for
// MaxMind downloads.
GeoIPBasicAuth string
inbound *dataset.View[ipcohort.Cohort]
outbound *dataset.View[ipcohort.Cohort]
whitelist *dataset.View[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.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.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":
fmt.Fprintf(os.Stdout, "check-ip %s\n", version)
os.Exit(0)
case "help", "-help", "--help":
fmt.Fprintf(os.Stdout, "check-ip %s\n\n", version)
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")
}
// 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 = httpcache.BasicAuth(conf.AccountID, conf.LicenseKey)
}
// 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.New(), 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}}
geoSet := dataset.NewSet(
&httpcache.Cacher{
URL: geoip.DownloadBase + "/GeoLite2-City/download?suffix=tar.gz",
Path: filepath.Join(maxmindDir, geoip.TarGzName(geoip.CityEdition)),
MaxAge: 3 * 24 * time.Hour,
Header: authHeader,
},
&httpcache.Cacher{
URL: geoip.DownloadBase + "/GeoLite2-ASN/download?suffix=tar.gz",
Path: filepath.Join(maxmindDir, geoip.TarGzName(geoip.ASNEdition)),
MaxAge: 3 * 24 * time.Hour,
Header: authHeader,
},
)
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, polled for mtime changes.
// A match here overrides any block decision from the blocklists.
var whitelistSet *dataset.Set
var loadWhitelist func()
if cfg.WhitelistPath != "" {
whitelistSet = dataset.NewSet(dataset.PollFiles(cfg.WhitelistPath))
cfg.whitelist = addCohort(whitelistSet, func(_ context.Context) (*ipcohort.Cohort, error) {
return ipcohort.LoadFile(cfg.WhitelistPath)
})
loadWhitelist = func() {
fmt.Fprint(os.Stderr, "Loading whitelist... ")
t := time.Now()
if err := whitelistSet.Load(context.Background()); err != nil {
fmt.Fprintln(os.Stderr)
log.Printf("whitelist: %v", err)
if !asyncServe {
os.Exit(1)
}
return
}
fmt.Fprintf(os.Stderr, "%s (entries=%s)\n",
time.Since(t).Round(time.Millisecond),
commafy(cfg.whitelist.Value().Size()),
)
}
if !asyncServe {
loadWhitelist()
}
}
// 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)
for _, ip := range ips {
cfg.writeText(os.Stdout, cfg.lookup(ip))
}
if cfg.Bind == "" {
return
}
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()
if asyncServe {
go loadBlocklists()
if loadWhitelist != nil {
go loadWhitelist()
}
}
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 whitelistSet != nil {
go whitelistSet.Tick(ctx, refreshInterval, func(err error) {
log.Printf("whitelist 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 c.whitelist != nil && c.whitelist.Value().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
}
// writeText renders res as human-readable plain text.
func (c *IPCheck) writeText(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.writeText(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 c.whitelist != nil {
datasets["whitelist"] = cohortStatus(c.whitelist)
}
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
}