feat: wire up asn and geoip checkers

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

cmd/anubis/main.go

@@ -30,11 +30,13 @@ import (
 	"github.com/TecharoHQ/anubis"
 	"github.com/TecharoHQ/anubis/data"
 	"github.com/TecharoHQ/anubis/internal"
+	"github.com/TecharoHQ/anubis/internal/thoth"
 	libanubis "github.com/TecharoHQ/anubis/lib"
 	botPolicy "github.com/TecharoHQ/anubis/lib/policy"
 	"github.com/TecharoHQ/anubis/lib/policy/config"
 	"github.com/TecharoHQ/anubis/web"
 	"github.com/facebookgo/flagenv"
+	_ "github.com/joho/godotenv/autoload"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )
 
@@ -66,6 +68,9 @@ var (
 	ogCacheConsiderHost      = flag.Bool("og-cache-consider-host", false, "enable or disable the use of the host in the Open Graph tag cache")
 	extractResources         = flag.String("extract-resources", "", "if set, extract the static resources to the specified folder")
 	webmasterEmail           = flag.String("webmaster-email", "", "if set, displays webmaster's email on the reject page for appeals")
+
+	thothURL   = flag.String("thoth-url", "", "if set, URL for Thoth, the IP reputation database for Anubis")
+	thothToken = flag.String("thoth-token", "", "if set, API token for Thoth, the IP reputation database for Anubis")
 )
 
 func keyFromHex(value string) (ed25519.PrivateKey, error) {
@@ -220,7 +225,19 @@ func main() {
 		}
 	}
 
-	policy, err := libanubis.LoadPoliciesOrDefault(*policyFname, *challengeDifficulty)
+	ctx := context.Background()
+
+	if *thothURL != "" && *thothToken != "" {
+		slog.Debug("connecting to Thoth")
+		thothClient, err := thoth.New(ctx, *thothURL, *thothToken)
+		if err != nil {
+			log.Fatalf("can't dial thoth at %s: %v", *thothURL, err)
+		}
+
+		ctx = thoth.With(ctx, thothClient)
+	}
+
+	policy, err := libanubis.LoadPoliciesOrDefault(ctx, *policyFname, *challengeDifficulty)
 	if err != nil {
 		log.Fatalf("can't parse policy file: %v", err)
 	}

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,8 +71,15 @@ func New(ctx context.Context, thothURL, apiToken string) (*Client, error) {
 }
 
 func (c *Client) Close() error {
+	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 {
 	asnMap := map[uint32]struct{}{}

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

lib/anubis_test.go

@@ -14,6 +14,8 @@ import (
 	"github.com/TecharoHQ/anubis"
 	"github.com/TecharoHQ/anubis/data"
 	"github.com/TecharoHQ/anubis/internal"
+	"github.com/TecharoHQ/anubis/internal/thoth"
+	"github.com/TecharoHQ/anubis/internal/thoth/thothmock"
 	"github.com/TecharoHQ/anubis/lib/policy"
 	"github.com/TecharoHQ/anubis/lib/policy/config"
 )
@@ -21,7 +23,11 @@ import (
 func loadPolicies(t *testing.T, fname string) *policy.ParsedConfig {
 	t.Helper()
 
-	anubisPolicy, err := LoadPoliciesOrDefault(t.Context(), fname, anubis.DefaultDifficulty)
+	thothCli := &thoth.Client{}
+	thothCli.WithIPToASNService(thothmock.MockIpToASNService())
+	ctx := thoth.With(t.Context(), thothCli)
+
+	anubisPolicy, err := LoadPoliciesOrDefault(ctx, fname, anubis.DefaultDifficulty)
 	if err != nil {
 		t.Fatal(err)
 	}

lib/policy/config/config.go

@@ -91,7 +91,9 @@ func (b BotConfig) Valid() error {
 	allFieldsEmpty := b.UserAgentRegex == nil &&
 		b.PathRegex == nil &&
 		len(b.RemoteAddr) == 0 &&
-		len(b.HeadersRegex) == 0
+		len(b.HeadersRegex) == 0 &&
+		b.ASNs == nil &&
+		b.GeoIP == nil
 
 	if allFieldsEmpty && b.Expression == nil {
 		errs = append(errs, ErrBotMustHaveUserAgentOrPath)

lib/policy/config/testdata/good/challenge_cloudflare.yaml

@@ -0,0 +1,6 @@
+bots:
+  - name: challenge-cloudflare
+    action: CHALLENGE
+    asns:
+      match:
+        - 13335 # Cloudflare

lib/policy/config/testdata/good/geoip_us.yaml

@@ -0,0 +1,6 @@
+bots:
+  - name: compute-tarrif-us
+    action: CHALLENGE
+    geoip:
+      countries:
+        - US

lib/policy/policy.go

@@ -119,6 +119,15 @@ func ParseConfig(ctx context.Context, fin io.Reader, fname string, defaultDiffic
 			cl = append(cl, tc.ASNCheckerFor(b.ASNs.Match))
 		}
 
+		if b.GeoIP != nil {
+			if !hasThothClient {
+				validationErrs = append(validationErrs, fmt.Errorf("%w: %w", ErrMisconfiguration, ErrNoThothClient))
+				continue
+			}
+
+			cl = append(cl, tc.GeoIPCheckerFor(b.GeoIP.Countries))
+		}
+
 		if b.Challenge == nil {
 			parsedBot.Challenge = &config.ChallengeRules{
 				Difficulty: defaultDifficulty,

lib/policy/policy_test.go

@@ -7,6 +7,8 @@ import (
 
 	"github.com/TecharoHQ/anubis"
 	"github.com/TecharoHQ/anubis/data"
+	"github.com/TecharoHQ/anubis/internal/thoth"
+	"github.com/TecharoHQ/anubis/internal/thoth/thothmock"
 )
 
 func TestDefaultPolicyMustParse(t *testing.T) {
@@ -16,7 +18,11 @@ func TestDefaultPolicyMustParse(t *testing.T) {
 	}
 	defer fin.Close()
 
-	if _, err := ParseConfig(t.Context(), fin, "botPolicies.json", anubis.DefaultDifficulty); err != nil {
+	thothCli := &thoth.Client{}
+	thothCli.WithIPToASNService(thothmock.MockIpToASNService())
+	ctx := thoth.With(t.Context(), thothCli)
+
+	if _, err := ParseConfig(ctx, fin, "botPolicies.json", anubis.DefaultDifficulty); err != nil {
 		t.Fatalf("can't parse config: %v", err)
 	}
 }
@@ -36,7 +42,11 @@ func TestGoodConfigs(t *testing.T) {
 			}
 			defer fin.Close()
 
-			if _, err := ParseConfig(t.Context(), fin, fin.Name(), anubis.DefaultDifficulty); err != nil {
+			thothCli := &thoth.Client{}
+			thothCli.WithIPToASNService(thothmock.MockIpToASNService())
+			ctx := thoth.With(t.Context(), thothCli)
+
+			if _, err := ParseConfig(ctx, fin, fin.Name(), anubis.DefaultDifficulty); err != nil {
 				t.Fatal(err)
 			}
 		})
@@ -58,7 +68,11 @@ func TestBadConfigs(t *testing.T) {
 			}
 			defer fin.Close()
 
-			if _, err := ParseConfig(t.Context(), fin, fin.Name(), anubis.DefaultDifficulty); err == nil {
+			thothCli := &thoth.Client{}
+			thothCli.WithIPToASNService(thothmock.MockIpToASNService())
+			ctx := thoth.With(t.Context(), thothCli)
+
+			if _, err := ParseConfig(ctx, fin, fin.Name(), anubis.DefaultDifficulty); err == nil {
 				t.Fatal(err)
 			} else {
 				t.Log(err)