Fix 4 critical issues from PR #2 (per-user credentials)

- Shell injection: sanitize API key values with shlex.quote() before
  rendering into SLURM template
- Async blocking: wrap sync run_verification_suite() with
  asyncio.to_thread() in async route handler
- DB migration: add _migrate() to handle existing databases missing
  the encrypted_keys column
- Missing dependency: add cryptography>=41.0 to webapp optional deps
- Add tests for shell injection prevention and DB migration

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

Author: JihaoXin

ark/webapp/db.py

@@ -6,6 +6,7 @@
 from datetime import datetime
 from typing import Optional
 
+from sqlalchemy import text
 from sqlmodel import Field, Session, SQLModel, create_engine, delete, select
 
 
@@ -56,11 +57,22 @@ class Feedback(SQLModel, table=True):
 _engine = None
 
 
+def _migrate(engine):
+    """Run lightweight schema migrations for SQLite."""
+    with engine.connect() as conn:
+        rows = conn.execute(text("PRAGMA table_info(user)")).fetchall()
+        existing = {row[1] for row in rows}
+        if "encrypted_keys" not in existing:
+            conn.execute(text("ALTER TABLE user ADD COLUMN encrypted_keys TEXT"))
+            conn.commit()
+
+
 def get_engine(db_path: str):
     global _engine
     if _engine is None:
         _engine = create_engine(f"sqlite:///{db_path}", echo=False)
         SQLModel.metadata.create_all(_engine)
+        _migrate(_engine)
     return _engine
 
 

ark/webapp/jobs.py

@@ -4,6 +4,7 @@
 
 import os
 import re
+import shlex
 import shutil
 import signal
 import subprocess
@@ -98,6 +99,8 @@ def submit_job(
     partition = settings.slurm_partition or _auto_partition()
     account = settings.slurm_account or _auto_account()
 
+    safe_api_keys = {k: shlex.quote(v) for k, v in (api_keys or {}).items()}
+
     template_text = _SLURM_TEMPLATE.read_text()
     script = Template(template_text).render(
         project_id=project_id,
@@ -110,7 +113,7 @@ def submit_job(
         gres=settings.slurm_gres,
         cpus_per_task=settings.slurm_cpus_per_task,
         conda_env=settings.slurm_conda_env,
-        api_keys=api_keys or {},
+        api_keys=safe_api_keys,
     )
 
     if api_keys:

ark/webapp/routes.py

@@ -854,7 +854,7 @@ async def api_save_user_settings(request: Request):
     # run_verification_suite will verify the updated keys.
 
     # even if they weren't all updated in this request.
-    verification_results = run_verification_suite(user.id, settings.projects_root, current_keys)
+    verification_results = await asyncio.to_thread(run_verification_suite, user.id, settings.projects_root, current_keys)
 
     # Revert failed keys
     # 1. LLM API Keys

ark/webapp/slurm_template.sh

@@ -19,10 +19,10 @@ export PATH="$HOME/.local/bin:$HOME/texlive/2025/bin/x86_64-linux:$PATH"
 conda activate {{ conda_env }}
 {% for k, v in api_keys.items() %}
 {% if k == "claude_oauth_token" %}
-export CLAUDE_CODE_OAUTH_TOKEN="{{ v }}"
+export CLAUDE_CODE_OAUTH_TOKEN={{ v }}
 {% elif k.endswith("_api_key") or k in ("gemini", "anthropic", "openai") %}
 {% set env_key = k.upper() ~ "_API_KEY" if "_api_key" not in k.lower() else k.upper() %}
-export {{ env_key }}="{{ v }}"
+export {{ env_key }}={{ v }}
 {% endif %}
 {% endfor %}
 export HOME="{{ project_dir }}"

pyproject.toml

@@ -21,6 +21,7 @@ webapp = [
     "jinja2>=3.1",
     "python-dotenv>=1.0",
     "itsdangerous>=2.1",
+    "cryptography>=41.0",
 ]
 
 

tests/test_security.py

@@ -1,5 +1,10 @@
 import pytest
 import json
+import shlex
+import sqlite3
+import tempfile
+from pathlib import Path
+
 from ark.webapp.crypto import encrypt_text, decrypt_text
 
 class TestCredentialSecurity:
@@ -15,30 +20,30 @@ def test_users(self):
     def test_encryption_isolation(self, test_users):
         """Verify that identical plaintexts result in different ciphertexts for different users."""
         test_secret = "sk-ant-1234567890abcdef"
-        
+
         user_a = test_users["user_a"]
         user_b = test_users["user_b"]
-        
+
         enc_a = encrypt_text(test_secret, user_a)
         enc_b = encrypt_text(test_secret, user_b)
-        
+
         assert enc_a != enc_b, "Ciphertexts must be different for different users (salting check)"
         assert len(enc_a) > 0
         assert len(enc_b) > 0
 
     def test_decryption_isolation(self, test_users):
         """Verify that a user cannot decrypt another user's ciphertext."""
         test_secret = "sk-ant-isolation-test"
-        
+
         user_a = test_users["user_a"]
         user_b = test_users["user_b"]
-        
+
         enc_a = encrypt_text(test_secret, user_a)
-        
+
         # Successful decryption for owner
         dec_a = decrypt_text(enc_a, user_a)
         assert dec_a == test_secret
-        
+
         # Failed decryption for other user
         dec_b = decrypt_text(enc_a, user_b)
         assert dec_b == "", "User B should not be able to decrypt User A's data"
@@ -54,3 +59,97 @@ def test_invalid_ciphertext(self, test_users):
         """Verify that invalid ciphertexts return empty strings rather than crashing."""
         user_a = test_users["user_a"]
         assert decrypt_text("not-a-valid-fernet-token", user_a) == ""
+
+
+class TestSlurmInjection:
+    """Tests that shell metacharacters in API keys are safely escaped."""
+
+    def test_shlex_quote_prevents_injection(self):
+        """Malicious key values must be escaped before rendering into SLURM template."""
+        from jinja2 import Template
+
+        template_text = Path(__file__).parent.parent / "ark" / "webapp" / "slurm_template.sh"
+        template = Template(template_text.read_text())
+
+        malicious_keys = {
+            "claude_oauth_token": '"; rm -rf / #',
+            "gemini": "$(whoami)",
+            "openai": "key with `backticks`",
+        }
+        safe_keys = {k: shlex.quote(v) for k, v in malicious_keys.items()}
+
+        rendered = template.render(
+            project_id="test",
+            project_dir="/tmp/test",
+            log_dir="/tmp/test/logs",
+            mode="paper",
+            max_iterations=2,
+            partition="batch",
+            account="",
+            gres="",
+            cpus_per_task=1,
+            conda_env="ark",
+            api_keys=safe_keys,
+        )
+
+        # shlex.quote wraps values in single quotes — verify safe versions appear
+        # and dangerous unquoted patterns do not
+        assert "CLAUDE_CODE_OAUTH_TOKEN='\"" in rendered  # single-quoted, not double-quoted
+        assert "GEMINI_API_KEY='$(whoami)'" in rendered    # $(whoami) inside single quotes = safe
+        assert "OPENAI_API_KEY='key with `backticks`'" in rendered
+
+        # The raw values must NOT appear in double quotes (which would allow expansion)
+        assert 'CLAUDE_CODE_OAUTH_TOKEN="' not in rendered
+        assert 'GEMINI_API_KEY="' not in rendered
+        assert 'OPENAI_API_KEY="' not in rendered
+
+
+class TestDbMigration:
+    """Tests that the DB migration adds the encrypted_keys column."""
+
+    def test_migrate_adds_column(self):
+        """Column is added to an existing table missing it."""
+        from ark.webapp.db import _migrate
+        from sqlalchemy import create_engine, text
+
+        with tempfile.TemporaryDirectory() as tmp:
+            db_path = Path(tmp) / "test.db"
+            conn = sqlite3.connect(str(db_path))
+            conn.execute(
+                "CREATE TABLE user (id TEXT PRIMARY KEY, email TEXT, name TEXT)"
+            )
+            conn.commit()
+            conn.close()
+
+            engine = create_engine(f"sqlite:///{db_path}", echo=False)
+            try:
+                _migrate(engine)
+
+                with engine.connect() as c:
+                    rows = c.execute(text("PRAGMA table_info(user)")).fetchall()
+                    cols = {row[1] for row in rows}
+
+                assert "encrypted_keys" in cols
+            finally:
+                engine.dispose()
+
+    def test_migrate_idempotent(self):
+        """Running migration twice does not fail."""
+        from ark.webapp.db import _migrate
+        from sqlalchemy import create_engine, text
+
+        with tempfile.TemporaryDirectory() as tmp:
+            db_path = Path(tmp) / "test.db"
+            conn = sqlite3.connect(str(db_path))
+            conn.execute(
+                "CREATE TABLE user (id TEXT PRIMARY KEY, email TEXT, encrypted_keys TEXT)"
+            )
+            conn.commit()
+            conn.close()
+
+            engine = create_engine(f"sqlite:///{db_path}", echo=False)
+            try:
+                _migrate(engine)  # should not raise
+                _migrate(engine)  # second call, still should not raise
+            finally:
+                engine.dispose()