support VAULT_RUN_MODE envvar (#104)

Author: raymonstah

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 IMPROVEMENTS:
 * Requests from the extension to Vault now set the User-Agent field accordingly.
+* Introduced a `VAULT_RUN_MODE` environment variable to allow user to run in proxy mode, file mode, or both. 
+  * The default value is 'default', which runs in *both* proxy and file mode.
 
 ## 0.9.0 (February 23, 2023)
 

internal/config/info.go

@@ -4,4 +4,5 @@ const (
 	ExtensionName    = "vault-lambda-extension"
 	ExtensionVersion = "0.9.0"
 	VaultLogLevel    = "VAULT_LOG_LEVEL" // Optional, one of TRACE, DEBUG, INFO, WARN, ERROR, OFF
+	VaultRunMode     = "VAULT_RUN_MODE"
 )

internal/extension/client.go

@@ -8,7 +8,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"net/http"
 )
 
@@ -89,7 +89,7 @@ func (e *Client) Register(ctx context.Context, filename string) (*RegisterRespon
 		return nil, fmt.Errorf("request failed with status %s", httpRes.Status)
 	}
 	defer httpRes.Body.Close()
-	body, err := ioutil.ReadAll(httpRes.Body)
+	body, err := io.ReadAll(httpRes.Body)
 	if err != nil {
 		return nil, err
 	}
@@ -122,7 +122,7 @@ func (e *Client) NextEvent(ctx context.Context) (*NextEventResponse, error) {
 		return nil, fmt.Errorf("request failed with status %s", httpRes.Status)
 	}
 	defer httpRes.Body.Close()
-	body, err := ioutil.ReadAll(httpRes.Body)
+	body, err := io.ReadAll(httpRes.Body)
 	if err != nil {
 		return nil, err
 	}

internal/proxy/cache.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"encoding/hex"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"strings"
@@ -74,7 +75,7 @@ func NewCache(cc config.CacheConfig) *Cache {
 // constructs the CacheKey for this request and token and returns the SHA256
 // hash
 func makeRequestHash(logger hclog.Logger, r *http.Request, token string) (string, error) {
-	reqBody, err := ioutil.ReadAll(r.Body)
+	reqBody, err := io.ReadAll(r.Body)
 	if err != nil {
 		if r.Body != nil {
 			if err := r.Body.Close(); err != nil {

internal/proxy/proxy_test.go

@@ -7,7 +7,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -81,7 +81,7 @@ func TestProxy(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, http.StatusOK, resp.StatusCode)
 		defer resp.Body.Close()
-		body, err := ioutil.ReadAll(resp.Body)
+		body, err := io.ReadAll(resp.Body)
 		require.NoError(t, err)
 		var secret api.Secret
 		require.NoError(t, json.Unmarshal(body, &secret), string(body))

internal/runmode/runmode.go

@@ -0,0 +1,38 @@
+// Package runmode determines if the vault lambda extension should run in default mode, file mode, or proxy mode.
+// default mode: uses both file and proxy mode
+// file mode: writes secrets to disk
+// proxy mode: forwards requests to a Vault server
+package runmode
+
+import "strings"
+
+type Mode string
+
+var (
+	ModeDefault Mode = "default"
+	ModeFile    Mode = "file"
+	ModeProxy   Mode = "proxy"
+)
+
+var modes = map[Mode]struct{}{
+	ModeFile:    {},
+	ModeProxy:   {},
+	ModeDefault: {},
+}
+
+func ParseMode(rm string) Mode {
+	_, ok := modes[Mode(strings.ToLower(rm))]
+	if !ok {
+		return ModeDefault
+	}
+
+	return Mode(rm)
+}
+
+func (m Mode) HasModeProxy() bool {
+	return m == ModeDefault || m == ModeProxy
+}
+
+func (m Mode) HasModeFile() bool {
+	return m == ModeDefault || m == ModeFile
+}

internal/vault/client.go

@@ -8,7 +8,7 @@ import (
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"os"
 	"strings"
 	"sync"
@@ -105,7 +105,6 @@ func (c *Client) Token(ctx context.Context) (string, error) {
 
 // login authenticates to Vault using IAM auth, and sets the client's token.
 func (c *Client) login(ctx context.Context) error {
-
 	authConfig := config.AuthConfigFromEnv()
 	roleToAssumeArn := authConfig.AssumedRoleArn
 
@@ -114,16 +113,13 @@ func (c *Client) login(ctx context.Context) error {
 	/* If passing in a role (through VAULT_ASSUMED_ROLE_ARN enviornment variable)
 	to be assumed for Vault authentication, use it instead of the function execution role */
 	if roleToAssumeArn != "" {
-
 		c.logger.Debug(fmt.Sprintf("Trying to assume role with arn of %s to authenticate with Vault", roleToAssumeArn))
-
 		sessionName := "vault_auth"
 
 		result, err := c.stsSvc.AssumeRole(&sts.AssumeRoleInput{
 			RoleArn:         &roleToAssumeArn,
 			RoleSessionName: &sessionName,
 		})
-
 		if err != nil {
 			return fmt.Errorf("failed to assume role with arn of %s %w", roleToAssumeArn, err)
 		}
@@ -163,7 +159,7 @@ func (c *Client) login(ctx context.Context) error {
 		return err
 	}
 
-	body, err := ioutil.ReadAll(req.HTTPRequest.Body)
+	body, err := io.ReadAll(req.HTTPRequest.Body)
 	if err != nil {
 		return err
 	}

main.go

@@ -8,7 +8,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io/ioutil"
 	"net"
 	"net/http"
 	"os"
@@ -24,6 +23,7 @@ import (
 	"github.com/hashicorp/vault-lambda-extension/internal/config"
 	"github.com/hashicorp/vault-lambda-extension/internal/extension"
 	"github.com/hashicorp/vault-lambda-extension/internal/proxy"
+	"github.com/hashicorp/vault-lambda-extension/internal/runmode"
 	"github.com/hashicorp/vault-lambda-extension/internal/vault"
 	"github.com/hashicorp/vault/api"
 )
@@ -33,19 +33,36 @@ func main() {
 		Level: hclog.LevelFromString(os.Getenv(config.VaultLogLevel)),
 	})
 
-	err := realMain(logger.Named(config.VaultLogLevel))
-	if err != nil {
+	runMode := runmode.ModeDefault
+	if runModeEnv := os.Getenv(config.VaultRunMode); runModeEnv != "" {
+		runMode = runmode.ParseMode(runModeEnv)
+	}
+
+	h := newHandler(logger.Named(config.ExtensionName), runMode)
+	if err := h.handle(); err != nil {
 		logger.Error("Fatal error, exiting", "error", err)
 		os.Exit(1)
 	}
 }
 
-func realMain(logger hclog.Logger) error {
+func newHandler(logger hclog.Logger, runMode runmode.Mode) *handler {
+	return &handler{
+		logger:  logger,
+		runMode: runMode,
+	}
+}
+
+type handler struct {
+	logger  hclog.Logger
+	runMode runmode.Mode
+}
+
+func (h *handler) handle() error {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
 	var wg sync.WaitGroup
-	srv, err := runExtension(ctx, logger, &wg)
+	cleanup, err := h.runExtension(ctx, &wg)
 	if err != nil {
 		return err
 	}
@@ -57,16 +74,16 @@ func realMain(logger hclog.Logger) error {
 		interruptChannel := make(chan os.Signal, 1)
 		signal.Notify(interruptChannel, syscall.SIGTERM, syscall.SIGINT)
 		select {
-		case s := <-interruptChannel:
-			logger.Info("Received signal, exiting", "signal", s)
+		case sig := <-interruptChannel:
+			h.logger.Info("Received signal, exiting", "signal", sig)
 		case <-shutdownChannel:
-			logger.Info("Received shutdown event, exiting")
+			h.logger.Info("Received shutdown event, exiting")
 		}
 
 		cancel()
-		if err := srv.Shutdown(context.Background()); err != nil {
+		if err := cleanup(context.Background()); err != nil {
 			// Error from closing listeners, or context timeout:
-			logger.Error("HTTP server shutdown error", "error", err)
+			h.logger.Error("HTTP server shutdown error", "error", err)
 		}
 	}()
 
@@ -76,20 +93,20 @@ func realMain(logger hclog.Logger) error {
 		return err
 	}
 
-	processEvents(ctx, logger, extensionClient)
+	processEvents(ctx, h.logger, extensionClient)
 
 	// Once processEvents returns, signal that it's time to shutdown.
 	shutdownChannel <- struct{}{}
 
 	// Ensure we wait for the HTTP server to gracefully shut down.
 	wg.Wait()
-	logger.Info("Graceful shutdown complete")
+	h.logger.Info("Graceful shutdown complete")
 
 	return nil
 }
 
-func runExtension(ctx context.Context, logger hclog.Logger, wg *sync.WaitGroup) (*http.Server, error) {
-	logger.Info("Initialising")
+func (h *handler) runExtension(ctx context.Context, wg *sync.WaitGroup) (func(context.Context) error, error) {
+	h.logger.Info("Initialising")
 
 	authConfig := config.AuthConfigFromEnv()
 	vaultConfig := api.DefaultConfig()
@@ -114,7 +131,7 @@ func runExtension(ctx context.Context, logger hclog.Logger, wg *sync.WaitGroup)
 	} else {
 		ses = session.Must(session.NewSession())
 	}
-	client, err := vault.NewClient(config.ExtensionName, config.ExtensionVersion, logger.Named("vault-client"), vaultConfig, authConfig, ses)
+	client, err := vault.NewClient(config.ExtensionName, config.ExtensionVersion, h.logger.Named("vault-client"), vaultConfig, authConfig, ses)
 	if err != nil {
 		return nil, fmt.Errorf("error getting client: %w", err)
 	} else if client == nil {
@@ -135,34 +152,41 @@ func runExtension(ctx context.Context, logger hclog.Logger, wg *sync.WaitGroup)
 
 	client.VaultClient = client.VaultClient.WithRequestCallbacks(api.RequireState(newState), vault.UserAgentRequestCallback(uaFunc)).WithResponseCallbacks()
 
-	err = writePreconfiguredSecrets(client.VaultClient)
-	if err != nil {
-		return nil, err
+	if h.runMode.HasModeFile() {
+		if err := writePreconfiguredSecrets(client.VaultClient); err != nil {
+			return nil, err
+		}
 	}
 
 	// clear out eventual consistency helpers
 	client.VaultClient = client.VaultClient.WithRequestCallbacks().WithResponseCallbacks()
 
-	ln, err := net.Listen("tcp", "127.0.0.1:8200")
-	if err != nil {
-		return nil, fmt.Errorf("failed to listen on port 8200: %w", err)
-	}
-	srv := proxy.New(logger.Named("proxy"), client, config.CacheConfigFromEnv())
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		logger.Info("Starting HTTP proxy server")
-		err = srv.Serve(ln)
-		if err != http.ErrServerClosed {
-			logger.Error("HTTP server shutdown unexpectedly", "error", err)
+	cleanupFunc := func(context.Context) error { return nil }
+	if h.runMode.HasModeProxy() {
+		ln, err := net.Listen("tcp", "127.0.0.1:8200")
+		if err != nil {
+			return nil, fmt.Errorf("failed to listen on port 8200: %w", err)
 		}
-	}()
-
-	logger.Info("Initialised")
+		srv := proxy.New(h.logger.Named("proxy"), client, config.CacheConfigFromEnv())
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			h.logger.Info("Starting HTTP proxy server")
+			err = srv.Serve(ln)
+			if err != http.ErrServerClosed {
+				h.logger.Error("HTTP server shutdown unexpectedly", "error", err)
+			}
+		}()
+		cleanupFunc = func(ctx context.Context) error {
+			return srv.Shutdown(ctx)
+		}
+	}
 
-	return srv, nil
+	h.logger.Info("Initialised")
+	return cleanupFunc, nil
 }
 
+// writePreconfiguredSecrets writes secrets to disk.
 func writePreconfiguredSecrets(client *api.Client) error {
 	configuredSecrets, err := config.ParseConfiguredSecrets()
 	if err != nil {
@@ -178,18 +202,18 @@ func writePreconfiguredSecrets(client *api.Client) error {
 
 		content, err := json.MarshalIndent(secret, "", "  ")
 		if err != nil {
-			return err
+			return fmt.Errorf("unable to marshal json: %w", err)
 		}
+
 		dir := path.Dir(s.FilePath)
 		if _, err = os.Stat(dir); os.IsNotExist(err) {
-			err = os.MkdirAll(dir, 0755)
-			if err != nil {
+			if err := os.MkdirAll(dir, 0755); err != nil {
 				return fmt.Errorf("failed to create directory %q for secret %s: %s", dir, s.Name(), err)
 			}
 		}
-		err = ioutil.WriteFile(s.FilePath, content, 0644)
-		if err != nil {
-			return err
+
+		if err := os.WriteFile(s.FilePath, content, 0644); err != nil {
+			return fmt.Errorf("error writing file: %w", err)
 		}
 	}