feat: add runtime configuration for one-shot token protection via AWF_ONE_SHOT_TOKENS

containers/agent/one-shot-token/README.md

@@ -2,13 +2,15 @@
 
 ## Overview
 
-The one-shot token library is an `LD_PRELOAD` shared library that provides **single-use access** to sensitive GitHub token environment variables. When a process reads a protected token via `getenv()`, the library returns the value once and immediately unsets the environment variable, preventing subsequent reads.
+The one-shot token library is an `LD_PRELOAD` shared library that provides **single-use access** to sensitive token environment variables. When a process reads a protected token via `getenv()`, the library returns the value once and immediately unsets the environment variable, preventing subsequent reads.
 
 This protects against malicious code that might attempt to exfiltrate tokens after the legitimate application has already consumed them.
 
-## Protected Environment Variables
+## Configuration
 
-The library intercepts access to these token variables:
+### Default Protected Tokens
+
+By default, the library protects these token variables:
 
 **GitHub:**
 - `COPILOT_GITHUB_TOKEN`
@@ -29,6 +31,24 @@ The library intercepts access to these token variables:
 **Codex:**
 - `CODEX_API_KEY`
 
+### Custom Token List
+
+You can configure a custom list of tokens to protect using the `AWF_ONE_SHOT_TOKENS` environment variable:
+
+```bash
+# Protect custom tokens instead of defaults
+export AWF_ONE_SHOT_TOKENS="MY_API_KEY,MY_SECRET_TOKEN,CUSTOM_AUTH_KEY"
+
+# Run your command with the library preloaded
+LD_PRELOAD=/usr/local/lib/one-shot-token.so ./your-program
+```
+
+**Important notes:**
+- When `AWF_ONE_SHOT_TOKENS` is set, **only** the tokens in this list are protected (defaults are not included)
+- Use comma-separated token names (whitespace is automatically trimmed)
+- Maximum of 100 tokens can be protected
+- The configuration is read once at library initialization (first `getenv()` call)
+
 ## How It Works
 
 ### The LD_PRELOAD Mechanism
@@ -154,11 +174,13 @@ This produces `one-shot-token.so` in the current directory.
 
 ## Testing
 
+### Basic Test (Default Tokens)
+
 ```bash
 # Build the library
 ./build.sh
 
-# Test with a simple program
+# Test with a default protected token
 export GITHUB_TOKEN="test-token-12345"
 LD_PRELOAD=./one-shot-token.so bash -c '
   echo "First read: $(printenv GITHUB_TOKEN)"
@@ -168,11 +190,57 @@ LD_PRELOAD=./one-shot-token.so bash -c '
 
 Expected output:
 ```
+[one-shot-token] Initialized with 11 default token(s)
 [one-shot-token] Token GITHUB_TOKEN accessed and cleared
 First read: test-token-12345
 Second read:
 ```
 
+### Custom Token Test
+
+```bash
+# Build the library
+./build.sh
+
+# Test with custom tokens
+export AWF_ONE_SHOT_TOKENS="MY_API_KEY,SECRET_TOKEN"
+export MY_API_KEY="secret-value-123"
+export SECRET_TOKEN="another-secret"
+
+LD_PRELOAD=./one-shot-token.so bash -c '
+  echo "First MY_API_KEY: $(printenv MY_API_KEY)"
+  echo "Second MY_API_KEY: $(printenv MY_API_KEY)"
+  echo "First SECRET_TOKEN: $(printenv SECRET_TOKEN)"
+  echo "Second SECRET_TOKEN: $(printenv SECRET_TOKEN)"
+'
+```
+
+Expected output:
+```
+[one-shot-token] Initialized with 2 custom token(s) from AWF_ONE_SHOT_TOKENS
+[one-shot-token] Token MY_API_KEY accessed and cleared
+First MY_API_KEY: secret-value-123
+Second MY_API_KEY:
+[one-shot-token] Token SECRET_TOKEN accessed and cleared
+First SECRET_TOKEN: another-secret
+Second SECRET_TOKEN:
+```
+
+### Integration with AWF
+
+When using the library with AWF (Agentic Workflow Firewall):
+
+```bash
+# Use default tokens
+sudo awf --allow-domains github.com -- your-command
+
+# Use custom tokens
+export AWF_ONE_SHOT_TOKENS="MY_TOKEN,CUSTOM_API_KEY"
+sudo -E awf --allow-domains github.com -- your-command
+```
+
+Note: The `AWF_ONE_SHOT_TOKENS` variable must be exported before running `awf` so it's available when the library initializes.
+
 ## Security Considerations
 
 ### What This Protects Against

containers/agent/one-shot-token/one-shot-token.c

@@ -5,6 +5,10 @@
  * On first access, returns the real value and immediately unsets the variable.
  * Subsequent calls return NULL, preventing token reuse by malicious code.
  *
+ * Configuration:
+ *   AWF_ONE_SHOT_TOKENS - Comma-separated list of token names to protect
+ *   If not set, uses built-in defaults
+ *
  * Compile: gcc -shared -fPIC -o one-shot-token.so one-shot-token.c -ldl
  * Usage: LD_PRELOAD=/path/to/one-shot-token.so ./your-program
  */
@@ -15,9 +19,10 @@
 #include <string.h>
 #include <pthread.h>
 #include <stdio.h>
+#include <ctype.h>
 
-/* Sensitive token environment variable names */
-static const char *SENSITIVE_TOKENS[] = {
+/* Default sensitive token environment variable names */
+static const char *DEFAULT_SENSITIVE_TOKENS[] = {
     /* GitHub tokens */
     "COPILOT_GITHUB_TOKEN",
     "GITHUB_TOKEN",
@@ -36,12 +41,22 @@ static const char *SENSITIVE_TOKENS[] = {
     NULL
 };
 
+/* Maximum number of tokens we can track (for static allocation) */
+#define MAX_TOKENS 100
+
+/* Runtime token list (populated from AWF_ONE_SHOT_TOKENS or defaults) */
+static char *sensitive_tokens[MAX_TOKENS];
+static int num_tokens = 0;
+
 /* Track which tokens have been accessed (one flag per token) */
-static int token_accessed[sizeof(SENSITIVE_TOKENS) / sizeof(SENSITIVE_TOKENS[0])] = {0};
+static int token_accessed[MAX_TOKENS] = {0};
 
 /* Mutex for thread safety */
 static pthread_mutex_t token_mutex = PTHREAD_MUTEX_INITIALIZER;
 
+/* Initialization flag */
+static int tokens_initialized = 0;
+
 /* Pointer to the real getenv function */
 static char *(*real_getenv)(const char *name) = NULL;
 
@@ -57,12 +72,88 @@ static void init_real_getenv(void) {
     }
 }
 
+/**
+ * Initialize the token list from AWF_ONE_SHOT_TOKENS environment variable
+ * or use defaults if not set. This is called once at first getenv() call.
+ * Note: This function must be called with token_mutex held.
+ */
+static void init_token_list(void) {
+    if (tokens_initialized) {
+        return;
+    }
+
+    /* Get the configuration from environment */
+    const char *config = real_getenv("AWF_ONE_SHOT_TOKENS");
+
+    if (config != NULL && config[0] != '\0') {
+        /* Parse comma-separated token list */
+        char *config_copy = strdup(config);
+        if (config_copy == NULL) {
+            fprintf(stderr, "[one-shot-token] ERROR: Failed to allocate memory for token list\n");
+            abort();
+        }
+
+        char *token = strtok(config_copy, ",");
+        while (token != NULL && num_tokens < MAX_TOKENS) {
+            /* Trim leading whitespace */
+            while (*token && isspace((unsigned char)*token)) token++;
+
+            /* Trim trailing whitespace (only if string is non-empty) */
+            if (strlen(token) > 0) {
+                char *end = token + strlen(token) - 1;
+                while (end > token && isspace((unsigned char)*end)) {
+                    *end = '\0';
+                    end--;
+                }
+            }
+
+            if (*token != '\0') {
+                sensitive_tokens[num_tokens] = strdup(token);
+                if (sensitive_tokens[num_tokens] == NULL) {
+                    fprintf(stderr, "[one-shot-token] ERROR: Failed to allocate memory for token name\n");
+                    /* Clean up previously allocated tokens */
+                    for (int i = 0; i < num_tokens; i++) {
+                        free(sensitive_tokens[i]);
+                    }
+                    free(config_copy);
+                    abort();
+                }
+                num_tokens++;
+            }
+
+            token = strtok(NULL, ",");
+        }
+
+        free(config_copy);
+
+        fprintf(stderr, "[one-shot-token] Initialized with %d custom token(s) from AWF_ONE_SHOT_TOKENS\n", num_tokens);
+    } else {
+        /* Use default token list */
+        for (int i = 0; DEFAULT_SENSITIVE_TOKENS[i] != NULL && num_tokens < MAX_TOKENS; i++) {
+            sensitive_tokens[num_tokens] = strdup(DEFAULT_SENSITIVE_TOKENS[i]);
+            if (sensitive_tokens[num_tokens] == NULL) {
+                fprintf(stderr, "[one-shot-token] ERROR: Failed to allocate memory for default token name\n");
+                /* Clean up previously allocated tokens */
+                for (int j = 0; j < num_tokens; j++) {
+                    free(sensitive_tokens[j]);
+                }
+                abort();
+            }
+            num_tokens++;
+        }
+
+        fprintf(stderr, "[one-shot-token] Initialized with %d default token(s)\n", num_tokens);
+    }
+
+    tokens_initialized = 1;
+}
+
 /* Check if a variable name is a sensitive token */
 static int get_token_index(const char *name) {
     if (name == NULL) return -1;
 
-    for (int i = 0; SENSITIVE_TOKENS[i] != NULL; i++) {
-        if (strcmp(name, SENSITIVE_TOKENS[i]) == 0) {
+    for (int i = 0; i < num_tokens; i++) {
+        if (strcmp(name, sensitive_tokens[i]) == 0) {
             return i;
         }
     }
@@ -81,16 +172,22 @@ static int get_token_index(const char *name) {
 char *getenv(const char *name) {
     init_real_getenv();
 
+    /* Initialize token list on first call (thread-safe) */
+    pthread_mutex_lock(&token_mutex);
+    if (!tokens_initialized) {
+        init_token_list();
+    }
+
+    /* Get token index while holding mutex to avoid race with initialization */
     int token_idx = get_token_index(name);
 
-    /* Not a sensitive token - pass through */
+    /* Not a sensitive token - release mutex and pass through */
     if (token_idx < 0) {
+        pthread_mutex_unlock(&token_mutex);
         return real_getenv(name);
     }
 
-    /* Sensitive token - handle one-shot access */
-    pthread_mutex_lock(&token_mutex);
-
+    /* Sensitive token - handle one-shot access (mutex already held) */
     char *result = NULL;
 
     if (!token_accessed[token_idx]) {