diff --git a/go.mod b/go.mod index aeb00106..11303384 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( github.com/TecharoHQ/thoth-proto v0.2.0 github.com/a-h/templ v0.3.865 github.com/facebookgo/flagenv v0.0.0-20160425205200-fcd59fca7456 + github.com/gaissmai/bart v0.20.4 github.com/golang-jwt/jwt/v5 v5.2.2 github.com/google/cel-go v0.25.0 github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus v1.0.1 diff --git a/go.sum b/go.sum index 78972927..9537e66a 100644 --- a/go.sum +++ b/go.sum @@ -101,6 +101,8 @@ github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHk github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/fsnotify/fsnotify v1.8.0 h1:dAwr6QBTBZIkG8roQaJjGof0pp0EeF+tNV7YBP3F/8M= github.com/fsnotify/fsnotify v1.8.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0= +github.com/gaissmai/bart v0.20.4 h1:Ik47r1fy3jRVU+1eYzKSW3ho2UgBVTVnUS8O993584U= +github.com/gaissmai/bart v0.20.4/go.mod h1:cEed+ge8dalcbpi8wtS9x9m2hn/fNJH5suhdGQOHnYk= github.com/gliderlabs/ssh v0.3.8 h1:a4YXD1V7xMF9g5nTkdfnja3Sxy1PVDCj1Zg4Wb8vY6c= github.com/gliderlabs/ssh v0.3.8/go.mod h1:xYoytBv1sV0aL3CavoDuJIQNURXkkfPA/wxQ1pL1fAU= github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 h1:+zs/tPmkDkHx3U66DAb0lQFJrpS6731Oaa12ikc+DiI= diff --git a/internal/thoth/asnchecker_test.go b/internal/thoth/asnchecker_test.go index ce164e8b..ba22489b 100644 --- a/internal/thoth/asnchecker_test.go +++ b/internal/thoth/asnchecker_test.go @@ -4,6 +4,8 @@ import ( "fmt" "net/http/httptest" "testing" + + iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1" ) func TestASNChecker(t *testing.T) { @@ -57,3 +59,20 @@ func TestASNChecker(t *testing.T) { }) } } + +func BenchmarkWithCache(b *testing.B) { + cli := loadSecrets(b) + req := &iptoasnv1.LookupRequest{IpAddress: "1.1.1.1"} + + _, err := cli.iptoasn.Lookup(b.Context(), req) + if err != nil { + b.Error(err) + } + + for b.Loop() { + _, err := cli.iptoasn.Lookup(b.Context(), req) + if err != nil { + b.Error(err) + } + } +} diff --git a/internal/thoth/cachediptoasn.go b/internal/thoth/cachediptoasn.go new file mode 100644 index 00000000..6b5d3a0e --- /dev/null +++ b/internal/thoth/cachediptoasn.go @@ -0,0 +1,58 @@ +package thoth + +import ( + "context" + "errors" + "fmt" + "log/slog" + "net/netip" + + iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1" + "github.com/gaissmai/bart" + "google.golang.org/grpc" +) + +type IPToASNWithCache struct { + next iptoasnv1.IpToASNServiceClient + table *bart.Table[*iptoasnv1.LookupResponse] +} + +func NewIpToASNWithCache(next iptoasnv1.IpToASNServiceClient) *IPToASNWithCache { + return &IPToASNWithCache{ + next: next, + table: &bart.Table[*iptoasnv1.LookupResponse]{}, + } +} + +func (ip2asn *IPToASNWithCache) Lookup(ctx context.Context, lr *iptoasnv1.LookupRequest, opts ...grpc.CallOption) (*iptoasnv1.LookupResponse, error) { + addr, err := netip.ParseAddr(lr.GetIpAddress()) + if err != nil { + return nil, fmt.Errorf("input is not an IP address: %w", err) + } + + cachedResponse, ok := ip2asn.table.Lookup(addr) + if ok { + return cachedResponse, nil + } + + resp, err := ip2asn.next.Lookup(ctx, lr, opts...) + if err != nil { + return nil, err + } + + var errs []error + for _, cidr := range resp.GetCidr() { + pfx, err := netip.ParsePrefix(cidr) + if err != nil { + errs = append(errs, err) + continue + } + ip2asn.table.Insert(pfx, resp) + } + + if len(errs) != 0 { + slog.Error("errors parsing IP prefixes", "err", errors.Join(errs...)) + } + + return resp, nil +} diff --git a/internal/thoth/thoth.go b/internal/thoth/thoth.go index 5b6d329c..eb0f36e3 100644 --- a/internal/thoth/thoth.go +++ b/internal/thoth/thoth.go @@ -60,7 +60,7 @@ func New(ctx context.Context, thothURL, apiToken string) (*Client, error) { return &Client{ conn: conn, health: hc, - iptoasn: iptoasnv1.NewIpToASNServiceClient(conn), + iptoasn: NewIpToASNWithCache(iptoasnv1.NewIpToASNServiceClient(conn)), }, nil } diff --git a/internal/thoth/thoth_test.go b/internal/thoth/thoth_test.go index c5f9aacc..53a1d259 100644 --- a/internal/thoth/thoth_test.go +++ b/internal/thoth/thoth_test.go @@ -7,7 +7,7 @@ import ( "github.com/joho/godotenv" ) -func loadSecrets(t *testing.T) *Client { +func loadSecrets(t testing.TB) *Client { if err := godotenv.Load(); err != nil { t.Skip(".env not defined, can't load thoth secrets") }