set user agent for outbound requests (#98)

Author: kpcraig

CHANGELOG.md

@@ -1,5 +1,8 @@
 ## Unreleased
 
+IMPROVEMENTS:
+* Requests from the extension to Vault now set the User-Agent field accordingly.
+
 ## 0.9.0 (February 23, 2023)
 
 IMPROVEMENTS:

internal/config/info.go

@@ -0,0 +1,7 @@
+package config
+
+const (
+	ExtensionName    = "vault-lambda-extension"
+	ExtensionVersion = "0.9.0"
+	VaultLogLevel    = "VAULT_LOG_LEVEL" // Optional, one of TRACE, DEBUG, INFO, WARN, ERROR, OFF
+)

internal/config/useragent.go

@@ -0,0 +1,12 @@
+package config
+
+import (
+	"fmt"
+	"runtime"
+)
+
+// GetUserAgentBase returns a base user agent string with the given user agent name and version in the form:
+// vault-client-go/0.0.1 (Darwin arm64; Go go1.19.2)
+func GetUserAgentBase(clientName string, clientVersion string) string {
+	return fmt.Sprintf("%s/%s (%s %s; Go %s)", clientName, clientVersion, runtime.GOOS, runtime.GOARCH, runtime.Version())
+}

internal/proxy/proxy.go

@@ -83,6 +83,10 @@ func proxyHandler(logger hclog.Logger, client *vault.Client, cache *Cache) func(
 			}
 		}
 
+		// add user agent header
+		ua := config.GetUserAgentBase(client.Name, client.Version)
+		fwReq.Header.Set("User-Agent", ua+"; requesting from proxy")
+
 		resp, err := client.VaultConfig.HttpClient.Do(fwReq)
 		if err != nil {
 			http.Error(w, fmt.Sprintf("failed to proxy request: %s", err), http.StatusBadGateway)

internal/proxy/proxy_test.go

@@ -30,6 +30,7 @@ type vaultResponse struct {
 }
 
 var (
+	vaultRequests       []*http.Request
 	fakeVaultResponse   vaultResponse
 	vaultResponseFooBar = vaultResponse{
 		secret: &api.Secret{
@@ -69,9 +70,14 @@ func TestProxy(t *testing.T) {
 	defer close()
 
 	t.Run("happy path bare http client", func(t *testing.T) {
+		// reset request array
+		vaultRequests = []*http.Request{}
 		fakeVaultResponse = vaultResponseFooBar
 		resp, err := http.Get(fmt.Sprintf("http://%s/v1/secret/data/foo", proxyAddr))
 
+		// the stored request should be the one _from the proxy_ since it's stored by
+		// the (fake) vault.
+		require.NotEmpty(t, vaultRequests[0].Header.Get("User-Agent"))
 		require.NoError(t, err)
 		require.Equal(t, http.StatusOK, resp.StatusCode)
 		defer resp.Body.Close()
@@ -83,12 +89,15 @@ func TestProxy(t *testing.T) {
 	})
 
 	t.Run("happy path with vault client", func(t *testing.T) {
+		vaultRequests = []*http.Request{}
 		fakeVaultResponse = vaultResponseFooBar
 		proxyVaultClient, err := api.NewClient(&api.Config{
 			Address: "http://" + proxyAddr,
 		})
 		require.NoError(t, err)
 		resp, err := proxyVaultClient.Logical().Read("secret/data/foo")
+
+		require.NotEmpty(t, vaultRequests[0].Header.Get("User-Agent"))
 		require.NoError(t, err)
 		require.Equal(t, "bar", resp.Data["foo"])
 	})
@@ -103,6 +112,7 @@ func TestProxy(t *testing.T) {
 	})
 
 	t.Run("failed upstream request should give 502", func(t *testing.T) {
+
 		fakeVaultResponse = vaultResponseFooBar
 		resp, err := http.Get(fmt.Sprintf("http://%s/FailedTransport", proxyAddr))
 		require.NoError(t, err)
@@ -114,7 +124,7 @@ func startProxy(t *testing.T, vaultAddress string, ses *session.Session) (string
 	vaultConfig := api.DefaultConfig()
 	require.NoError(t, vaultConfig.Error)
 	vaultConfig.Address = vaultAddress
-	client, err := vault.NewClient(hclog.NewNullLogger(), vaultConfig, config.AuthConfig{}, ses)
+	client, err := vault.NewClient("", "", hclog.NewNullLogger(), vaultConfig, config.AuthConfig{}, ses)
 	require.NoError(t, err)
 	client.VaultConfig.Address = vaultAddress
 	ln, err := net.Listen("tcp", "127.0.0.1:0")
@@ -130,6 +140,10 @@ func startProxy(t *testing.T, vaultAddress string, ses *session.Session) (string
 func fakeVault() *httptest.Server {
 	vaultResponsePtr := &fakeVaultResponse
 	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// after handling, save the request so we can inspect it in test cases
+		defer func() {
+			vaultRequests = append(vaultRequests, r)
+		}()
 		switch {
 		case strings.Contains(r.URL.Path, "login"):
 			b, err := json.Marshal(vaultLoginResponse)

internal/vault/client.go

@@ -32,6 +32,9 @@ const (
 
 // Client holds api.Client and handles state required to renew tokens and re-auth as required.
 type Client struct {
+	Name    string
+	Version string
+
 	mtx sync.Mutex
 
 	VaultClient *api.Client
@@ -50,7 +53,7 @@ type Client struct {
 
 // NewClient uses the AWS IAM auth method configured in a Vault cluster to
 // authenticate the execution role and create a Vault API client.
-func NewClient(logger hclog.Logger, vaultConfig *api.Config, authConfig config.AuthConfig, awsSes *session.Session) (*Client, error) {
+func NewClient(name, version string, logger hclog.Logger, vaultConfig *api.Config, authConfig config.AuthConfig, awsSes *session.Session) (*Client, error) {
 	vaultClient, err := api.NewClient(vaultConfig)
 	if err != nil {
 		return nil, fmt.Errorf("error making extension: %w", err)
@@ -64,6 +67,8 @@ func NewClient(logger hclog.Logger, vaultConfig *api.Config, authConfig config.A
 	client := &Client{
 		VaultClient: vaultClient,
 		VaultConfig: vaultConfig,
+		Name:        name,
+		Version:     version,
 
 		logger:     logger,
 		stsSvc:     sts.New(awsSes),
@@ -238,3 +243,11 @@ func parseTokenExpiryGracePeriod() (time.Duration, error) {
 
 	return expiryGracePeriod, nil
 }
+
+// UserAgentRequestCallback takes a function that returns a user agent string and will invoke that function to set
+// the user agent string on the request.
+func UserAgentRequestCallback(agentFunc func(request *api.Request) string) api.RequestCallback {
+	return func(req *api.Request) {
+		req.Headers.Set("User-Agent", agentFunc(req))
+	}
+}

internal/vault/client_test.go

@@ -296,6 +296,49 @@ func TestParseTokenExpiryGracePeriod(t *testing.T) {
 	require.Error(t, err)
 }
 
+const userAgent = "abcd"
+
+func TestUserAgentHeaderAddition(t *testing.T) {
+	vault := fakeVault()
+	generateVaultClient := func() *api.Client {
+		vaultClient, err := api.NewClient(&api.Config{
+			Address: vault.URL,
+		})
+		require.NoError(t, err)
+		return vaultClient
+	}
+
+	t.Run("Ensure request contains header if decorator set", func(t *testing.T) {
+		vaultRequests = []*http.Request{}
+		vaultClient := generateVaultClient()
+		vaultClient.SetToken(t.Name())
+		c := Client{
+			VaultClient: vaultClient,
+			logger:      hclog.Default(),
+			//stsSvc:      stsSvc,
+
+			tokenRenewable: true,
+			tokenExpiry:    time.Now().Add(time.Hour),
+			tokenTTL:       10 * time.Hour,
+		}
+		secretFunc = generateSecretFunc(t, []*api.Secret{
+			with10hLease,
+		})
+		c.VaultClient = c.VaultClient.WithRequestCallbacks(UserAgentRequestCallback(fakeUserAgent))
+
+		_, err := c.Token(context.Background())
+		require.NoError(t, err)
+
+		// validate request was set and the user agent is what we expect
+		require.Equal(t, 1, len(vaultRequests))
+		require.Equal(t, userAgent, vaultRequests[0].Header.Get("User-Agent"))
+	})
+}
+
+func fakeUserAgent(_ *api.Request) string {
+	return userAgent
+}
+
 func fakeVault() *httptest.Server {
 	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		defer func() {

main.go

@@ -28,17 +28,12 @@ import (
 	"github.com/hashicorp/vault/api"
 )
 
-const (
-	extensionName = "vault-lambda-extension"
-	vaultLogLevel = "VAULT_LOG_LEVEL" // Optional, one of TRACE, DEBUG, INFO, WARN, ERROR, OFF
-)
-
 func main() {
 	logger := hclog.New(&hclog.LoggerOptions{
-		Level: hclog.LevelFromString(os.Getenv(vaultLogLevel)),
+		Level: hclog.LevelFromString(os.Getenv(config.VaultLogLevel)),
 	})
 
-	err := realMain(logger.Named(extensionName))
+	err := realMain(logger.Named(config.VaultLogLevel))
 	if err != nil {
 		logger.Error("Fatal error, exiting", "error", err)
 		os.Exit(1)
@@ -76,7 +71,7 @@ func realMain(logger hclog.Logger) error {
 	}()
 
 	extensionClient := extension.NewClient(os.Getenv("AWS_LAMBDA_RUNTIME_API"))
-	_, err = extensionClient.Register(ctx, extensionName)
+	_, err = extensionClient.Register(ctx, config.ExtensionName)
 	if err != nil {
 		return err
 	}
@@ -119,7 +114,7 @@ func runExtension(ctx context.Context, logger hclog.Logger, wg *sync.WaitGroup)
 	} else {
 		ses = session.Must(session.NewSession())
 	}
-	client, err := vault.NewClient(logger.Named("vault-client"), vaultConfig, authConfig, ses)
+	client, err := vault.NewClient(config.ExtensionName, config.ExtensionVersion, logger.Named("vault-client"), vaultConfig, authConfig, ses)
 	if err != nil {
 		return nil, fmt.Errorf("error getting client: %w", err)
 	} else if client == nil {
@@ -134,7 +129,11 @@ func runExtension(ctx context.Context, logger hclog.Logger, wg *sync.WaitGroup)
 		return nil, fmt.Errorf("error logging in to Vault: %w", err)
 	}
 
-	client.VaultClient = client.VaultClient.WithRequestCallbacks(api.RequireState(newState)).WithResponseCallbacks()
+	uaFunc := func(request *api.Request) string {
+		return config.GetUserAgentBase(config.ExtensionName, config.ExtensionVersion) + "; writing to temp file"
+	}
+
+	client.VaultClient = client.VaultClient.WithRequestCallbacks(api.RequireState(newState), vault.UserAgentRequestCallback(uaFunc)).WithResponseCallbacks()
 
 	err = writePreconfiguredSecrets(client.VaultClient)
 	if err != nil {