feat(sdk): Include auth token in grpc

sdk/token_adding_interceptor.go

@@ -0,0 +1,115 @@
+package sdk
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/sha256"
+	"encoding/base64"
+	"fmt"
+	"log/slog"
+	"time"
+
+	"github.com/lestrrat-go/jwx/v2/jwk"
+	"github.com/lestrrat-go/jwx/v2/jws"
+	"github.com/lestrrat-go/jwx/v2/jwt"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+const (
+	JTILength            = 14
+	JWTExpirationMinutes = 10
+)
+
+func newOutgoingInterceptor(t AccessTokenSource) tokenAddingInterceptor {
+	return tokenAddingInterceptor{tokenSource: t}
+}
+
+type tokenAddingInterceptor struct {
+	tokenSource AccessTokenSource
+}
+
+func (i tokenAddingInterceptor) addCredentials(ctx context.Context,
+	method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+	newMetadata := make([]string, 0)
+	accessToken, err := i.tokenSource.AccessToken()
+	if err == nil {
+		newMetadata = append(newMetadata, "Authorization", fmt.Sprintf("DPoP %s", accessToken))
+	} else {
+		slog.Error("error getting access token: %w. Request will be unauthenticated", err)
+	}
+
+	dpopTok, err := i.getDPOPToken(method, string(accessToken))
+	if err == nil {
+		newMetadata = append(newMetadata, "DPoP", dpopTok)
+	} else {
+		slog.Error("error adding dpop token to outgoing request. Request will not have DPoP token", err)
+	}
+
+	newCtx := metadata.AppendToOutgoingContext(ctx, newMetadata...)
+
+	err = invoker(newCtx, method, req, reply, cc, opts...)
+
+	// this is the error from the RPC service. we can determine when the current token is no longer valid
+	// by inspecting this error
+	return err
+}
+
+func (i tokenAddingInterceptor) getDPOPToken(method, accessToken string) (string, error) {
+	tok, err := i.tokenSource.MakeToken(func(key jwk.Key) ([]byte, error) {
+		jtiBytes := make([]byte, JTILength)
+		_, err := rand.Read(jtiBytes)
+		if err != nil {
+			return nil, fmt.Errorf("error creating jti for dpop jwt: %w", err)
+		}
+
+		publicKey, err := key.PublicKey()
+		if err != nil {
+			return nil, fmt.Errorf("error getting public key from DPOP key: %w", err)
+		}
+
+		headers := jws.NewHeaders()
+		err = headers.Set(jws.JWKKey, publicKey)
+		if err != nil {
+			return nil, fmt.Errorf("error setting the key on the DPOP token: %w", err)
+		}
+		err = headers.Set(jws.TypeKey, "dpop+jwt")
+		if err != nil {
+			return nil, fmt.Errorf("error setting the type on the DPOP token: %w", err)
+		}
+		err = headers.Set(jws.AlgorithmKey, key.Algorithm())
+		if err != nil {
+			return nil, fmt.Errorf("error setting the algorithm on the DPOP token: %w", err)
+		}
+
+		h := sha256.New()
+		h.Write([]byte(accessToken))
+		ath := h.Sum(nil)
+
+		dpopTok, err := jwt.NewBuilder().
+			Claim("htu", method).
+			Claim("htm", "POST").
+			Claim("ath", base64.StdEncoding.EncodeToString(ath)).
+			Claim("jti", base64.StdEncoding.EncodeToString(jtiBytes)).
+			IssuedAt(time.Now()).
+			Expiration(time.Now().Add(time.Minute * JWTExpirationMinutes)).
+			Build()
+
+		if err != nil {
+			return nil, fmt.Errorf("error creating dpop jwt: %w", err)
+		}
+
+		signedToken, err := jwt.Sign(dpopTok, jwt.WithKey(key.Algorithm(), key, jws.WithProtectedHeaders(headers)))
+		if err != nil {
+			return nil, fmt.Errorf("error signing dpop jwt: %w", err)
+		}
+
+		return signedToken, nil
+	})
+
+	if err != nil {
+		return "", fmt.Errorf("error creating DPOP token in interceptor: %w", err)
+	}
+
+	return string(tok), nil
+}

sdk/token_adding_interceptor_test.go

@@ -0,0 +1,187 @@
+package sdk
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/base64"
+	"errors"
+	"net"
+	"slices"
+	"testing"
+
+	gocrypto "crypto"
+
+	"github.com/lestrrat-go/jwx/v2/jwa"
+	"github.com/lestrrat-go/jwx/v2/jwk"
+	"github.com/lestrrat-go/jwx/v2/jws"
+	"github.com/lestrrat-go/jwx/v2/jwt"
+	kas "github.com/opentdf/backend-go/pkg/access"
+	"github.com/opentdf/platform/sdk/internal/crypto"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/test/bufconn"
+	"google.golang.org/protobuf/types/known/wrapperspb"
+)
+
+func TestAddingTokensToOutgoingRequest(t *testing.T) {
+	_, key, _, _ := getNewDPoPKey()
+	ts := FakeTokenSource{
+		key:         key,
+		accessToken: "thisisafakeaccesstoken",
+	}
+	server := FakeAccessServiceServer{}
+	oo := newOutgoingInterceptor(&ts)
+
+	client, stop := runServer(context.Background(), &server, oo)
+	defer stop()
+
+	_, err := client.Info(context.Background(), &kas.InfoRequest{})
+	if err != nil {
+		t.Fatalf("error making call: %v", err)
+	}
+
+	if len(server.accessToken) != 1 || server.accessToken[0] != "DPoP thisisafakeaccesstoken" {
+		t.Fatalf("got incorrect access token: %v", server.accessToken)
+	}
+
+	if len(server.dpopToken) != 1 {
+		t.Fatalf("Got incorrect dpop token headers: %v", server.dpopToken)
+	}
+
+	dpopToken := server.dpopToken[0]
+
+	alg, ok := key.Algorithm().(jwa.SignatureAlgorithm)
+	if !ok {
+		t.Fatalf("got a bad signing algorithm")
+	}
+
+	_, err = jws.Verify([]byte(dpopToken), jws.WithKey(alg, key))
+	if err != nil {
+		t.Fatalf("error verifying signature: %v", err)
+	}
+
+	parsedSignature, _ := jws.Parse([]byte(dpopToken))
+
+	if len(parsedSignature.Signatures()) == 0 {
+		t.Fatalf("didn't get signature from jwt")
+	}
+
+	sig := parsedSignature.Signatures()[0]
+	tokenKey, ok := sig.ProtectedHeaders().Get("jwk")
+	if !ok {
+		t.Fatalf("didn't get error getting key from token")
+	}
+
+	tp, _ := tokenKey.(jwk.Key).Thumbprint(gocrypto.SHA256)
+	ktp, _ := key.Thumbprint(gocrypto.SHA256)
+	if !slices.Equal(tp, ktp) {
+		t.Fatalf("got the wrong key from the token")
+	}
+
+	parsedToken, _ := jwt.Parse([]byte(dpopToken), jwt.WithVerify(false))
+
+	if method, _ := parsedToken.Get("htm"); method.(string) != "POST" {
+		t.Fatalf("we got a bad method: %v", method)
+	}
+
+	if path, _ := parsedToken.Get("htu"); path.(string) != "/access.AccessService/Info" {
+		t.Fatalf("we got a bad method: %v", path)
+	}
+
+	h := sha256.New()
+	h.Write([]byte("thisisafakeaccesstoken"))
+	expectedHash := base64.URLEncoding.EncodeToString(h.Sum(nil))
+
+	if ath, _ := parsedToken.Get("ath"); ath.(string) != expectedHash {
+		t.Fatalf("got invalid ath claim in token: %v", ath)
+	}
+}
+
+func Test_InvalidCredentials_StillSendMessage(t *testing.T) {
+	ts := FakeTokenSource{key: nil}
+	server := FakeAccessServiceServer{}
+	oo := newOutgoingInterceptor(&ts)
+
+	client, stop := runServer(context.Background(), &server, oo)
+	defer stop()
+
+	_, err := client.Info(context.Background(), &kas.InfoRequest{})
+
+	if err != nil {
+		t.Fatalf("got an error when sending the message")
+	}
+}
+
+type FakeAccessServiceServer struct {
+	accessToken []string
+	dpopToken   []string
+	kas.UnimplementedAccessServiceServer
+}
+
+func (f *FakeAccessServiceServer) Info(ctx context.Context, _ *kas.InfoRequest) (*kas.InfoResponse, error) {
+	if md, ok := metadata.FromIncomingContext(ctx); ok {
+		f.accessToken = md.Get("authorization")
+		f.dpopToken = md.Get("dpop")
+	}
+
+	return &kas.InfoResponse{}, nil
+}
+func (f *FakeAccessServiceServer) PublicKey(context.Context, *kas.PublicKeyRequest) (*kas.PublicKeyResponse, error) {
+	return &kas.PublicKeyResponse{}, status.Error(codes.Unauthenticated, "no public key for you")
+}
+func (f *FakeAccessServiceServer) LegacyPublicKey(context.Context, *kas.LegacyPublicKeyRequest) (*wrapperspb.StringValue, error) {
+	return &wrapperspb.StringValue{}, nil
+}
+func (f *FakeAccessServiceServer) Rewrap(context.Context, *kas.RewrapRequest) (*kas.RewrapResponse, error) {
+	return &kas.RewrapResponse{}, nil
+}
+
+type FakeTokenSource struct {
+	key         jwk.Key
+	accessToken string
+}
+
+func (fts *FakeTokenSource) AccessToken() (AccessToken, error) {
+	return AccessToken(fts.accessToken), nil
+}
+func (*FakeTokenSource) AsymDecryption() crypto.AsymDecryption {
+	return crypto.AsymDecryption{}
+}
+func (fts *FakeTokenSource) MakeToken(f func(jwk.Key) ([]byte, error)) ([]byte, error) {
+	if fts.key == nil {
+		return nil, errors.New("no such key")
+	}
+	return f(fts.key)
+}
+func (*FakeTokenSource) DPOPPublicKeyPEM() string {
+	return ""
+}
+func (*FakeTokenSource) RefreshAccessToken() error {
+	return nil
+}
+
+func runServer(ctx context.Context, //nolint:ireturn // this is pretty concrete
+	f *FakeAccessServiceServer, oo tokenAddingInterceptor) (kas.AccessServiceClient, func()) {
+	buffer := 1024 * 1024
+	listener := bufconn.Listen(buffer)
+
+	s := grpc.NewServer()
+	kas.RegisterAccessServiceServer(s, f)
+	go func() {
+		if err := s.Serve(listener); err != nil {
+			panic(err)
+		}
+	}()
+
+	conn, _ := grpc.DialContext(ctx, "", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
+		return listener.Dial()
+	}), grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithBlock(), grpc.WithUnaryInterceptor(oo.addCredentials))
+
+	client := kas.NewAccessServiceClient(conn)
+
+	return client, s.Stop
+}