mirror of
https://github.com/therootcompany/golib.git
synced 2026-11-04 06:55:00 +00:00
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:
parent
c2f5dbeeca
commit
bea4c3de0a
36
net/ipcohort/README.md
Normal file
36
net/ipcohort/README.md
Normal 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.
|
||||||
34
net/ipcohort/cmd/ipcohort-contains/format.go
Normal file
34
net/ipcohort/cmd/ipcohort-contains/format.go
Normal 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"`
|
||||||
|
}
|
||||||
185
net/ipcohort/cmd/ipcohort-contains/main.go
Normal file
185
net/ipcohort/cmd/ipcohort-contains/main.go
Normal 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
|
||||||
|
}
|
||||||
87
net/ipcohort/cmd/ipcohort-contains/term.go
Normal file
87
net/ipcohort/cmd/ipcohort-contains/term.go
Normal 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
3
net/ipcohort/go.mod
Normal file
@ -0,0 +1,3 @@
|
|||||||
|
module github.com/therootcompany/golib/net/ipcohort
|
||||||
|
|
||||||
|
go 1.26.0
|
||||||
0
net/ipcohort/go.sum
Normal file
0
net/ipcohort/go.sum
Normal file
264
net/ipcohort/ipcohort.go
Normal file
264
net/ipcohort/ipcohort.go
Normal 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()) // 0–32
|
||||||
|
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
|
||||||
|
}
|
||||||
159
net/ipcohort/ipcohort_test.go
Normal file
159
net/ipcohort/ipcohort_test.go
Normal 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")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
Loading…
x
Reference in New Issue
Block a user