feat: wire up asn and geoip checkers

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

internal/thoth/asnchecker.go

@@ -15,7 +15,7 @@ type ASNChecker struct {
 }
 
 func (asnc *ASNChecker) Check(r *http.Request) (bool, error) {
-	ctx, cancel := context.WithTimeout(r.Context(), 50*time.Millisecond)
+	ctx, cancel := context.WithTimeout(r.Context(), 500*time.Millisecond)
 	defer cancel()
 
 	ipInfo, err := asnc.iptoasn.Lookup(ctx, &iptoasnv1.LookupRequest{
@@ -25,6 +25,10 @@ func (asnc *ASNChecker) Check(r *http.Request) (bool, error) {
 		return false, err
 	}
 
+	if !ipInfo.GetAnnounced() {
+		return false, nil
+	}
+
 	_, ok := asnc.asns[uint32(ipInfo.GetAsNumber())]
 
 	return ok, nil

internal/thoth/asnchecker_test.go

@@ -5,11 +5,11 @@ import (
 	"net/http/httptest"
 	"testing"
 
-	"github.com/TecharoHQ/anubis/lib/policy"
+	"github.com/TecharoHQ/anubis/lib/policy/checker"
 	iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
 )
 
-var _ policy.Checker = &ASNChecker{}
+var _ checker.Impl = &ASNChecker{}
 
 func TestASNChecker(t *testing.T) {
 	cli := loadSecrets(t)

internal/thoth/cachediptoasn.go

@@ -18,10 +18,36 @@ type IPToASNWithCache struct {
 }
 
 func NewIpToASNWithCache(next iptoasnv1.IpToASNServiceClient) *IPToASNWithCache {
-	return &IPToASNWithCache{
+	result := &IPToASNWithCache{
 		next:  next,
 		table: &bart.Table[*iptoasnv1.LookupResponse]{},
 	}
+
+	for _, pfx := range []netip.Prefix{
+		netip.MustParsePrefix("10.0.0.0/8"),         // RFC 1918
+		netip.MustParsePrefix("172.16.0.0/12"),      // RFC 1918
+		netip.MustParsePrefix("192.168.0.0/16"),     // RFC 1918
+		netip.MustParsePrefix("127.0.0.0/8"),        // Loopback
+		netip.MustParsePrefix("169.254.0.0/16"),     // Link-local
+		netip.MustParsePrefix("100.64.0.0/10"),      // CGNAT
+		netip.MustParsePrefix("192.0.0.0/24"),       // Protocol assignments
+		netip.MustParsePrefix("192.0.2.0/24"),       // TEST-NET-1
+		netip.MustParsePrefix("198.18.0.0/15"),      // Benchmarking
+		netip.MustParsePrefix("198.51.100.0/24"),    // TEST-NET-2
+		netip.MustParsePrefix("203.0.113.0/24"),     // TEST-NET-3
+		netip.MustParsePrefix("240.0.0.0/4"),        // Reserved
+		netip.MustParsePrefix("255.255.255.255/32"), // Broadcast
+		netip.MustParsePrefix("fc00::/7"),           // Unique local address
+		netip.MustParsePrefix("fe80::/10"),          // Link-local
+		netip.MustParsePrefix("::1/128"),            // Loopback
+		netip.MustParsePrefix("::/128"),             // Unspecified
+		netip.MustParsePrefix("100::/64"),           // Discard-only
+		netip.MustParsePrefix("2001:db8::/32"),      // Documentation
+	} {
+		result.table.Insert(pfx, &iptoasnv1.LookupResponse{Announced: false})
+	}
+
+	return result
 }
 
 func (ip2asn *IPToASNWithCache) Lookup(ctx context.Context, lr *iptoasnv1.LookupRequest, opts ...grpc.CallOption) (*iptoasnv1.LookupResponse, error) {

internal/thoth/geoipchecker.go

@@ -16,7 +16,7 @@ type GeoIPChecker struct {
 }
 
 func (gipc *GeoIPChecker) Check(r *http.Request) (bool, error) {
-	ctx, cancel := context.WithTimeout(r.Context(), 50*time.Millisecond)
+	ctx, cancel := context.WithTimeout(r.Context(), 500*time.Millisecond)
 	defer cancel()
 
 	ipInfo, err := gipc.iptoasn.Lookup(ctx, &iptoasnv1.LookupRequest{
@@ -26,6 +26,10 @@ func (gipc *GeoIPChecker) Check(r *http.Request) (bool, error) {
 		return false, err
 	}
 
+	if !ipInfo.GetAnnounced() {
+		return false, nil
+	}
+
 	_, ok := gipc.countries[strings.ToLower(ipInfo.GetCountryCode())]
 
 	return ok, nil

internal/thoth/thoth.go

@@ -71,7 +71,14 @@ func New(ctx context.Context, thothURL, apiToken string) (*Client, error) {
 }
 
 func (c *Client) Close() error {
-	return c.conn.Close()
+	if c.conn != nil {
+		return c.conn.Close()
+	}
+	return nil
+}
+
+func (c *Client) WithIPToASNService(impl iptoasnv1.IpToASNServiceClient) {
+	c.iptoasn = impl
 }
 
 func (c *Client) ASNCheckerFor(asns []uint32) checker.Impl {

internal/thoth/thothmock/iptoasn.go

@@ -0,0 +1,44 @@
+package thothmock
+
+import (
+	"context"
+
+	iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func MockIpToASNService() *IpToASNService {
+	responses := map[string]*iptoasnv1.LookupResponse{
+		"1.1.1.1": {
+			Announced:   true,
+			AsNumber:    13335,
+			Cidr:        []string{"1.1.1.0/24"},
+			CountryCode: "US",
+			Description: "Cloudflare",
+		},
+		"2.2.2.2": {
+			Announced:   true,
+			AsNumber:    420,
+			Cidr:        []string{"2.2.2.0/24"},
+			CountryCode: "CA",
+			Description: "test canada",
+		},
+	}
+
+	return &IpToASNService{Responses: responses}
+}
+
+type IpToASNService struct {
+	Responses map[string]*iptoasnv1.LookupResponse
+}
+
+func (ip2asn *IpToASNService) Lookup(ctx context.Context, lr *iptoasnv1.LookupRequest, opts ...grpc.CallOption) (*iptoasnv1.LookupResponse, error) {
+	resp, ok := ip2asn.Responses[lr.GetIpAddress()]
+	if !ok {
+		return nil, status.Error(codes.NotFound, "IP address not found in mock")
+	}
+
+	return resp, nil
+}