From 714363fb072e01e968fc051b32f14f2b9f9cb478 Mon Sep 17 00:00:00 2001
From: Caleb Doxsey <cdoxsey@pomerium.com>
Date: Thu, 30 Jul 2020 09:40:55 -0600
Subject: [PATCH] grpcutil: add functions for JWTs in gRPC metadata (#1165)

---
 pkg/grpcutil/grpcutil.go      | 23 +++++++++++++++++++++++
 pkg/grpcutil/grpcutil_test.go | 22 ++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/pkg/grpcutil/grpcutil.go b/pkg/grpcutil/grpcutil.go
index 9dd664a67..baeead3b6 100644
--- a/pkg/grpcutil/grpcutil.go
+++ b/pkg/grpcutil/grpcutil.go
@@ -29,3 +29,26 @@ func SessionIDFromGRPCRequest(ctx context.Context) (sessionID string, ok bool) {
 
 	return sessionIDs[0], true
 }
+
+// JWTMetadataKey is the key in the metadata.
+const JWTMetadataKey = "jwt"
+
+// WithOutgoingJWT appends a metadata header for the JWT to a context.
+func WithOutgoingJWT(ctx context.Context, rawjwt string) context.Context {
+	return metadata.AppendToOutgoingContext(ctx, JWTMetadataKey, rawjwt)
+}
+
+// JWTFromGRPCRequest returns the JWT from the gRPC request.
+func JWTFromGRPCRequest(ctx context.Context) (rawjwt string, ok bool) {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return "", false
+	}
+
+	rawjwts := md.Get(JWTMetadataKey)
+	if len(rawjwts) == 0 {
+		return "", false
+	}
+
+	return rawjwts[0], true
+}
diff --git a/pkg/grpcutil/grpcutil_test.go b/pkg/grpcutil/grpcutil_test.go
index ce076cfd0..f9f181c6a 100644
--- a/pkg/grpcutil/grpcutil_test.go
+++ b/pkg/grpcutil/grpcutil_test.go
@@ -27,3 +27,25 @@ func TestSessionIDFromGRPCRequest(t *testing.T) {
 	assert.True(t, ok)
 	assert.Equal(t, "EXAMPLE", sessionID)
 }
+
+func TestWithOutgoingJWT(t *testing.T) {
+	rawjwt := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
+	ctx := context.Background()
+	ctx = WithOutgoingJWT(ctx, rawjwt)
+	md, ok := metadata.FromOutgoingContext(ctx)
+	if !assert.True(t, ok) {
+		return
+	}
+	assert.Equal(t, []string{rawjwt}, md.Get("jwt"))
+}
+
+func TestJWTFromGRPCRequest(t *testing.T) {
+	rawjwt := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
+	ctx := context.Background()
+	ctx = metadata.NewIncomingContext(ctx, metadata.MD{
+		"jwt": {rawjwt},
+	})
+	found, ok := JWTFromGRPCRequest(ctx)
+	assert.True(t, ok)
+	assert.Equal(t, rawjwt, found)
+}