feat(thoth): make ASNChecker instances

Signed-off-by: Xe Iaso <me@xeiaso.net>
This commit is contained in:
Xe Iaso
2025-05-21 15:51:11 -04:00
parent 86ad85909c
commit 4aea22fac5
3 changed files with 22 additions and 3 deletions
+2 -2
View File
@@ -10,7 +10,7 @@ import (
type ASNChecker struct { type ASNChecker struct {
iptoasn iptoasnv1.IpToASNServiceClient iptoasn iptoasnv1.IpToASNServiceClient
asns map[int]struct{} asns map[uint32]struct{}
hash string hash string
} }
@@ -25,7 +25,7 @@ func (asnc *ASNChecker) Check(r *http.Request) (bool, error) {
return false, err return false, err
} }
_, ok := asnc.asns[int(ipInfo.GetAsNumber())] _, ok := asnc.asns[uint32(ipInfo.GetAsNumber())]
return ok, nil return ok, nil
} }
+1 -1
View File
@@ -16,7 +16,7 @@ func TestASNChecker(t *testing.T) {
asnc := &ASNChecker{ asnc := &ASNChecker{
iptoasn: cli.iptoasn, iptoasn: cli.iptoasn,
asns: map[int]struct{}{ asns: map[uint32]struct{}{
13335: {}, 13335: {},
}, },
hash: "foobar", hash: "foobar",
+19
View File
@@ -4,8 +4,11 @@ import (
"context" "context"
"crypto/tls" "crypto/tls"
"fmt" "fmt"
"strings"
"time" "time"
"github.com/TecharoHQ/anubis/internal"
"github.com/TecharoHQ/anubis/lib/policy"
iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1" iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
grpcprom "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus" grpcprom "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/timeout" "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 { func (c *Client) Close() error {
return c.conn.Close() 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()),
}
}