tomswallapi/internal/asnexpand/iplocate_db.go

package asnexpand

import (
	"archive/zip"
	"bufio"
	"bytes"
	"context"
	"encoding/csv"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"
)

// DefaultIPLocateDBURL is the iplocate ip-to-asn CSV database, proxied through
// the artifactapi github remote. Remote proxies are served at
// /api/v1/remote/<name>/<path>. The files are Git-LFS, so the /raw/ path
// redirects to the LFS media host — artifactapi follows that redirect.
const DefaultIPLocateDBURL = "https://artifactapi.k8s.syd1.au.unkin.net/api/v1/remote/github/iplocate/ip-address-databases/raw/main/ip-to-asn/ip-to-asn.csv.zip"

// IPLocateDB expands ASNs by reading iplocate's ip-to-asn CSV (CIDR,asn,...). It
// downloads and indexes the whole database once, then serves every ASN group
// from the in-memory index, rebuilding when the index goes stale. This is far
// cheaper than a per-ASN API call and needs no API key.
type IPLocateDB struct {
	URL  string
	TTL  time.Duration
	HTTP *http.Client

	mu       sync.Mutex
	index    map[string][]string
	loadedAt time.Time
}

// NewIPLocateDB builds a DB-backed expander. An empty url uses the default.
func NewIPLocateDB(url string) *IPLocateDB {
	if url == "" {
		url = DefaultIPLocateDBURL
	}
	return &IPLocateDB{
		URL:  url,
		TTL:  24 * time.Hour,
		HTTP: &http.Client{Timeout: 5 * time.Minute},
	}
}

// Prefixes implements Expander: it returns the CIDRs announced by the ASN.
func (d *IPLocateDB) Prefixes(ctx context.Context, asn string) ([]string, error) {
	if err := d.ensureIndex(ctx); err != nil {
		return nil, err
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.index[normalizeASN(asn)], nil
}

// ensureIndex (re)builds the index when it is missing or stale. On a refresh
// failure it keeps the prior index (fail-safe) and only errors when there is
// nothing cached yet.
func (d *IPLocateDB) ensureIndex(ctx context.Context) error {
	d.mu.Lock()
	fresh := d.index != nil && time.Since(d.loadedAt) < d.TTL
	hadIndex := d.index != nil
	d.mu.Unlock()
	if fresh {
		return nil
	}

	idx, err := d.download(ctx)
	if err != nil {
		if hadIndex {
			return nil
		}
		return err
	}
	d.mu.Lock()
	d.index = idx
	d.loadedAt = time.Now()
	d.mu.Unlock()
	return nil
}

func (d *IPLocateDB) download(ctx context.Context) (map[string][]string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, d.URL, nil)
	if err != nil {
		return nil, err
	}
	resp, err := d.HTTP.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("iplocate db: status %d", resp.StatusCode)
	}
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	return indexFromZip(data)
}

// indexFromZip finds the single .csv in the archive and indexes it.
func indexFromZip(data []byte) (map[string][]string, error) {
	zr, err := zip.NewReader(bytes.NewReader(data), int64(len(data)))
	if err != nil {
		return nil, fmt.Errorf("iplocate db: opening zip: %w", err)
	}
	for _, f := range zr.File {
		if !strings.HasSuffix(f.Name, ".csv") {
			continue
		}
		rc, err := f.Open()
		if err != nil {
			return nil, err
		}
		defer rc.Close()
		return buildIndex(rc)
	}
	return nil, fmt.Errorf("iplocate db: no .csv entry in archive")
}

// buildIndex reads the ip-to-asn CSV (header: network,asn,country_code,name,org,
// domain) and returns a map of ASN -> announced CIDRs.
func buildIndex(r io.Reader) (map[string][]string, error) {
	cr := csv.NewReader(bufio.NewReaderSize(r, 1<<20))
	cr.ReuseRecord = true
	cr.FieldsPerRecord = -1 // tolerate quoted commas (org names)

	header, err := cr.Read()
	if err != nil {
		return nil, fmt.Errorf("iplocate db: reading header: %w", err)
	}
	netIdx := columnIndex(header, "network")
	asnIdx := columnIndex(header, "asn")
	if netIdx < 0 || asnIdx < 0 {
		return nil, fmt.Errorf("iplocate db: header missing network/asn columns: %v", header)
	}

	index := map[string][]string{}
	for {
		rec, err := cr.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, fmt.Errorf("iplocate db: reading row: %w", err)
		}
		if netIdx >= len(rec) || asnIdx >= len(rec) {
			continue
		}
		// Records are reused across Read calls, so clone the retained fields.
		asn := strings.Clone(rec[asnIdx])
		index[asn] = append(index[asn], strings.Clone(rec[netIdx]))
	}
	return index, nil
}

func columnIndex(header []string, name string) int {
	for i, h := range header {
		if strings.EqualFold(strings.TrimSpace(h), name) {
			return i
		}
	}
	return -1
}

// normalizeASN strips an optional AS prefix so "AS13335" and "13335" both match
// the bare numbers in the database.
func normalizeASN(asn string) string {
	return strings.TrimPrefix(strings.ToUpper(strings.TrimSpace(asn)), "AS")
}