feat(thoth): cached ip to asn checker

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

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