feat(net/ipcohort): IP block/allow cohort filter

IPv4 host + CIDR membership testing backed by sorted /32 binary search
and a linear scan of CIDR ranges. CSV/text loader with bounded reads
(maxBytes is required so callers pick a limit appropriate to their
input source — 1 KB allowlists through 100+ MB threat feeds). Returns
errors for unparseable IPs so callers can decide fail-open vs
fail-closed; partial cohorts are returned alongside parse errors.
This commit is contained in:
AJ ONeal 2026-05-02 14:34:54 -06:00
parent c2f5dbeeca
commit bea4c3de0a
No known key found for this signature in database
8 changed files with 768 additions and 0 deletions

36
net/ipcohort/README.md Normal file
View File

@ -0,0 +1,36 @@
# [ipcohort](https://github.com/therootcompany/golib/tree/main/net/ipcohort)
A memory-efficient, fast IP cohort checker for blacklists, whitelists, and ad cohorts.
- 4 bytes per /32 host, 5 bytes per CIDR range (8 with alignment padding)
- O(log n) binary search across both hosts and CIDR ranges
- immutable cohorts — callers swap via `atomic.Pointer` for lock-free reads
- requires CIDR sources to be an anti-chain (no entry contained in another;
a /8 supersedes any /24 inside it). Most curated blocklists already
ship pre-normalized this way.
## Example
Check if an IP address belongs to a cohort (such as a blacklist):
```go
cohort, err := ipcohort.LoadFile("/srv/data/inbound.txt")
if err != nil {
log.Fatalf("load: %v", err)
}
blocked, err := cohort.Contains("92.255.85.72")
if err != nil {
log.Fatalf("parse: %v", err)
}
if blocked {
fmt.Println("BLOCKED")
os.Exit(1)
}
fmt.Println("allowed")
```
`Cohort.Contains(string)` parses the address each call and returns an error
on unparseable input — callers decide fail-open vs fail-closed. If you
already have a `netip.Addr` (e.g. from `netip.ParseAddr` on a request peer),
use `Cohort.ContainsAddr(netip.Addr)` to skip the parse and the error.

View File

@ -0,0 +1,34 @@
package main
import (
"fmt"
"os"
)
const (
formatPretty = "pretty"
formatTSV = "tsv"
formatCSV = "csv"
formatJSON = "json"
)
// parseFormat resolves the --format flag value, auto-detecting based on
// stdout (pretty on TTY, tsv when piped) when s is empty.
func parseFormat(s string) (string, error) {
switch s {
case "":
if isTTYish(os.Stdout) {
return formatPretty, nil
}
return formatTSV, nil
case formatPretty, formatTSV, formatCSV, formatJSON:
return s, nil
default:
return "", fmt.Errorf("unknown format %q (want pretty, tsv, csv, json)", s)
}
}
type result struct {
IP string `json:"ip"`
Found bool `json:"found"`
}

View File

@ -0,0 +1,185 @@
// ipcohort-contains checks whether one or more IP addresses appear in a set
// of cohort files (plain text, one IP/CIDR per line).
//
// Usage:
//
// ipcohort-contains [flags] <file>... -- <ip>...
// ipcohort-contains [flags] --ip <ip> <file>...
// echo "<ip>" | ipcohort-contains <file>...
//
// Exit code: 0 if all queried IPs are found, 1 if any are not found, 2 on error.
package main
import (
"bufio"
"errors"
"flag"
"fmt"
"io"
"os"
"strings"
"time"
"github.com/therootcompany/golib/net/ipcohort"
)
// Replaced by goreleaser / ldflags at build time.
var (
name = "ipcohort-contains"
version = "0.0.0-dev"
commit = "0000000"
date = "0001-01-01"
licenseYear = "2021-present"
licenseOwner = "AJ ONeal <aj@therootcompany.com>"
licenseType = "MPL-2.0"
)
type Config struct {
IP string
Format string
}
func main() {
cfg := Config{}
fs := flag.NewFlagSet(os.Args[0], flag.ContinueOnError)
fs.StringVar(&cfg.IP, "ip", "", "IP address to check (alternative to -- separator)")
fs.StringVar(&cfg.Format, "format", "", "output format: pretty, tsv, csv, json (default: auto)")
fs.Usage = func() {
fmt.Fprintf(os.Stderr, "Usage: %s [flags] <file>... -- <ip>...\n", os.Args[0])
fmt.Fprintf(os.Stderr, " %s --ip <ip> <file>...\n", os.Args[0])
fmt.Fprintf(os.Stderr, " echo <ip> | %s <file>...\n", os.Args[0])
fs.PrintDefaults()
fmt.Fprintln(os.Stderr, "Exit: 0=all found, 1=not found, 2=error")
}
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(2)
}
filePaths, ips := splitArgs(fs.Args(), cfg.IP)
if len(filePaths) == 0 {
fmt.Fprintln(os.Stderr, "error: at least one file path required")
fs.Usage()
os.Exit(2)
}
cohort := loadCohort(filePaths)
if len(ips) == 0 {
var err error
ips, err = readIPsFromStdin()
if err != nil {
fmt.Fprintf(os.Stderr, "error reading stdin: %v\n", err)
os.Exit(2)
}
}
if len(ips) == 0 {
fmt.Fprintln(os.Stderr, "error: no IP addresses to check")
fs.Usage()
os.Exit(2)
}
fmt.Fprintln(os.Stderr)
format, err := parseFormat(cfg.Format)
if err != nil {
fmt.Fprintf(os.Stderr, "error: %v\n", err)
os.Exit(2)
}
results, allFound := check(cohort, ips)
if err := writeTable(format, results); err != nil {
fmt.Fprintf(os.Stderr, "error: %v\n", err)
os.Exit(2)
}
if !allFound {
os.Exit(1)
}
}
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)
}
// splitArgs separates positional args into file paths and IPs, using either
// the --ip flag (single IP) or a `--` separator (multiple IPs). When neither
// is set, all args are file paths and IPs come from stdin.
func splitArgs(args []string, ipFlag string) (filePaths, ips []string) {
if ipFlag != "" {
return args, []string{ipFlag}
}
for i, a := range args {
if a == "--" {
return args[:i], args[i+1:]
}
}
return args, nil
}
func loadCohort(filePaths []string) *ipcohort.Cohort {
fmt.Fprint(os.Stderr, "Loading cohort... ")
t := time.Now()
cohort, err := ipcohort.LoadFiles(filePaths...)
if cohort == nil {
fmt.Fprintln(os.Stderr)
fmt.Fprintf(os.Stderr, "error: %v\n", err)
os.Exit(2)
}
if err != nil {
fmt.Fprintln(os.Stderr)
fmt.Fprintf(os.Stderr, "warning: %v\n", err)
}
fmt.Fprintf(os.Stderr, "%s (entries=%s)\n",
time.Since(t).Round(time.Millisecond),
commafy(cohort.Size()),
)
return cohort
}
func readIPsFromStdin() ([]string, error) {
var ips []string
sc := bufio.NewScanner(os.Stdin)
for sc.Scan() {
line := strings.TrimSpace(sc.Text())
if line == "" || strings.HasPrefix(line, "#") {
continue
}
ips = append(ips, line)
}
return ips, sc.Err()
}
func check(cohort *ipcohort.Cohort, ips []string) (results []result, allFound bool) {
results = make([]result, 0, len(ips))
allFound = true
for _, ip := range ips {
found, err := cohort.Contains(ip)
if err != nil {
fmt.Fprintf(os.Stderr, "error: %v\n", err)
os.Exit(2)
}
results = append(results, result{IP: ip, Found: found})
if !found {
allFound = false
}
}
return results, allFound
}

View File

@ -0,0 +1,87 @@
package main
import (
"encoding/csv"
"encoding/json"
"fmt"
"os"
"strconv"
"strings"
)
// isTTYish reports whether f is a terminal (character 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
}
// writeTable renders results in the chosen format to stdout.
func writeTable(format string, results []result) error {
switch format {
case formatPretty:
w := 0
for _, r := range results {
if len(r.IP) > w {
w = len(r.IP)
}
}
for _, r := range results {
fmt.Printf("%-*s %s\n", w, r.IP, statusLabel(r.Found))
}
case formatTSV:
for _, r := range results {
fmt.Printf("%s\t%s\n", r.IP, statusLabel(r.Found))
}
case formatCSV:
cw := csv.NewWriter(os.Stdout)
_ = cw.Write([]string{"ip", "status"})
for _, r := range results {
if err := cw.Write([]string{r.IP, statusLabel(r.Found)}); err != nil {
return err
}
}
cw.Flush()
return cw.Error()
case formatJSON:
enc := json.NewEncoder(os.Stdout)
enc.SetIndent("", " ")
return enc.Encode(results)
}
return nil
}
func statusLabel(found bool) string {
if found {
return "FOUND"
}
return "NOT FOUND"
}
// commafy renders n with thousands separators, e.g. 1234567 -> "1,234,567".
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()
}

3
net/ipcohort/go.mod Normal file
View File

@ -0,0 +1,3 @@
module github.com/therootcompany/golib/net/ipcohort
go 1.26.0

0
net/ipcohort/go.sum Normal file
View File

264
net/ipcohort/ipcohort.go Normal file
View File

@ -0,0 +1,264 @@
package ipcohort
import (
"cmp"
"encoding/binary"
"encoding/csv"
"errors"
"fmt"
"io"
"net/netip"
"os"
"slices"
"strings"
)
// IPv4Net represents a subnet or single address (/32).
type IPv4Net struct {
networkBE uint32
prefix uint8
}
func NewIPv4Net(ip4be uint32, prefix uint8) IPv4Net {
return IPv4Net{
networkBE: ip4be,
prefix: prefix,
}
}
func (r IPv4Net) Contains(ip uint32) bool {
mask := uint32(0xFFFFFFFF) << (32 - r.prefix)
return (ip & mask) == r.networkBE
}
// Cohort is an immutable, read-only set of IPv4 addresses and subnets.
// Contains is safe for concurrent use without locks.
//
// hosts holds sorted /32 addresses for O(log n) binary search.
// nets holds CIDR ranges (prefix < 32), sorted by networkBE.
//
// Invariant: nets is an anti-chain — no net contains another. Cohort
// sources are expected to ship pre-normalized (a /8 supersedes any /24
// inside it; the /24 is omitted). Under this invariant at most one net
// can match any IP, so Contains does a single binary search + one
// Contains check rather than walking candidates.
type Cohort struct {
hosts []uint32
nets []IPv4Net
}
func sortNets(nets []IPv4Net) {
slices.SortFunc(nets, func(a, b IPv4Net) int {
return cmp.Compare(a.networkBE, b.networkBE)
})
}
// Size returns the total number of entries (hosts + nets).
func (c *Cohort) Size() int {
return len(c.hosts) + len(c.nets)
}
// Contains reports whether ipStr falls within any host or subnet in the cohort.
// Returns an error if ipStr is unparseable; callers decide fail-open vs
// fail-closed for invalid input. IPv6 addresses parse but always return false
// (cohort is IPv4-only).
func (c *Cohort) Contains(ipStr string) (bool, error) {
ip, err := netip.ParseAddr(ipStr)
if err != nil {
return false, fmt.Errorf("parse %q: %w", ipStr, err)
}
return c.ContainsAddr(ip), nil
}
// ContainsAddr reports whether ip falls within any host or subnet in the cohort.
// IPv6 addresses always return false (cohort is IPv4-only).
func (c *Cohort) ContainsAddr(ip netip.Addr) bool {
if !ip.Is4() {
return false
}
ip4 := ip.As4()
ipU32 := binary.BigEndian.Uint32(ip4[:])
if _, found := slices.BinarySearch(c.hosts, ipU32); found {
return true
}
// Under the anti-chain invariant, at most one net can contain ipU32,
// and (when one does) it's the net with the largest networkBE <= ipU32.
// Binary-search for the upper bound, then check the immediate predecessor.
hi, _ := slices.BinarySearchFunc(c.nets, ipU32, func(n IPv4Net, target uint32) int {
if n.networkBE > target {
return 1
}
return -1
})
if hi == 0 {
return false
}
return c.nets[hi-1].Contains(ipU32)
}
// Parse builds a Cohort from a list of IP/CIDR strings. Returns an error
// listing every unparseable entry; the caller decides whether to proceed
// with a partial cohort by inspecting the returned (non-nil) Cohort.
func Parse(prefixList []string) (*Cohort, error) {
var hosts []uint32
var nets []IPv4Net
var errs []error
for i, raw := range prefixList {
ipv4net, err := ParseIPv4(raw)
if err != nil {
errs = append(errs, fmt.Errorf("entry %d %q: %w", i, raw, err))
continue
}
if ipv4net.prefix == 32 {
hosts = append(hosts, ipv4net.networkBE)
} else {
nets = append(nets, ipv4net)
}
}
slices.Sort(hosts)
sortNets(nets)
c := &Cohort{hosts: hosts, nets: nets}
if len(errs) > 0 {
return c, fmt.Errorf("ipcohort.Parse: %d invalid entries: %w",
len(errs), errors.Join(errs...))
}
return c, nil
}
func ParseIPv4(raw string) (ipv4net IPv4Net, err error) {
var ippre netip.Prefix
var ip netip.Addr
if strings.Contains(raw, "/") {
ippre, err = netip.ParsePrefix(raw)
if err != nil {
return ipv4net, err
}
} else {
ip, err = netip.ParseAddr(raw)
if err != nil {
return ipv4net, err
}
ippre = netip.PrefixFrom(ip, 32)
}
addr := ippre.Addr()
if !addr.Is4() {
return ipv4net, fmt.Errorf("not an IPv4 address: %s", raw)
}
ip4 := addr.As4()
prefix := uint8(ippre.Bits()) // 032
return NewIPv4Net(
binary.BigEndian.Uint32(ip4[:]),
prefix,
), nil
}
// LoadFile reads path and parses it into a Cohort. On open or read
// errors, returns nil + error. On parse errors only, returns the partial
// cohort + a joined error.
func LoadFile(path string) (*Cohort, error) {
f, err := os.Open(path)
if err != nil {
return nil, fmt.Errorf("could not load %q: %w", path, err)
}
defer func() { _ = f.Close() }()
return ParseCSV(f)
}
// LoadFiles loads and merges multiple files into one Cohort. Useful when
// hosts and networks are stored in separate files.
//
// On open or read errors, returns nil + error. On parse errors only,
// returns the partial cohort + a joined error so callers can choose to
// proceed with what loaded.
func LoadFiles(paths ...string) (*Cohort, error) {
var hosts []uint32
var nets []IPv4Net
var parseErrs []error
for _, path := range paths {
c, err := LoadFile(path)
if c == nil {
return nil, err
}
if err != nil {
parseErrs = append(parseErrs, fmt.Errorf("%s: %w", path, err))
}
hosts = append(hosts, c.hosts...)
nets = append(nets, c.nets...)
}
slices.Sort(hosts)
sortNets(nets)
c := &Cohort{hosts: hosts, nets: nets}
if len(parseErrs) > 0 {
return c, errors.Join(parseErrs...)
}
return c, nil
}
// ParseCSV parses CSV cohort data from r.
func ParseCSV(r io.Reader) (*Cohort, error) {
csvReader := csv.NewReader(r)
csvReader.FieldsPerRecord = -1
return ReadAll(csvReader)
}
// ReadAll reads CSV records from r and builds a Cohort. Returns an error
// listing every unparseable entry (with line number); the returned Cohort
// is non-nil and contains every entry that did parse, so callers can choose
// to proceed with a partial cohort.
func ReadAll(r *csv.Reader) (*Cohort, error) {
var hosts []uint32
var nets []IPv4Net
var errs []error
for {
record, err := r.Read()
if err == io.EOF {
break
}
if err != nil {
return nil, fmt.Errorf("csv read error: %w", err)
}
if len(record) == 0 {
continue
}
raw := strings.TrimSpace(record[0])
if raw == "" || strings.HasPrefix(raw, "#") {
continue
}
ipv4net, err := ParseIPv4(raw)
if err != nil {
line, _ := r.FieldPos(0)
errs = append(errs, fmt.Errorf("line %d %q: %w", line, raw, err))
continue
}
if ipv4net.prefix == 32 {
hosts = append(hosts, ipv4net.networkBE)
} else {
nets = append(nets, ipv4net)
}
}
slices.Sort(hosts)
sortNets(nets)
c := &Cohort{hosts: hosts, nets: nets}
if len(errs) > 0 {
return c, fmt.Errorf("ipcohort.ReadAll: %d invalid entries: %w",
len(errs), errors.Join(errs...))
}
return c, nil
}

View File

@ -0,0 +1,159 @@
package ipcohort_test
import (
"os"
"path/filepath"
"testing"
"github.com/therootcompany/golib/net/ipcohort"
)
func mustContain(t *testing.T, c *ipcohort.Cohort, ip string) bool {
t.Helper()
found, err := c.Contains(ip)
if err != nil {
t.Fatalf("Contains(%q): unexpected err: %v", ip, err)
}
return found
}
func TestContains_SingleHosts(t *testing.T) {
c, err := ipcohort.Parse([]string{"1.2.3.4", "5.6.7.8", "10.0.0.1"})
if err != nil {
t.Fatal(err)
}
hits := []string{"1.2.3.4", "5.6.7.8", "10.0.0.1"}
misses := []string{"1.2.3.5", "5.6.7.7", "10.0.0.2", "0.0.0.0"}
for _, ip := range hits {
if !mustContain(t, c, ip) {
t.Errorf("expected %s to be in cohort", ip)
}
}
for _, ip := range misses {
if mustContain(t, c, ip) {
t.Errorf("expected %s NOT to be in cohort", ip)
}
}
}
func TestContains_CIDRRanges(t *testing.T) {
c, err := ipcohort.Parse([]string{"10.0.0.0/8", "192.168.1.0/24"})
if err != nil {
t.Fatal(err)
}
hits := []string{
"10.0.0.0", "10.0.0.1", "10.255.255.255",
"192.168.1.0", "192.168.1.1", "192.168.1.254", "192.168.1.255",
}
misses := []string{
"9.255.255.255", "11.0.0.0",
"192.168.0.255", "192.168.2.0",
}
for _, ip := range hits {
if !mustContain(t, c, ip) {
t.Errorf("expected %s to be in cohort (CIDR)", ip)
}
}
for _, ip := range misses {
if mustContain(t, c, ip) {
t.Errorf("expected %s NOT to be in cohort (CIDR)", ip)
}
}
}
// AntiChain: nets sorted by networkBE, none contained in another. Verifies
// the bsearch+single-check path picks the right net for IPs whose binary
// search lands on a non-matching neighbor.
func TestContains_AntiChain(t *testing.T) {
c, err := ipcohort.Parse([]string{
"9.0.0.0/24", // narrow, just before the /8
"10.0.0.0/8", // broad
"20.0.0.0/24", // narrow, after the /8
"30.0.0.0/16", // mid
})
if err != nil {
t.Fatal(err)
}
hits := []string{
"9.0.0.50",
"10.5.5.5", // bsearch lands on 10.0.0.0/8 — match
"10.255.255.255",
"20.0.0.5",
"30.0.99.1",
}
misses := []string{
"8.255.255.255", // before everything
"9.0.1.0", // past 9.0.0.0/24, before /8
"11.0.0.0", // past /8, before 20.0.0.0/24
"20.0.1.0", // past /24
"31.0.0.0", // past /16
}
for _, ip := range hits {
if !mustContain(t, c, ip) {
t.Errorf("expected %s to be in cohort", ip)
}
}
for _, ip := range misses {
if mustContain(t, c, ip) {
t.Errorf("expected %s NOT to be in cohort", ip)
}
}
}
func TestContains_Empty(t *testing.T) {
c, err := ipcohort.Parse(nil)
if err != nil {
t.Fatal(err)
}
if mustContain(t, c, "1.2.3.4") {
t.Error("empty cohort should not contain anything")
}
}
func TestLoadFile(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "list.txt")
content := "# comment\n1.2.3.4\n10.0.0.0/8\n\n5.6.7.8\n"
if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
t.Fatal(err)
}
c, err := ipcohort.LoadFile(path)
if err != nil {
t.Fatal(err)
}
if c.Size() != 3 {
t.Errorf("Size() = %d, want 3", c.Size())
}
if !mustContain(t, c, "1.2.3.4") {
t.Error("missing 1.2.3.4")
}
if !mustContain(t, c, "10.5.5.5") {
t.Error("missing CIDR member 10.5.5.5")
}
}
func TestLoadFiles_Merge(t *testing.T) {
dir := t.TempDir()
f1 := filepath.Join(dir, "singles.txt")
f2 := filepath.Join(dir, "networks.txt")
os.WriteFile(f1, []byte("1.2.3.4\n5.6.7.8\n"), 0o644)
os.WriteFile(f2, []byte("192.168.0.0/24\n"), 0o644)
c, err := ipcohort.LoadFiles(f1, f2)
if err != nil {
t.Fatal(err)
}
if c.Size() != 3 {
t.Errorf("Size() = %d, want 3", c.Size())
}
if !mustContain(t, c, "192.168.0.100") {
t.Error("missing merged CIDR member")
}
}