worker/app/models/rkn_ip.go

package models

import (
	"fmt"
	"log"
	"net"
	"strings"

	"github.com/davecgh/go-spew/spew"

	"rocketgit.ru/rsmon/worker/app/models/concerns"
	"rocketgit.ru/rsmon/worker/internal/netaddr"
)

// RknIP stores a single CIDR from the ru-blocked IP list. The cidr column
// uses Postgres's native cidr type so the (>>) containment operator works
// directly inside IsRknIPBlocked queries — GiST index recommended for any
// table large enough to amortize the planner cost; see EnsureRknIndexes.
type RknIP struct {
	concerns.Model
	Network *netaddr.Cidr `json:"network" gorm:"type:cidr;"`
}

// TableName pins the underlying table name.
func (RknIP) TableName() string { return "rkn_ips" }

// PanicOnErr panics with a spew-formatted error dump if err is non-nil.
// Kept for callers that still use the old add-by-row InsertInBatches path
// (none after this commit, but kept in case external scripts reference it).
func PanicOnErr(err error) {
	if err != nil {
		spew.Dump(err)
		panic(err)
	}
}

// FindRknIP returns true iff the IP inside `ip` (treated as a /32 or /128
// host) falls inside any CIDR already stored in rkn_ips. Kept for callers
// that already construct an internal/netaddr.Inet.
func FindRknIP(ip netaddr.Inet) (bool, error) {
	ipstr := ip.Inet.String()
	if !strings.Contains(ipstr, "/") {
		if ip.Inet.To4() != nil {
			ipstr += "/32"
		} else {
			ipstr += "/128"
		}
	}
	_, network, err := net.ParseCIDR(ipstr)
	if err != nil {
		return false, err
	}
	cidr := netaddr.Cidr{Cidr: *network, Valid: true}
	var matched []RknIP
	if err := DB().Raw("SELECT id FROM rkn_ips WHERE network >>= ?", &cidr).Scan(&matched).Error; err != nil {
		return false, err
	}
	return len(matched) > 0, nil
}

// AddRknIP is the legacy per-row insert path. Deprecated: callers should
// invoke ReplaceRknIPs from the rkn updater. Kept around so existing cron
// scripts that import the symbol keep compiling.
func AddRknIP(data []string, count int) {
	tx := DB().Begin()
	for k, dataIP := range data {
		if count > 0 && k > count-1 {
			break
		}
		if !strings.Contains(dataIP, "/") {
			dataIP += "/32"
		}
		_, network, err := net.ParseCIDR(dataIP)
		PanicOnErr(err)

		cidr := netaddr.Cidr{Cidr: *network, Valid: true}
		var rknIPs []RknIP
		err = tx.Raw("SELECT id FROM rkn_ips where network >>= ?", &cidr).Scan(&rknIPs).Error
		PanicOnErr(err)
		if len(rknIPs) == 0 {
			rknIPs = []RknIP{}
			err = tx.Raw("SELECT id FROM rkn_ips where network << ?", &cidr).Scan(&rknIPs).Error
			PanicOnErr(err)
			if len(rknIPs) > 0 {
				for _, r := range rknIPs {
					PanicOnErr(tx.Delete(&r).Error)
				}
			}
			PanicOnErr(tx.Create(&RknIP{Network: &cidr}).Error)
		}
		log.Println("created:", dataIP)
	}
	PanicOnErr(tx.Commit().Error)
}

// ReplaceRknIPs deletes every existing row and bulk-inserts the given CIDRs
// in batches of 1000. Whole operation runs in a single transaction so a
// partially-applied update can never leave the table in a hybrid state.
//
// CIDR.parseCIDR-loop uses net.ParseCIDR to canonicalise the address —
// upstream .dat files occasionally contain range/mask pairs that aren't
// already reduced (e.g. 192.168.0.0/16 spelled as 192.168.5.0/16); the
// Postgres cidr type normalises on insert.
func ReplaceRknIPs(nets []*net.IPNet) error {
	tx := DB().Begin()
	if tx.Error != nil {
		return tx.Error
	}

	rows := make([]RknIP, 0, len(nets))
	seen := make(map[string]struct{}, len(nets))
	for _, n := range nets {
		if n == nil || n.IP == nil {
			continue
		}
		// Canonicalise by routing through net.ParseCIDR. This drops the
		// host bits (a common bug in upstream dumps where a /24 range
		// is written with the .5 host bit set) and stamps the right
		// address family flag for Postgres.
		canonical := n.String()
		if _, parsed, err := net.ParseCIDR(canonical); err == nil {
			n = parsed
			canonical = parsed.String()
		}
		if _, ok := seen[canonical]; ok {
			continue
		}
		seen[canonical] = struct{}{}
		cidr := netaddr.Cidr{Cidr: *n, Valid: true}
		rows = append(rows, RknIP{Network: &cidr})
	}

	// Wipe + re-insert in one transaction.
	if err := tx.Where("1 = 1").Delete(&RknIP{}).Error; err != nil {
		_ = tx.Rollback().Error
		return err
	}
	const batchSize = 1000
	for start := 0; start < len(rows); start += batchSize {
		end := start + batchSize
		if end > len(rows) {
			end = len(rows)
		}
		if err := tx.CreateInBatches(rows[start:end], batchSize).Error; err != nil {
			_ = tx.Rollback().Error
			return err
		}
	}
	return tx.Commit().Error
}

// IsRknIPBlocked returns true iff `ip` (any textual form ParseCIDR accepts)
// falls inside any CIDR stored in the rkn_ips table. The query uses the
// cidr >>= inet containment operator — see EnsureRknIndexes for the GiST
// index that makes this fast at scale.
func IsRknIPBlocked(ip string) (bool, error) {
	ip = strings.TrimSpace(ip)
	if ip == "" {
		return false, nil
	}
	var count int64
	if err := DB().Raw("SELECT COUNT(*) FROM rkn_ips WHERE network >>= ?::inet", ip).Scan(&count).Error; err != nil {
		return false, err
	}
	return count > 0, nil
}

// EnsureRknIndexes creates the indexes that AutoMigrate cannot express —
// the GiST index on rkn_ips.network uses cidr >>= cidr containment (the
// expression index `network` already covers equality and prefix ranges,
// but the planner benefits from a GiST for `network >>= <other cidr>`
// queries against ~30k rows). The unique index on rkn_domains.domain is
// also declared in the GORM tag, this function only adds what GORM can't
// (GiST) and is idempotent so it's safe to call repeatedly during boot
// or migration.
func EnsureRknIndexes() error {
	// GiST on cidr requires the btree_gist contrib — its `cidr_ops`
	// opclass exposes cidr to GiST. CREATE EXTENSION IF NOT EXISTS is
	// idempotent.
	if err := DB().Exec("CREATE EXTENSION IF NOT EXISTS btree_gist").Error; err != nil {
		return fmt.Errorf("ensure btree_gist: %w", err)
	}
	// rkn_ips GiST index on the cidr column supports the >>= containment
	// operator that IsRknIPBlocked uses. Without it a 30k-row table makes
	// every IP check a sequential scan; with it each check is an index
	// probe.
	if err := DB().Exec(
		"CREATE INDEX IF NOT EXISTS idx_rkn_ips_network ON rkn_ips USING gist (network)",
	).Error; err != nil {
		return fmt.Errorf("ensure idx_rkn_ips_network: %w", err)
	}
	return nil
}