chore: expose thoth in lib (#911)

Imports a patch previously exclusive to Botstopper.

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

lib/anubis_test.go

@@ -17,9 +17,9 @@ import (
 	"github.com/TecharoHQ/anubis"
 	"github.com/TecharoHQ/anubis/data"
 	"github.com/TecharoHQ/anubis/internal"
-	"github.com/TecharoHQ/anubis/internal/thoth/thothmock"
 	"github.com/TecharoHQ/anubis/lib/policy"
 	"github.com/TecharoHQ/anubis/lib/policy/config"
+	"github.com/TecharoHQ/anubis/lib/thoth/thothmock"
 )
 
 func init() {

lib/config_test.go

@@ -7,8 +7,8 @@ import (
 	"testing"
 
 	"github.com/TecharoHQ/anubis"
-	"github.com/TecharoHQ/anubis/internal/thoth/thothmock"
 	"github.com/TecharoHQ/anubis/lib/policy"
+	"github.com/TecharoHQ/anubis/lib/thoth/thothmock"
 )
 
 func TestInvalidChallengeMethod(t *testing.T) {

lib/policy/policy.go

@@ -8,10 +8,10 @@ import (
 	"log/slog"
 	"sync/atomic"
 
-	"github.com/TecharoHQ/anubis/internal/thoth"
 	"github.com/TecharoHQ/anubis/lib/policy/checker"
 	"github.com/TecharoHQ/anubis/lib/policy/config"
 	"github.com/TecharoHQ/anubis/lib/store"
+	"github.com/TecharoHQ/anubis/lib/thoth"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 

lib/policy/policy_test.go

@@ -7,7 +7,7 @@ import (
 
 	"github.com/TecharoHQ/anubis"
 	"github.com/TecharoHQ/anubis/data"
-	"github.com/TecharoHQ/anubis/internal/thoth/thothmock"
+	"github.com/TecharoHQ/anubis/lib/thoth/thothmock"
 )
 
 func TestDefaultPolicyMustParse(t *testing.T) {

lib/thoth/asnchecker.go

@@ -0,0 +1,69 @@
+package thoth
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/TecharoHQ/anubis/internal"
+	"github.com/TecharoHQ/anubis/lib/policy/checker"
+	iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
+)
+
+func (c *Client) ASNCheckerFor(asns []uint32) checker.Impl {
+	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.FastHash(sb.String()),
+	}
+}
+
+type ASNChecker struct {
+	iptoasn iptoasnv1.IpToASNServiceClient
+	asns    map[uint32]struct{}
+	hash    string
+}
+
+func (asnc *ASNChecker) Check(r *http.Request) (bool, error) {
+	ctx, cancel := context.WithTimeout(r.Context(), 500*time.Millisecond)
+	defer cancel()
+
+	ipInfo, err := asnc.iptoasn.Lookup(ctx, &iptoasnv1.LookupRequest{
+		IpAddress: r.Header.Get("X-Real-Ip"),
+	})
+	if err != nil {
+		switch {
+		case errors.Is(err, context.DeadlineExceeded):
+			slog.Debug("error contacting thoth", "err", err, "actionable", false)
+			return false, nil
+		default:
+			slog.Error("error contacting thoth, please contact support", "err", err, "actionable", true)
+			return false, nil
+		}
+	}
+
+	// If IP is not publicly announced, return false
+	if !ipInfo.GetAnnounced() {
+		return false, nil
+	}
+
+	_, ok := asnc.asns[uint32(ipInfo.GetAsNumber())]
+
+	return ok, nil
+}
+
+func (asnc *ASNChecker) Hash() string {
+	return asnc.hash
+}

lib/thoth/asnchecker_test.go

@@ -0,0 +1,81 @@
+package thoth_test
+
+import (
+	"fmt"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/TecharoHQ/anubis/lib/policy/checker"
+	"github.com/TecharoHQ/anubis/lib/thoth"
+	iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
+)
+
+var _ checker.Impl = &thoth.ASNChecker{}
+
+func TestASNChecker(t *testing.T) {
+	cli := loadSecrets(t)
+
+	asnc := cli.ASNCheckerFor([]uint32{13335})
+
+	for _, cs := range []struct {
+		ipAddress string
+		wantMatch bool
+		wantError bool
+	}{
+		{
+			ipAddress: "1.1.1.1",
+			wantMatch: true,
+			wantError: false,
+		},
+		{
+			ipAddress: "2.2.2.2",
+			wantMatch: false,
+			wantError: false,
+		},
+		{
+			ipAddress: "taco",
+			wantMatch: false,
+			wantError: false,
+		},
+		{
+			ipAddress: "127.0.0.1",
+			wantMatch: false,
+			wantError: false,
+		},
+	} {
+		t.Run(fmt.Sprintf("%v", cs), func(t *testing.T) {
+			req := httptest.NewRequest("GET", "/", nil)
+			req.Header.Set("X-Real-Ip", cs.ipAddress)
+
+			match, err := asnc.Check(req)
+
+			if match != cs.wantMatch {
+				t.Errorf("Wanted match: %v, got: %v", cs.wantMatch, match)
+			}
+
+			switch {
+			case err != nil && !cs.wantError:
+				t.Errorf("Did not want error but got: %v", err)
+			case err == nil && cs.wantError:
+				t.Error("Wanted error but got none")
+			}
+		})
+	}
+}
+
+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)
+		}
+	}
+}

lib/thoth/auth.go

@@ -0,0 +1,39 @@
+package thoth
+
+import (
+	"context"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+func authUnaryClientInterceptor(token string) grpc.UnaryClientInterceptor {
+	return func(
+		ctx context.Context,
+		method string,
+		req interface{},
+		reply interface{},
+		cc *grpc.ClientConn,
+		invoker grpc.UnaryInvoker,
+		opts ...grpc.CallOption,
+	) error {
+		md := metadata.Pairs("authorization", "Bearer "+token)
+		ctx = metadata.NewOutgoingContext(ctx, md)
+		return invoker(ctx, method, req, reply, cc, opts...)
+	}
+}
+
+func authStreamClientInterceptor(token string) grpc.StreamClientInterceptor {
+	return func(
+		ctx context.Context,
+		desc *grpc.StreamDesc,
+		cc *grpc.ClientConn,
+		method string,
+		streamer grpc.Streamer,
+		opts ...grpc.CallOption,
+	) (grpc.ClientStream, error) {
+		md := metadata.Pairs("authorization", "Bearer "+token)
+		ctx = metadata.NewOutgoingContext(ctx, md)
+		return streamer(ctx, desc, cc, method, opts...)
+	}
+}

lib/thoth/cachediptoasn.go

@@ -0,0 +1,84 @@
+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 {
+	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) {
+	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
+}

lib/thoth/context.go

@@ -0,0 +1,14 @@
+package thoth
+
+import "context"
+
+type ctxKey struct{}
+
+func With(ctx context.Context, cli *Client) context.Context {
+	return context.WithValue(ctx, ctxKey{}, cli)
+}
+
+func FromContext(ctx context.Context) (*Client, bool) {
+	cli, ok := ctx.Value(ctxKey{}).(*Client)
+	return cli, ok
+}

lib/thoth/geoipchecker.go

@@ -0,0 +1,68 @@
+package thoth
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/TecharoHQ/anubis/lib/policy/checker"
+	iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
+)
+
+func (c *Client) GeoIPCheckerFor(countries []string) checker.Impl {
+	countryMap := map[string]struct{}{}
+	var sb strings.Builder
+	fmt.Fprintln(&sb, "GeoIPChecker")
+	for _, cc := range countries {
+		countryMap[cc] = struct{}{}
+		fmt.Fprintln(&sb, cc)
+	}
+
+	return &GeoIPChecker{
+		IPToASN:   c.IPToASN,
+		Countries: countryMap,
+		hash:      sb.String(),
+	}
+}
+
+type GeoIPChecker struct {
+	IPToASN   iptoasnv1.IpToASNServiceClient
+	Countries map[string]struct{}
+	hash      string
+}
+
+func (gipc *GeoIPChecker) Check(r *http.Request) (bool, error) {
+	ctx, cancel := context.WithTimeout(r.Context(), 500*time.Millisecond)
+	defer cancel()
+
+	ipInfo, err := gipc.IPToASN.Lookup(ctx, &iptoasnv1.LookupRequest{
+		IpAddress: r.Header.Get("X-Real-Ip"),
+	})
+	if err != nil {
+		switch {
+		case errors.Is(err, context.DeadlineExceeded):
+			slog.Debug("error contacting thoth", "err", err, "actionable", false)
+			return false, nil
+		default:
+			slog.Error("error contacting thoth, please contact support", "err", err, "actionable", true)
+			return false, nil
+		}
+	}
+
+	// If IP is not publicly announced, return false
+	if !ipInfo.GetAnnounced() {
+		return false, nil
+	}
+
+	_, ok := gipc.Countries[strings.ToLower(ipInfo.GetCountryCode())]
+
+	return ok, nil
+}
+
+func (gipc *GeoIPChecker) Hash() string {
+	return gipc.hash
+}

lib/thoth/geoipchecker_test.go

@@ -0,0 +1,63 @@
+package thoth_test
+
+import (
+	"fmt"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/TecharoHQ/anubis/lib/policy/checker"
+	"github.com/TecharoHQ/anubis/lib/thoth"
+)
+
+var _ checker.Impl = &thoth.GeoIPChecker{}
+
+func TestGeoIPChecker(t *testing.T) {
+	cli := loadSecrets(t)
+
+	asnc := cli.GeoIPCheckerFor([]string{"us"})
+
+	for _, cs := range []struct {
+		ipAddress string
+		wantMatch bool
+		wantError bool
+	}{
+		{
+			ipAddress: "1.1.1.1",
+			wantMatch: true,
+			wantError: false,
+		},
+		{
+			ipAddress: "2.2.2.2",
+			wantMatch: false,
+			wantError: false,
+		},
+		{
+			ipAddress: "taco",
+			wantMatch: false,
+			wantError: false,
+		},
+		{
+			ipAddress: "127.0.0.1",
+			wantMatch: false,
+			wantError: false,
+		},
+	} {
+		t.Run(fmt.Sprintf("%v", cs), func(t *testing.T) {
+			req := httptest.NewRequest("GET", "/", nil)
+			req.Header.Set("X-Real-Ip", cs.ipAddress)
+
+			match, err := asnc.Check(req)
+
+			if match != cs.wantMatch {
+				t.Errorf("Wanted match: %v, got: %v", cs.wantMatch, match)
+			}
+
+			switch {
+			case err != nil && !cs.wantError:
+				t.Errorf("Did not want error but got: %v", err)
+			case err == nil && cs.wantError:
+				t.Error("Wanted error but got none")
+			}
+		})
+	}
+}

lib/thoth/thoth.go

@@ -0,0 +1,79 @@
+package thoth
+
+import (
+	"context"
+	"crypto/tls"
+	"fmt"
+	"time"
+
+	"github.com/TecharoHQ/anubis"
+	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"
+	"github.com/prometheus/client_golang/prometheus"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
+	healthv1 "google.golang.org/grpc/health/grpc_health_v1"
+)
+
+type Client struct {
+	conn    *grpc.ClientConn
+	health  healthv1.HealthClient
+	IPToASN iptoasnv1.IpToASNServiceClient
+}
+
+func New(ctx context.Context, thothURL, apiToken string, plaintext bool) (*Client, error) {
+	clMetrics := grpcprom.NewClientMetrics(
+		grpcprom.WithClientHandlingTimeHistogram(
+			grpcprom.WithHistogramBuckets([]float64{0.001, 0.01, 0.1, 0.3, 0.6, 1, 3, 6, 9, 20, 30, 60, 90, 120}),
+		),
+	)
+	prometheus.DefaultRegisterer.Register(clMetrics)
+
+	do := []grpc.DialOption{
+		grpc.WithChainUnaryInterceptor(
+			timeout.UnaryClientInterceptor(500*time.Millisecond),
+			clMetrics.UnaryClientInterceptor(),
+			authUnaryClientInterceptor(apiToken),
+		),
+		grpc.WithChainStreamInterceptor(
+			clMetrics.StreamClientInterceptor(),
+			authStreamClientInterceptor(apiToken),
+		),
+		grpc.WithUserAgent(fmt.Sprint("Techaro/anubis:", anubis.Version)),
+	}
+
+	if plaintext {
+		do = append(do, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	} else {
+		do = append(do, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{})))
+	}
+
+	conn, err := grpc.NewClient(
+		thothURL,
+		do...,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("can't dial thoth at %s: %w", thothURL, err)
+	}
+
+	hc := healthv1.NewHealthClient(conn)
+
+	return &Client{
+		conn:    conn,
+		health:  hc,
+		IPToASN: NewIpToASNWithCache(iptoasnv1.NewIpToASNServiceClient(conn)),
+	}, nil
+}
+
+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
+}

lib/thoth/thoth_test.go

@@ -0,0 +1,36 @@
+package thoth_test
+
+import (
+	"os"
+	"testing"
+
+	"github.com/TecharoHQ/anubis/lib/thoth"
+	"github.com/TecharoHQ/anubis/lib/thoth/thothmock"
+	"github.com/joho/godotenv"
+)
+
+func loadSecrets(t testing.TB) *thoth.Client {
+	t.Helper()
+
+	if err := godotenv.Load(); err != nil {
+		t.Log("using mock thoth")
+		result := &thoth.Client{}
+		result.WithIPToASNService(thothmock.MockIpToASNService())
+		return result
+	}
+
+	cli, err := thoth.New(t.Context(), os.Getenv("THOTH_URL"), os.Getenv("THOTH_API_KEY"), false)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	return cli
+}
+
+func TestNew(t *testing.T) {
+	cli := loadSecrets(t)
+
+	if err := cli.Close(); err != nil {
+		t.Fatal(err)
+	}
+}

lib/thoth/thothmock/iptoasn.go

@@ -0,0 +1,59 @@
+package thothmock
+
+import (
+	"context"
+	"net/netip"
+
+	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{
+		"127.0.0.1": {Announced: false},
+		"::1":       {Announced: false},
+		"10.10.10.10": {
+			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",
+		},
+		"1.1.1.1": {
+			Announced:   true,
+			AsNumber:    13335,
+			Cidr:        []string{"1.1.1.0/24"},
+			CountryCode: "US",
+			Description: "Cloudflare",
+		},
+	}
+
+	return &IpToASNService{Responses: responses}
+}
+
+type IpToASNService struct {
+	iptoasnv1.UnimplementedIpToASNServiceServer
+	Responses map[string]*iptoasnv1.LookupResponse
+}
+
+func (ip2asn *IpToASNService) Lookup(ctx context.Context, lr *iptoasnv1.LookupRequest, opts ...grpc.CallOption) (*iptoasnv1.LookupResponse, error) {
+	if _, err := netip.ParseAddr(lr.GetIpAddress()); err != nil {
+		return nil, err
+	}
+
+	resp, ok := ip2asn.Responses[lr.GetIpAddress()]
+	if !ok {
+		return nil, status.Error(codes.NotFound, "IP address not found in mock")
+	}
+
+	return resp, nil
+}

lib/thoth/thothmock/withthothmock.go

@@ -0,0 +1,17 @@
+package thothmock
+
+import (
+	"context"
+	"testing"
+
+	"github.com/TecharoHQ/anubis/lib/thoth"
+)
+
+func WithMockThoth(t *testing.T) context.Context {
+	t.Helper()
+
+	thothCli := &thoth.Client{}
+	thothCli.WithIPToASNService(MockIpToASNService())
+	ctx := thoth.With(t.Context(), thothCli)
+	return ctx
+}