From d277757bdb6743f4b302e8e8b80b3da8d6f03e53 Mon Sep 17 00:00:00 2001
From: Ben Arena <barena@figure.com>
Date: Thu, 30 Jun 2022 09:38:03 -0700
Subject: [PATCH] optionally inject the wallet public key as a header

---
 README.md                   |  3 ++-
 cmd/jwt-wallet/main_test.go |  2 ++
 jwt-wallet.go               | 16 +++++++++++-----
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index ca9d7a5..d365be4 100644
--- a/README.md
+++ b/README.md
@@ -48,7 +48,8 @@ Configuration options:
 * `authHeader` - The name of the request header containing the JWT. Defaults to "Authorization". 
   - Requires the Bearer token format e.g {authHeader} Bearer {jwt}
 * `accessHeader` - The name of the header to inject with the wallet access JSON. Defaults to "x-wallet-access". Only injected when rbac configuration is set
-* `senderHeader` - The name of the header to inject with the addr claim of the user signed JWT. If not set then will not inject
+* `senderHeader` - The name of the header to inject with the `addr` claim of the user signed JWT. If not set then will not inject
+* `publicKeyHeader` - The name of the header to inject with the public key (`sub` claim) of the wallet. If not set then will not inject.
 
 *= Required to get delegated access rights
 
diff --git a/cmd/jwt-wallet/main_test.go b/cmd/jwt-wallet/main_test.go
index 7193708..2eb3e52 100644
--- a/cmd/jwt-wallet/main_test.go
+++ b/cmd/jwt-wallet/main_test.go
@@ -148,10 +148,12 @@ func TestValidJwt(t *testing.T) {
 	})
 	assert.NoError(t, err)
 
+	config.PublicKeyHeader = "x-wallet-public-key"
 	env.DoHttp(config)
 
 	assert.Equal(t, 200, env.ClientRes.Status)
 	assert.NotEmpty(t, env.ServiceReq.Headers.Get("x-wallet-access"))
+	assert.Equal(t, env.ServiceReq.Headers.Get("x-wallet-public-key"), claims.Subject)
 	assert.Empty(t, env.ServiceReq.Headers.Get("x-sender"))
 	assert.Equal(t, subjectJSONString, env.ServiceReq.Headers.Get("x-wallet-access"))
 }
diff --git a/jwt-wallet.go b/jwt-wallet.go
index 025ba31..d058015 100644
--- a/jwt-wallet.go
+++ b/jwt-wallet.go
@@ -18,11 +18,12 @@ import (
 )
 
 type Config struct {
-	RBAC         string `json:"rbac"`
-	APIKey       string `json:"apikey"`
-	AuthHeader   string `json:"authHeader"`
-	AccessHeader string `json:"accessHeader"`
-	SenderHeader string `json:"senderHeader"`
+	RBAC            string `json:"rbac"`
+	APIKey          string `json:"apikey"`
+	AuthHeader      string `json:"authHeader"`
+	AccessHeader    string `json:"accessHeader"`
+	SenderHeader    string `json:"senderHeader"`
+	PublicKeyHeader string `json:"publicKeyHeader"`
 }
 
 func New() interface{} {
@@ -83,6 +84,7 @@ func (conf Config) Access(kong *pdk.PDK) {
 		kong.Response.Exit(500, "something went wrong", x)
 		return
 	}
+
 	if conf.AccessHeader == "" {
 		conf.AccessHeader = "x-wallet-access"
 	}
@@ -95,6 +97,10 @@ func (conf Config) Access(kong *pdk.PDK) {
 		kong.ServiceRequest.SetHeader(conf.SenderHeader, sender)
 	}
 
+	if conf.PublicKeyHeader != "" {
+		kong.ServiceRequest.SetHeader(conf.PublicKeyHeader, tok.Claims.(*signing.Claims).Subject)
+	}
+
 	kong.Log.Warn(tok)
 
 }