feat(internal): add Thoth client and simple ASN checker

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

internal/thoth/asnchecker.go

@@ -0,0 +1,35 @@
+package thoth
+
+import (
+	"context"
+	"net/http"
+	"time"
+
+	iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
+)
+
+type ASNChecker struct {
+	iptoasn iptoasnv1.IpToASNServiceClient
+	asns    map[int]struct{}
+	hash    string
+}
+
+func (asnc *ASNChecker) Check(r *http.Request) (bool, error) {
+	ctx, cancel := context.WithTimeout(r.Context(), 50*time.Millisecond)
+	defer cancel()
+
+	ipInfo, err := asnc.iptoasn.Lookup(ctx, &iptoasnv1.LookupRequest{
+		IpAddress: r.Header.Get("X-Real-Ip"),
+	})
+	if err != nil {
+		return false, err
+	}
+
+	_, ok := asnc.asns[int(ipInfo.GetAsNumber())]
+
+	return ok, nil
+}
+
+func (asnc *ASNChecker) Hash() string {
+	return asnc.hash
+}

internal/thoth/asnchecker_test.go

@@ -0,0 +1,59 @@
+package thoth
+
+import (
+	"fmt"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestASNChecker(t *testing.T) {
+	cli := loadSecrets(t)
+
+	asnc := &ASNChecker{
+		iptoasn: cli.iptoasn,
+		asns: map[int]struct{}{
+			13335: struct{}{},
+		},
+		hash: "foobar",
+	}
+
+	for _, cs := range []struct {
+		ipAddress string
+		wantMatch bool
+		wantError bool
+	}{
+		{
+			ipAddress: "1.1.1.1",
+			wantMatch: true,
+			wantError: false,
+		},
+		{
+			ipAddress: "8.8.8.8",
+			wantMatch: false,
+			wantError: false,
+		},
+		{
+			ipAddress: "taco",
+			wantMatch: false,
+			wantError: true,
+		},
+	} {
+		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")
+			}
+		})
+	}
+}

internal/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...)
+	}
+}

internal/thoth/thoth.go

@@ -0,0 +1,69 @@
+package thoth
+
+import (
+	"context"
+	"crypto/tls"
+	"fmt"
+
+	iptoasnv1 "github.com/TecharoHQ/thoth-proto/gen/techaro/thoth/iptoasn/v1"
+	grpcprom "github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus"
+	"github.com/prometheus/client_golang/prometheus"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+	healthv1 "google.golang.org/grpc/health/grpc_health_v1"
+)
+
+type Client struct {
+	thothURL string
+
+	conn    *grpc.ClientConn
+	health  healthv1.HealthClient
+	iptoasn iptoasnv1.IpToASNServiceClient
+}
+
+func New(ctx context.Context, thothURL, apiToken string) (*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)
+
+	conn, err := grpc.DialContext(
+		ctx,
+		thothURL,
+		grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{})),
+		grpc.WithChainUnaryInterceptor(
+			clMetrics.UnaryClientInterceptor(),
+			authUnaryClientInterceptor(apiToken),
+		),
+		grpc.WithChainStreamInterceptor(
+			clMetrics.StreamClientInterceptor(),
+			authStreamClientInterceptor(apiToken),
+		),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("can't dial thoth at %s: %w", thothURL, err)
+	}
+
+	hc := healthv1.NewHealthClient(conn)
+
+	resp, err := hc.Check(ctx, &healthv1.HealthCheckRequest{})
+	if err != nil {
+		return nil, fmt.Errorf("can't verify thoth health at %s: %w", thothURL, err)
+	}
+
+	if resp.Status != healthv1.HealthCheckResponse_SERVING {
+		return nil, fmt.Errorf("thoth is not healthy, wanted %s but got %s", healthv1.HealthCheckResponse_SERVING, resp.Status)
+	}
+
+	return &Client{
+		conn:    conn,
+		health:  hc,
+		iptoasn: iptoasnv1.NewIpToASNServiceClient(conn),
+	}, nil
+}
+
+func (c *Client) Close() error {
+	return c.conn.Close()
+}

internal/thoth/thoth_test.go

@@ -0,0 +1,29 @@
+package thoth
+
+import (
+	"os"
+	"testing"
+
+	"github.com/joho/godotenv"
+)
+
+func loadSecrets(t *testing.T) *Client {
+	if err := godotenv.Load(); err != nil {
+		t.Skip(".env not defined, can't load thoth secrets")
+	}
+
+	cli, err := New(t.Context(), os.Getenv("THOTH_URL"), os.Getenv("THOTH_API_KEY"))
+	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)
+	}
+}