diff --git a/containers/agent/one-shot-token/README.md b/containers/agent/one-shot-token/README.md index dc2c0b2a..b00f09af 100644 --- a/containers/agent/one-shot-token/README.md +++ b/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 diff --git a/containers/agent/one-shot-token/one-shot-token.c b/containers/agent/one-shot-token/one-shot-token.c index c80a8f3e..c6dba82b 100644 --- a/containers/agent/one-shot-token/one-shot-token.c +++ b/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 #include #include +#include -/* 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]) {