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

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

@@ -6,9 +6,11 @@ The one-shot token library is an `LD_PRELOAD` shared library that provides **sin
 
 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,26 @@ 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 with valid tokens, **only** those tokens are protected (defaults are not included)
+- If `AWF_ONE_SHOT_TOKENS` is set but contains only whitespace or commas (e.g., `"   "` or `",,,"`), the library falls back to the default token list to maintain protection
+- 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)
+- Uses `strtok_r()` internally, which is thread-safe and won't interfere with application code using `strtok()`
+
 ## How It Works
 
 ### The LD_PRELOAD Mechanism
@@ -154,6 +176,8 @@ This produces `one-shot-token.so` in the current directory.
 
 ## Testing
 
+### Basic Test (Default Tokens)
+
 ```bash
 # Build the library
 ./build.sh
@@ -184,11 +208,57 @@ LD_PRELOAD=./one-shot-token.so ./test_getenv
 
 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,24 @@ static const char *SENSITIVE_TOKENS[] = {
     NULL
 };
 
+/* Maximum number of tokens we can track (for static allocation). This limit
+ * balances memory usage with practical needs - 100 tokens should be more than
+ * sufficient for any reasonable use case while keeping memory overhead low. */
+#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;
 
@@ -58,6 +75,95 @@ static void init_real_getenv_once(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 using strtok_r for thread safety */
+        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 *saveptr = NULL;
+        char *token = strtok_r(config_copy, ",", &saveptr);
+        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) */
+            size_t token_len = strlen(token);
+            if (token_len > 0) {
+                char *end = token + token_len - 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_r(NULL, ",", &saveptr);
+        }
+
+        free(config_copy);
+
+        /* If AWF_ONE_SHOT_TOKENS was set but resulted in zero tokens (e.g., ",,," or whitespace only),
+         * fall back to defaults to avoid silently disabling all protection */
+        if (num_tokens == 0) {
+            fprintf(stderr, "[one-shot-token] WARNING: AWF_ONE_SHOT_TOKENS was set but parsed to zero tokens\n");
+            fprintf(stderr, "[one-shot-token] WARNING: Falling back to default token list to maintain protection\n");
+            /* num_tokens is already 0 here; assignment is defensive programming for future refactoring */
+            num_tokens = 0;
+        } else {
+            fprintf(stderr, "[one-shot-token] Initialized with %d custom token(s) from AWF_ONE_SHOT_TOKENS\n", num_tokens);
+            tokens_initialized = 1;
+            return;
+        }
+    }
+
+    /* Use default token list (when AWF_ONE_SHOT_TOKENS is unset, empty, or parsed to zero tokens) */
+    /* Note: num_tokens should be 0 when we reach here */
+    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;
+}
 /* Ensure real_getenv is initialized (thread-safe) */
 static void init_real_getenv(void) {
     pthread_once(&getenv_init_once, init_real_getenv_once);
@@ -67,8 +173,8 @@ static void init_real_getenv(void) {
 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;
         }
     }
@@ -87,16 +193,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]) {