From 4aea22fac58f835d08a837f400dd79934e446e76 Mon Sep 17 00:00:00 2001
From: Xe Iaso <me@xeiaso.net>
Date: Wed, 21 May 2025 15:51:11 -0400
Subject: [PATCH] feat(thoth): make ASNChecker instances

Signed-off-by: Xe Iaso <me@xeiaso.net>
---
 internal/thoth/asnchecker.go      |  4 ++--
 internal/thoth/asnchecker_test.go |  2 +-
 internal/thoth/thoth.go           | 19 +++++++++++++++++++
 3 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/internal/thoth/asnchecker.go b/internal/thoth/asnchecker.go
index 2cd69dc9..553b2657 100644
--- a/internal/thoth/asnchecker.go
+++ b/internal/thoth/asnchecker.go
@@ -10,7 +10,7 @@ import (
 
 type ASNChecker struct {
 	iptoasn iptoasnv1.IpToASNServiceClient
-	asns    map[int]struct{}
+	asns    map[uint32]struct{}
 	hash    string
 }
 
@@ -25,7 +25,7 @@ func (asnc *ASNChecker) Check(r *http.Request) (bool, error) {
 		return false, err
 	}
 
-	_, ok := asnc.asns[int(ipInfo.GetAsNumber())]
+	_, ok := asnc.asns[uint32(ipInfo.GetAsNumber())]
 
 	return ok, nil
 }
diff --git a/internal/thoth/asnchecker_test.go b/internal/thoth/asnchecker_test.go
index 61857a63..16e3dc4f 100644
--- a/internal/thoth/asnchecker_test.go
+++ b/internal/thoth/asnchecker_test.go
@@ -16,7 +16,7 @@ func TestASNChecker(t *testing.T) {
 
 	asnc := &ASNChecker{
 		iptoasn: cli.iptoasn,
-		asns: map[int]struct{}{
+		asns: map[uint32]struct{}{
 			13335: {},
 		},
 		hash: "foobar",
diff --git a/internal/thoth/thoth.go b/internal/thoth/thoth.go
index c373e69d..43f2f27d 100644
--- a/internal/thoth/thoth.go
+++ b/internal/thoth/thoth.go
@@ -4,8 +4,11 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"strings"
 	"time"
 
+	"github.com/TecharoHQ/anubis/internal"
+	"github.com/TecharoHQ/anubis/lib/policy"
 	iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
 	grpcprom "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/timeout"
@@ -70,3 +73,19 @@ func New(ctx context.Context, thothURL, apiToken string) (*Client, error) {
 func (c *Client) Close() error {
 	return c.conn.Close()
 }
+
+func (c *Client) ASNCheckerFor(asns []uint32) policy.Checker {
+	asnMap := map[uint32]struct{}{}
+	var sb strings.Builder
+	fmt.Fprintln(&sb, "ASNChecker")
+	for _, asn := range asns {
+		asnMap[asn] = struct{}{}
+		fmt.Fprintln(&sb, "AS", asn)
+	}
+
+	return &ASNChecker{
+		iptoasn: c.iptoasn,
+		asns:    asnMap,
+		hash:    internal.SHA256sum(sb.String()),
+	}
+}