age: optional caching of key command output

age/keysource.go

@@ -10,7 +10,9 @@ import (
 	"os/exec"
 	"path/filepath"
 	"runtime"
+	"strconv"
 	"strings"
+	"sync"
 
 	"filippo.io/age"
 	"filippo.io/age/agessh"
@@ -43,6 +45,10 @@ const (
 	// SopsAgeSshPrivateKeyFileEnv can be set as an environment variable pointing to
 	// a private SSH key file.
 	SopsAgeSshPrivateKeyFileEnv = "SOPS_AGE_SSH_PRIVATE_KEY_FILE"
+	// SopsAgeKeyCmdCacheEnv can be set to a value considered true by
+	// strconv.ParseBool to execute each key command at most once per process,
+	// reusing its first output for every recipient.
+	SopsAgeKeyCmdCacheEnv = "SOPS_AGE_KEY_CMD_CACHE"
 	// SopsAgeKeyUserConfigPath is the default age keys file path in
 	// getUserConfigDir().
 	SopsAgeKeyUserConfigPath = "sops/age/keys.txt"
@@ -294,10 +300,44 @@ func (key *MasterKey) TypeToIdentifier() string {
 	return KeyTypeIdentifier
 }
 
+// cmdOutputCache memoises successful key command outputs per command string
+// when SopsAgeKeyCmdCacheEnv is set. Errors are never cached, and the lock is
+// not held while the command runs, so a slow or failing command cannot block
+// or poison other decrypts in long-lived processes such as the keyservice.
+var cmdOutputCache = struct {
+	sync.Mutex
+	results map[string][]byte
+}{results: make(map[string][]byte)}
+
 // getOutputFromCmd executes a shell command provided in param 'cmdString',
 // optionally adding env vars provided in param 'envVars',
-// and returns the command's output and error
+// and returns the command's output and error. If SopsAgeKeyCmdCacheEnv is
+// set to a true value, successful output is cached per command string for
+// the lifetime of the process, ignoring envVars on cache hits.
 func getOutputFromCmd(cmdString string, envVars []string) ([]byte, error) {
+	cache, err := strconv.ParseBool(os.Getenv(SopsAgeKeyCmdCacheEnv))
+	if err != nil || !cache {
+		return runCmd(cmdString, envVars)
+	}
+	cmdOutputCache.Lock()
+	out, ok := cmdOutputCache.results[cmdString]
+	cmdOutputCache.Unlock()
+	if ok {
+		return out, nil
+	}
+	out, err = runCmd(cmdString, envVars)
+	if err == nil {
+		cmdOutputCache.Lock()
+		cmdOutputCache.results[cmdString] = out
+		cmdOutputCache.Unlock()
+	}
+	return out, err
+}
+
+// runCmd executes a shell command provided in param 'cmdString',
+// optionally adding env vars provided in param 'envVars',
+// and returns the command's output and error
+func runCmd(cmdString string, envVars []string) ([]byte, error) {
 	var out []byte
 
 	args, err := shlex.Split(cmdString)

age/keysource_test.go

@@ -575,6 +575,79 @@ func TestMasterKey_loadIdentities(t *testing.T) {
 		assert.Nil(t, got)
 		assert.Len(t, unusedLocations, 7)
 	})
+
+	t.Run(SopsAgeKeyCmdCacheEnv, func(t *testing.T) {
+		tmpDir := t.TempDir()
+		// Overwrite to ensure local config is not picked up by tests
+		overwriteUserConfigDir(t, tmpDir)
+
+		countFile := filepath.Join(tmpDir, "count")
+		t.Setenv(SopsAgeKeyCmdEnv, fmt.Sprintf("bash -c 'echo x >> %s; echo %s'", countFile, mockIdentity))
+		t.Setenv(SopsAgeKeyCmdCacheEnv, "true")
+
+		// The command runs once; the second recipient is served from the cache.
+		for _, recipient := range []string{mockRecipient, mockRecipient + "abc"} {
+			key := &MasterKey{Recipient: recipient}
+			got, unusedLocations, errs := key.loadIdentities()
+			assert.Len(t, errs, 0)
+			assert.Len(t, got, 1)
+			assert.Len(t, unusedLocations, 6)
+		}
+		count, err := os.ReadFile(countFile)
+		assert.NoError(t, err)
+		assert.Equal(t, "x\n", string(count))
+	})
+
+	t.Run(SopsAgeKeyCmdCacheEnv+" does not cache errors", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		// Overwrite to ensure local config is not picked up by tests
+		overwriteUserConfigDir(t, tmpDir)
+
+		flagFile := filepath.Join(tmpDir, "flag")
+		countFile := filepath.Join(tmpDir, "count")
+		t.Setenv(SopsAgeKeyCmdEnv, fmt.Sprintf(
+			"bash -c 'echo x >> %s; if [ ! -f %s ]; then touch %s; exit 1; fi; echo %s'",
+			countFile, flagFile, flagFile, mockIdentity))
+		t.Setenv(SopsAgeKeyCmdCacheEnv, "true")
+
+		// First run fails and must not be cached.
+		key := &MasterKey{Recipient: mockRecipient}
+		got, _, errs := key.loadIdentities()
+		assert.Len(t, errs, 1)
+		assert.Nil(t, got)
+
+		// Second run succeeds and is cached; the third is served from cache.
+		for range 2 {
+			key = &MasterKey{Recipient: mockRecipient}
+			got, _, errs = key.loadIdentities()
+			assert.Len(t, errs, 0)
+			assert.Len(t, got, 1)
+		}
+		count, err := os.ReadFile(countFile)
+		assert.NoError(t, err)
+		assert.Equal(t, "x\nx\n", string(count))
+	})
+
+	t.Run(SopsAgeKeyCmdCacheEnv+" disabled", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		// Overwrite to ensure local config is not picked up by tests
+		overwriteUserConfigDir(t, tmpDir)
+
+		countFile := filepath.Join(tmpDir, "count")
+		t.Setenv(SopsAgeKeyCmdEnv, fmt.Sprintf("bash -c 'echo x >> %s; echo %s'", countFile, mockIdentity))
+
+		// Without the cache flag the command runs once per recipient.
+		for _, recipient := range []string{mockRecipient, mockRecipient + "abc"} {
+			key := &MasterKey{Recipient: recipient}
+			got, unusedLocations, errs := key.loadIdentities()
+			assert.Len(t, errs, 0)
+			assert.Len(t, got, 1)
+			assert.Len(t, unusedLocations, 6)
+		}
+		count, err := os.ReadFile(countFile)
+		assert.NoError(t, err)
+		assert.Equal(t, "x\nx\n", string(count))
+	})
 }
 
 // overwriteUserConfigDir sets the user config directory and the user home directory