tamabox/internal/db/ipgeo.go

// Copyright 2026 tamakyi. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package db

import (
	"context"
	"time"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// IPGeoRecord 是一个来源 IP 的地理定位缓存。
//
// 为什么必须落库而不是每次现查：
//   - 免费 IP 库都有配额（ip-api 免费版 45 次/分钟）。一次周报可能有几百上千个
//     不同来源 IP，现查必然撞限流，撞了就是「报告里一半 IP 没有位置」；
//   - 同一个攻击者会在多份报告里反复出现，缓存能让第二次之后零外部请求；
//   - 定位结果本身变化很慢（国家级几乎不变，城市级极少变），没必要重复付费。
//
// 缓存键用 shared.LimitKey（IPv4 精确、IPv6 折叠 /64）而不是原始 IP：
// 攻击者常在一个 /64 里轮换地址，按精确地址缓存会把同一个 C 段查上几十次，
// 既浪费配额又得到一模一样的结果。
type IPGeoRecord struct {
	// IPKey 是缓存键（shared.LimitKey 的产物）。主键，不用自增 id：
	// 这张表是纯 KV 缓存，没有任何按 id 遍历或引用的需求。
	IPKey string `gorm:"size:64;primarykey"`

	Country     string `gorm:"size:64"`
	CountryCode string `gorm:"size:8"`
	Region      string `gorm:"size:128"`
	City        string `gorm:"size:128"`
	// ISP 运营商 / 网络名（"China Telecom"、"Amazon.com, Inc."）。肉鸡报告里
	// 这个字段比国家更有信息量：能看出是家宽还是云主机。
	ISP string `gorm:"size:160"`
	// ASN 自治域号（"AS4134"）。同 ASN 大量命中通常意味着同一片被控网络。
	ASN string `gorm:"size:64"`

	// Lat / Lon 为定位坐标。坐标系随来源不同：
	//   tencent → GCJ-02（火星坐标，与腾讯地图对齐）
	//   ipapi   → WGS-84
	// 两者相差几百米，在世界地图气泡图尺度上无影响；放大到街道级才有意义。
	Lat float64
	Lon float64
	// HasCoords 标记坐标是否可用。IP 库对国家级、部分省级结果不给坐标，
	// 用 0,0 表示「未知」会被画到几内亚湾去 —— 必须显式区分。
	HasCoords bool

	// Source 记录这条数据来自哪个库（tencent / ipapi），便于事后判断精度。
	Source string `gorm:"size:16"`
	// Failed 标记查询失败。失败也要缓存：某个 IP 查不出来通常是确定性的
	// （保留地址、库未收录），反复重试只会白白消耗配额。
	Failed bool

	ResolvedAt time.Time `gorm:"index"`
}

// TableName 固定表名。
func (IPGeoRecord) TableName() string { return "ip_geo_cache" }

// IPGeos 地理定位缓存全局 store。
var IPGeos IPGeoStore

// IPGeoStore 地理定位缓存的读写。
type IPGeoStore interface {
	// GetMany 批量取缓存，返回 map[IPKey]*IPGeoRecord。未命中的键不出现在结果里。
	GetMany(ctx context.Context, keys []string) (map[string]*IPGeoRecord, error)
	// UpsertMany 批量写入或覆盖。
	UpsertMany(ctx context.Context, records []*IPGeoRecord) error
	// PurgeResolvedBefore 删除 resolved_at 早于 before 的记录（缓存过期）。
	//
	// 注意「过期」与「失败」的 TTL 不同，由调用方分两次调：
	// 成功结果留得久（90 天），失败结果留得短（3 天）—— 失败可能是临时故障
	// （对方 API 超时），留太久等于永久放弃这个 IP。
	PurgeResolvedBefore(ctx context.Context, before time.Time, onlyFailed bool) (int64, error)
	// Count 返回缓存条目数。
	Count(ctx context.Context) (int64, error)
}

type ipGeos struct {
	db *gorm.DB
}

// NewIPGeoStore 构造缓存 store。务必传 db.DB 而不是复用什么包级单例。
func NewIPGeoStore(db *gorm.DB) IPGeoStore {
	return &ipGeos{db: db}
}

func (s *ipGeos) GetMany(ctx context.Context, keys []string) (map[string]*IPGeoRecord, error) {
	out := make(map[string]*IPGeoRecord, len(keys))
	if len(keys) == 0 {
		return out, nil
	}

	// 分批 IN 查询：PostgreSQL 的占位符上限是 65535，MySQL 受 max_allowed_packet
	// 约束，一次塞上千个键虽然多数情况能过，但没必要赌。
	const batch = 500
	for start := 0; start < len(keys); start += batch {
		end := start + batch
		if end > len(keys) {
			end = len(keys)
		}
		var rows []*IPGeoRecord
		if err := s.db.WithContext(ctx).
			Where("ip_key IN ?", keys[start:end]).
			Find(&rows).Error; err != nil {
			return nil, err
		}
		for _, r := range rows {
			out[r.IPKey] = r
		}
	}
	return out, nil
}

func (s *ipGeos) UpsertMany(ctx context.Context, records []*IPGeoRecord) error {
	if len(records) == 0 {
		return nil
	}
	// 只更新定位字段，不更新 ip_key 本身（它是冲突判定列）。
	cols := []string{
		"country", "country_code", "region", "city", "isp", "asn",
		"lat", "lon", "has_coords", "source", "failed", "resolved_at",
	}
	return s.db.WithContext(ctx).
		Clauses(clause.OnConflict{
			Columns:   []clause.Column{{Name: "ip_key"}},
			DoUpdates: clause.AssignmentColumns(cols),
		}).
		CreateInBatches(records, 100).Error
}

func (s *ipGeos) PurgeResolvedBefore(ctx context.Context, before time.Time, onlyFailed bool) (int64, error) {
	q := s.db.WithContext(ctx).Where("resolved_at < ?", before)
	if onlyFailed {
		q = q.Where("failed = ?", true)
	} else {
		// 成功记录那一轮不要连失败的一起删：两者的 TTL 不同。
		q = q.Where("failed = ?", false)
	}
	res := q.Delete(&IPGeoRecord{})
	return res.RowsAffected, res.Error
}

func (s *ipGeos) Count(ctx context.Context) (int64, error) {
	var n int64
	err := s.db.WithContext(ctx).Model(&IPGeoRecord{}).Count(&n).Error
	return n, err
}