From abf094bcbc7ab2a3e1cfcd9783fcd8145927adde Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@windmill.dev>
Date: Thu, 6 Feb 2025 00:00:33 +0100
Subject: [PATCH] fix uv install for nsjail (#5222)

---
 backend/windmill-queue/src/jobs.rs            |  8 +++-
 .../nsjail/download.py.config.proto           |  8 ++++
 .../windmill-worker/src/python_executor.rs    | 43 +++++++++++--------
 3 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/backend/windmill-queue/src/jobs.rs b/backend/windmill-queue/src/jobs.rs
index 8c8b5e4957..72c0856943 100644
--- a/backend/windmill-queue/src/jobs.rs
+++ b/backend/windmill-queue/src/jobs.rs
@@ -898,7 +898,9 @@ pub async fn add_completed_job<T: Serialize + Send + Sync + ValidableJson>(
             sqlx::query_scalar!("SELECT premium FROM workspace WHERE id = $1", w_id)
                 .fetch_one(db)
                 .await
-                .map_err(|e| Error::internal_err(format!("fetching if {w_id} is premium: {e:#}")))?;
+                .map_err(|e| {
+                    Error::internal_err(format!("fetching if {w_id} is premium: {e:#}"))
+                })?;
         let _ = sqlx::query!(
                 "INSERT INTO usage (id, is_workspace, month_, usage) 
                 VALUES ($1, TRUE, EXTRACT(YEAR FROM current_date) * 12 + EXTRACT(MONTH FROM current_date), $2) 
@@ -2333,7 +2335,9 @@ pub async fn get_result_and_success_by_id_from_flow(
                 .fetch_optional(db)
                 .await?
                 .ok_or_else(|| {
-                    error::Error::internal_err(format!("Could not get success from flow job status"))
+                    error::Error::internal_err(format!(
+                        "Could not get success from flow job status"
+                    ))
                 })?
         }
     };
diff --git a/backend/windmill-worker/nsjail/download.py.config.proto b/backend/windmill-worker/nsjail/download.py.config.proto
index 6a93a2e2d3..5c2989417f 100644
--- a/backend/windmill-worker/nsjail/download.py.config.proto
+++ b/backend/windmill-worker/nsjail/download.py.config.proto
@@ -22,6 +22,7 @@ keep_caps: true
 keep_env: true
 mount_proc: true
 
+
 mount {
     src: "/bin"
     dst: "/bin"
@@ -74,6 +75,13 @@ mount {
     is_bind: true
 }
 
+mount {
+    src: "{TARGET_DIR}"
+    dst: "{TARGET_DIR}"
+    is_bind: true
+    rw: true
+}
+
 mount {
     src: "{PY_INSTALL_DIR}"
     dst: "{PY_INSTALL_DIR}"
diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs
index 8dc9ca0107..b25575e0cf 100644
--- a/backend/windmill-worker/src/python_executor.rs
+++ b/backend/windmill-worker/src/python_executor.rs
@@ -1615,6 +1615,7 @@ async fn spawn_uv_install(
     // If none, it is system python
     py_path: Option<String>,
     no_uv_install: bool,
+    worker_dir: &str,
 ) -> Result<tokio::process::Child, Error> {
     if !*DISABLE_NSJAIL {
         tracing::info!(
@@ -1649,13 +1650,30 @@ async fn spawn_uv_install(
         vars.push(("REQ", &req));
         vars.push(("TARGET", venv_p));
 
+        std::fs::create_dir_all(venv_p)?;
+        let nsjail_proto = format!("{req}.config.proto");
+        // Prepare NSJAIL
+        let _ = write_file(
+            job_dir,
+            &nsjail_proto,
+            &(if no_uv_install {
+                NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT_FALLBACK
+            } else {
+                NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT
+            })
+            .replace("{WORKER_DIR}", worker_dir)
+            .replace("{PY_INSTALL_DIR}", &PY_INSTALL_DIR)
+            .replace("{TARGET_DIR}", &venv_p)
+            .replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()),
+        )?;
+    
         let mut nsjail_cmd = Command::new(NSJAIL_PATH.as_str());
         nsjail_cmd
             .current_dir(job_dir)
             .env_clear()
             .envs(vars)
             .envs(PROXY_ENVS.clone())
-            .args(vec!["--config", "download.config.proto"])
+            .args(vec!["--config", &nsjail_proto])
             .stdout(Stdio::piped())
             .stderr(Stdio::piped());
         start_child_process(nsjail_cmd, NSJAIL_PATH.as_str()).await
@@ -1841,6 +1859,8 @@ pub async fn handle_python_reqs(
     // TODO: Remove (Deprecated)
     mut no_uv_install: bool,
 ) -> error::Result<Vec<String>> {
+    let worker_dir = worker_dir.to_string();
+
     let counter_arc = Arc::new(tokio::sync::Mutex::new(0));
     // Append logs with line like this:
     // [9/21]   +  requests==2.32.3            << (S3) |  in 57ms
@@ -1924,22 +1944,6 @@ pub async fn handle_python_reqs(
             .map(handle_ephemeral_token),
     );
 
-    // Prepare NSJAIL
-    if !*DISABLE_NSJAIL {
-        let _ = write_file(
-            job_dir,
-            "download.config.proto",
-            &(if no_uv_install {
-                NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT_FALLBACK
-            } else {
-                NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT
-            })
-            .replace("{WORKER_DIR}", &worker_dir)
-            .replace("{PY_INSTALL_DIR}", &PY_INSTALL_DIR)
-            .replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()),
-        )?;
-    };
-
     // Cached paths
     let mut req_with_penv: Vec<(String, String)> = vec![];
     // Requirements to pull (not cached)
@@ -2176,7 +2180,7 @@ pub async fn handle_python_reqs(
         let pip_indexes = pip_indexes.clone();
         let py_path = py_path.clone();
         let pids = pids.clone();
-
+        let worker_dir = worker_dir.clone();
         handles.push(task::spawn(async move {
             // permit will be dropped anyway if this thread exits at any point
             // so we dont have to drop it manually
@@ -2244,7 +2248,8 @@ pub async fn handle_python_reqs(
                 &job_dir,
                 pip_indexes,
                 py_path,
-                no_uv_install
+                no_uv_install,
+                &worker_dir
             ).await {
                 Ok(r) => r,
                 Err(e) => {