feat(thoth): cached ip to asn checker

Signed-off-by: Xe Iaso <me@xeiaso.net>

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

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=

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)
+		}
+	}
+}

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
+}

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
 }
 

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")
 	}