feat: port one-shot-token from C to Rust and improve container security

containers/agent/Dockerfile

@@ -64,16 +64,21 @@ RUN chmod +x /usr/local/bin/setup-iptables.sh /usr/local/bin/entrypoint.sh /usr/
 
 # Build one-shot-token LD_PRELOAD library for single-use token access
 # This prevents tokens from being read multiple times (e.g., by malicious code)
-COPY one-shot-token/one-shot-token.c /tmp/one-shot-token.c
+COPY one-shot-token/Cargo.toml one-shot-token/src/ /tmp/one-shot-token/
+COPY one-shot-token/src/ /tmp/one-shot-token/src/
 RUN set -eux; \
-    BUILD_PKGS="gcc libc6-dev"; \
     apt-get update && \
-    ( apt-get install -y --no-install-recommends $BUILD_PKGS || \
-      (rm -rf /var/lib/apt/lists/* && apt-get update && \
-       apt-get install -y --no-install-recommends $BUILD_PKGS) ) && \
-    gcc -shared -fPIC -O2 -Wall -o /usr/local/lib/one-shot-token.so /tmp/one-shot-token.c -ldl -lpthread && \
-    rm /tmp/one-shot-token.c && \
-    apt-get remove -y $BUILD_PKGS && \
+    apt-get install -y --no-install-recommends curl build-essential && \
+    # Install Rust via rustup (minimal profile)
+    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal && \
+    export PATH="/root/.cargo/bin:$PATH" && \
+    # Build the library
+    cd /tmp/one-shot-token && \
+    cargo build --release && \
+    cp target/release/libone_shot_token.so /usr/local/lib/one-shot-token.so && \
+    # Clean up Rust and build artifacts
+    rm -rf /tmp/one-shot-token /root/.cargo /root/.rustup && \
+    apt-get remove -y build-essential && \
     apt-get autoremove -y && \
     rm -rf /var/lib/apt/lists/*
 

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

@@ -1 +1,9 @@
+# Build output
 *.so
+
+# Rust build artifacts
+target/
+Cargo.lock
+
+# C build artifacts (legacy)
+*.o

containers/agent/one-shot-token/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "one-shot-token"
+version = "0.1.0"
+edition = "2021"
+description = "LD_PRELOAD library for one-shot access to sensitive environment variables"
+license = "MIT"
+
+[lib]
+name = "one_shot_token"
+crate-type = ["cdylib"]
+
+[dependencies]
+libc = "0.2"
+once_cell = "1.19"
+
+[profile.release]
+opt-level = 2
+lto = true
+strip = true

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

@@ -158,22 +158,22 @@ In chroot mode, the library must be accessible from within the chroot (host file
 
 ### In Docker (automatic)
 
-The Dockerfile compiles the library during image build:
+The Dockerfile compiles the Rust library during image build:
 
 ```dockerfile
-RUN gcc -shared -fPIC -O2 -Wall \
-    -o /usr/local/lib/one-shot-token.so \
-    /tmp/one-shot-token.c \
-    -ldl -lpthread
+RUN cargo build --release && \
+    cp target/release/libone_shot_token.so /usr/local/lib/one-shot-token.so
 ```
 
 ### Locally (for testing)
 
+Requires Rust toolchain (install via [rustup](https://rustup.rs/)):
+
 ```bash
 ./build.sh
 ```
 
-This produces `one-shot-token.so` in the current directory.
+This builds `target/release/libone_shot_token.so` and creates a symlink `one-shot-token.so` for backwards compatibility.
 
 ## Testing
 
@@ -285,12 +285,14 @@ This library is one layer in AWF's security model:
 
 ## Limitations
 
-- **x86_64 Linux only**: The library is compiled for x86_64 Ubuntu
+- **Linux only**: The library is compiled for Linux (x86_64 and potentially other architectures via Rust cross-compilation)
 - **glibc programs only**: Programs using musl libc or statically linked programs are not affected
 - **Single process**: Child processes inherit the LD_PRELOAD but have their own token state and cache (each starts fresh)
 
 ## Files
 
-- `one-shot-token.c` - Library source code
+- `src/lib.rs` - Library source code (Rust)
+- `Cargo.toml` - Rust package configuration
 - `build.sh` - Local build script
+- `one-shot-token.c` - Legacy C implementation (for reference)
 - `README.md` - This documentation

containers/agent/one-shot-token/build.sh

@@ -1,34 +1,39 @@
 #!/bin/bash
 # Build the one-shot-token LD_PRELOAD library
-# This script compiles the shared library for x86_64 Ubuntu
+# This script compiles the Rust shared library
 
 set -e
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SOURCE_FILE="${SCRIPT_DIR}/one-shot-token.c"
-OUTPUT_FILE="${SCRIPT_DIR}/one-shot-token.so"
+LINK_FILE="${SCRIPT_DIR}/one-shot-token.so"
 
-echo "[build] Compiling one-shot-token.so..."
+echo "[build] Building one-shot-token with Cargo..."
 
-# Compile as a shared library with position-independent code
-# -shared: create a shared library
-# -fPIC: position-independent code (required for shared libs)
-# -ldl: link with libdl for dlsym
-# -lpthread: link with pthread for mutex
-# -O2: optimize for performance
-# -Wall -Wextra: enable warnings
-gcc -shared -fPIC \
-    -O2 -Wall -Wextra \
-    -o "${OUTPUT_FILE}" \
-    "${SOURCE_FILE}" \
-    -ldl -lpthread
+cd "${SCRIPT_DIR}"
 
-echo "[build] Successfully built: ${OUTPUT_FILE}"
+# Build the release version
+cargo build --release
+
+# Determine the output file based on platform
+if [[ "$(uname)" == "Darwin" ]]; then
+    OUTPUT_FILE="${SCRIPT_DIR}/target/release/libone_shot_token.dylib"
+    echo "[build] Successfully built: ${OUTPUT_FILE} (macOS)"
+else
+    OUTPUT_FILE="${SCRIPT_DIR}/target/release/libone_shot_token.so"
+    echo "[build] Successfully built: ${OUTPUT_FILE}"
+
+    # Create symlink for backwards compatibility (Linux only)
+    if [[ -L "${LINK_FILE}" ]]; then
+        rm "${LINK_FILE}"
+    fi
+    ln -sf "target/release/libone_shot_token.so" "${LINK_FILE}"
+    echo "[build] Created symlink: ${LINK_FILE} -> target/release/libone_shot_token.so"
+fi
 
 # Verify it's a valid shared library
-if file "${OUTPUT_FILE}" | grep -q "shared object"; then
-    echo "[build] Verified: valid shared object"
+if file "${OUTPUT_FILE}" | grep -qE "shared object|dynamically linked"; then
+    echo "[build] Verified: valid shared library"
 else
-    echo "[build] ERROR: Output is not a valid shared object"
+    echo "[build] ERROR: Output is not a valid shared library"
     exit 1
 fi