use anyhow::Result;use collections::{HashMap, HashSet};use futures::{ channel::{mpsc, oneshot}, pin_mut, SinkExt, StreamExt,};use gpui::{AppContext, AsyncApp, Context, Entity, Task, WeakEntity};use mlua::{Lua, MultiValue, Table, UserData, UserDataMethods};use parking_lot::Mutex;use project::{search::SearchQuery, Fs, Project};use regex::Regex;use std::{ cell::RefCell, path::{Path, PathBuf}, sync::Arc,};use util::{paths::PathMatcher, ResultExt}; pub struct ScriptOutput { pub stdout: String,} struct ForegroundFn(Box<dyn FnOnce(WeakEntity<Session>, AsyncApp) + Send>); pub struct Session { project: Entity<Project>, // TODO Remove this fs_changes: Arc<Mutex<HashMap<PathBuf, Vec<u8>>>>, foreground_fns_tx: mpsc::Sender<ForegroundFn>, _invoke_foreground_fns: Task<()>,} impl Session { pub fn new(project: Entity<Project>, cx: &mut Context<Self>) -> Self { let (foreground_fns_tx, mut foreground_fns_rx) = mpsc::channel(128); Session { project, fs_changes: Arc::new(Mutex::new(HashMap::default())), foreground_fns_tx, _invoke_foreground_fns: cx.spawn(|this, cx| async move { while let Some(foreground_fn) = foreground_fns_rx.next().await { foreground_fn.0(this.clone(), cx.clone()); } }), } } /// Runs a Lua script in a sandboxed environment and returns the printed lines pub fn run_script( &mut self, script: String, cx: &mut Context<Self>, ) -> Task<Result<ScriptOutput>> { const SANDBOX_PREAMBLE: &str = include_str!("sandbox_preamble.lua"); // TODO Remove fs_changes let fs_changes = self.fs_changes.clone(); // TODO Honor all worktrees instead of the first one let root_dir = self .project .read(cx) .visible_worktrees(cx) .next() .map(|worktree| worktree.read(cx).abs_path()); let fs = self.project.read(cx).fs().clone(); let foreground_fns_tx = self.foreground_fns_tx.clone(); cx.background_spawn(async move { let lua = Lua::new(); lua.set_memory_limit(2 * 1024 * 1024 * 1024)?; // 2 GB let globals = lua.globals(); let stdout = Arc::new(Mutex::new(String::new())); globals.set( "sb_print", lua.create_function({ let stdout = stdout.clone(); move |_, args: MultiValue| Self::print(args, &stdout) })?, )?; globals.set( "search", lua.create_async_function({ let foreground_fns_tx = foreground_fns_tx.clone(); let fs = fs.clone(); move |lua, regex| { Self::search(lua, foreground_fns_tx.clone(), fs.clone(), regex) } })?, )?; globals.set( "sb_io_open", lua.create_function({ let fs_changes = fs_changes.clone(); let root_dir = root_dir.clone(); move |lua, (path_str, mode)| { Self::io_open(&lua, &fs_changes, root_dir.as_ref(), path_str, mode) } })?, )?; globals.set("user_script", script)?; lua.load(SANDBOX_PREAMBLE).exec_async().await?; // Drop Lua instance to decrement reference count. drop(lua); let stdout = Arc::try_unwrap(stdout) .expect("no more references to stdout") .into_inner(); Ok(ScriptOutput { stdout }) }) } /// Sandboxed print() function in Lua. fn print(args: MultiValue, stdout: &Mutex<String>) -> mlua::Result<()> { for (index, arg) in args.into_iter().enumerate() { // Lua's `print()` prints tab characters between each argument. if index > 0 { stdout.lock().push('\t'); } // If the argument's to_string() fails, have the whole function call fail. stdout.lock().push_str(&arg.to_string()?); } stdout.lock().push('\n'); Ok(()) } /// Sandboxed io.open() function in Lua. fn io_open( lua: &Lua, fs_changes: &Arc<Mutex<HashMap<PathBuf, Vec<u8>>>>, root_dir: Option<&Arc<Path>>, path_str: String, mode: Option<String>, ) -> mlua::Result<(Option<Table>, String)> { let root_dir = root_dir .ok_or_else(|| mlua::Error::runtime("cannot open file without a root directory"))?; let mode = mode.unwrap_or_else(|| "r".to_string()); // Parse the mode string to determine read/write permissions let read_perm = mode.contains('r'); let write_perm = mode.contains('w') || mode.contains('a') || mode.contains('+'); let append = mode.contains('a'); let truncate = mode.contains('w'); // This will be the Lua value returned from the `open` function. let file = lua.create_table()?; // Store file metadata in the file file.set("__path", path_str.clone())?; file.set("__mode", mode.clone())?; file.set("__read_perm", read_perm)?; file.set("__write_perm", write_perm)?; // Sandbox the path; it must be within root_dir let path: PathBuf = { let rust_path = Path::new(&path_str); // Get absolute path if rust_path.is_absolute() { // Check if path starts with root_dir prefix without resolving symlinks if !rust_path.starts_with(&root_dir) { return Ok(( None, format!( "Error: Absolute path {} is outside the current working directory", path_str ), )); } rust_path.to_path_buf() } else { // Make relative path absolute relative to cwd root_dir.join(rust_path) } }; // close method let close_fn = { let fs_changes = fs_changes.clone(); lua.create_function(move |_lua, file_userdata: mlua::Table| { let write_perm = file_userdata.get::<bool>("__write_perm")?; let path = file_userdata.get::<String>("__path")?; if write_perm { // When closing a writable file, record the content let content = file_userdata.get::<mlua::AnyUserData>("__content")?; let content_ref = content.borrow::<FileContent>()?; let content_vec = content_ref.0.borrow(); // Don't actually write to disk; instead, just update fs_changes. let path_buf = PathBuf::from(&path); fs_changes .lock() .insert(path_buf.clone(), content_vec.clone()); } Ok(true) })? }; file.set("close", close_fn)?; // If it's a directory, give it a custom read() and return early. if path.is_dir() { // TODO handle the case where we changed it in the in-memory fs // Create a special directory handle file.set("__is_directory", true)?; // Store directory entries let entries = match std::fs::read_dir(&path) { Ok(entries) => { let mut entry_names = Vec::new(); for entry in entries.flatten() { entry_names.push(entry.file_name().to_string_lossy().into_owned()); } entry_names } Err(e) => return Ok((None, format!("Error reading directory: {}", e))), }; // Save the list of entries file.set("__dir_entries", entries)?; file.set("__dir_position", 0usize)?; // Create a directory-specific read function let read_fn = lua.create_function(|_lua, file_userdata: mlua::Table| { let position = file_userdata.get::<usize>("__dir_position")?; let entries = file_userdata.get::<Vec<String>>("__dir_entries")?; if position >= entries.len() { return Ok(None); // No more entries } let entry = entries[position].clone(); file_userdata.set("__dir_position", position + 1)?; Ok(Some(entry)) })?; file.set("read", read_fn)?; // If we got this far, the directory was opened successfully return Ok((Some(file), String::new())); } let fs_changes_map = fs_changes.lock(); let is_in_changes = fs_changes_map.contains_key(&path); let file_exists = is_in_changes || path.exists(); let mut file_content = Vec::new(); if file_exists && !truncate { if is_in_changes { file_content = fs_changes_map.get(&path).unwrap().clone(); } else { // Try to read existing content if file exists and we're not truncating match std::fs::read(&path) { Ok(content) => file_content = content, Err(e) => return Ok((None, format!("Error reading file: {}", e))), } } } drop(fs_changes_map); // Unlock the fs_changes mutex. // If in append mode, position should be at the end let position = if append && file_exists { file_content.len() } else { 0 }; file.set("__position", position)?; file.set( "__content", lua.create_userdata(FileContent(RefCell::new(file_content)))?, )?; // Create file methods // read method let read_fn = { lua.create_function( |_lua, (file_userdata, format): (mlua::Table, Option<mlua::Value>)| { let read_perm = file_userdata.get::<bool>("__read_perm")?; if !read_perm { return Err(mlua::Error::runtime("File not open for reading")); } let content = file_userdata.get::<mlua::AnyUserData>("__content")?; let mut position = file_userdata.get::<usize>("__position")?; let content_ref = content.borrow::<FileContent>()?; let content_vec = content_ref.0.borrow(); if position >= content_vec.len() { return Ok(None); // EOF } match format { Some(mlua::Value::String(s)) => { let lossy_string = s.to_string_lossy(); let format_str: &str = lossy_string.as_ref(); // Only consider the first 2 bytes, since it's common to pass e.g. "*all" instead of "*a" match &format_str[0..2] { "*a" => { // Read entire file from current position let result = String::from_utf8_lossy(&content_vec[position..]) .to_string(); position = content_vec.len(); file_userdata.set("__position", position)?; Ok(Some(result)) } "*l" => { // Read next line let mut line = Vec::new(); let mut found_newline = false; while position < content_vec.len() { let byte = content_vec[position]; position += 1; if byte == b'\n' { found_newline = true; break; } // Skip \r in \r\n sequence but add it if it's alone if byte == b'\r' { if position < content_vec.len() && content_vec[position] == b'\n' { position += 1; found_newline = true; break; } } line.push(byte); } file_userdata.set("__position", position)?; if !found_newline && line.is_empty() && position >= content_vec.len() { return Ok(None); // EOF } let result = String::from_utf8_lossy(&line).to_string(); Ok(Some(result)) } "*n" => { // Try to parse as a number (number of bytes to read) match format_str.parse::<usize>() { Ok(n) => { let end = std::cmp::min(position + n, content_vec.len()); let bytes = &content_vec[position..end]; let result = String::from_utf8_lossy(bytes).to_string(); position = end; file_userdata.set("__position", position)?; Ok(Some(result)) } Err(_) => Err(mlua::Error::runtime(format!( "Invalid format: {}", format_str ))), } } "*L" => { // Read next line keeping the end of line let mut line = Vec::new(); while position < content_vec.len() { let byte = content_vec[position]; position += 1; line.push(byte); if byte == b'\n' { break; } // If we encounter a \r, add it and check if the next is \n if byte == b'\r' && position < content_vec.len() && content_vec[position] == b'\n' { line.push(content_vec[position]); position += 1; break; } } file_userdata.set("__position", position)?; if line.is_empty() && position >= content_vec.len() { return Ok(None); // EOF } let result = String::from_utf8_lossy(&line).to_string(); Ok(Some(result)) } _ => Err(mlua::Error::runtime(format!( "Unsupported format: {}", format_str ))), } } Some(mlua::Value::Number(n)) => { // Read n bytes let n = n as usize; let end = std::cmp::min(position + n, content_vec.len()); let bytes = &content_vec[position..end]; let result = String::from_utf8_lossy(bytes).to_string(); position = end; file_userdata.set("__position", position)?; Ok(Some(result)) } Some(_) => Err(mlua::Error::runtime("Invalid format")), None => { // Default is to read a line let mut line = Vec::new(); let mut found_newline = false; while position < content_vec.len() { let byte = content_vec[position]; position += 1; if byte == b'\n' { found_newline = true; break; } // Handle \r\n if byte == b'\r' { if position < content_vec.len() && content_vec[position] == b'\n' { position += 1; found_newline = true; break; } } line.push(byte); } file_userdata.set("__position", position)?; if !found_newline && line.is_empty() && position >= content_vec.len() { return Ok(None); // EOF } let result = String::from_utf8_lossy(&line).to_string(); Ok(Some(result)) } } }, )? }; file.set("read", read_fn)?; // write method let write_fn = { let fs_changes = fs_changes.clone(); lua.create_function(move |_lua, (file_userdata, text): (mlua::Table, String)| { let write_perm = file_userdata.get::<bool>("__write_perm")?; if !write_perm { return Err(mlua::Error::runtime("File not open for writing")); } let content = file_userdata.get::<mlua::AnyUserData>("__content")?; let position = file_userdata.get::<usize>("__position")?; let content_ref = content.borrow::<FileContent>()?; let mut content_vec = content_ref.0.borrow_mut(); let bytes = text.as_bytes(); // Ensure the vector has enough capacity if position + bytes.len() > content_vec.len() { content_vec.resize(position + bytes.len(), 0); } // Write the bytes for (i, &byte) in bytes.iter().enumerate() { content_vec[position + i] = byte; } // Update position let new_position = position + bytes.len(); file_userdata.set("__position", new_position)?; // Update fs_changes let path = file_userdata.get::<String>("__path")?; let path_buf = PathBuf::from(path);Lines 1–500 of 743. Load more
use anyhow::Result; use collections::{HashMap, HashSet}; use futures::{ channel::{mpsc, oneshot}, pin_mut, SinkExt, StreamExt, }; use gpui::{AppContext, AsyncApp, Context, Entity, Task, WeakEntity}; use mlua::{Lua, MultiValue, Table, UserData, UserDataMethods}; use parking_lot::Mutex; use project::{search::SearchQuery, Fs, Project}; use regex::Regex; use std::{ cell::RefCell, path::{Path, PathBuf}, sync::Arc, }; use util::{paths::PathMatcher, ResultExt}; pub struct ScriptOutput { pub stdout: String, } struct ForegroundFn(Box<dyn FnOnce(WeakEntity<Session>, AsyncApp) + Send>); pub struct Session { project: Entity<Project>, // TODO Remove this fs_changes: Arc<Mutex<HashMap<PathBuf, Vec<u8>>>>, foreground_fns_tx: mpsc::Sender<ForegroundFn>, _invoke_foreground_fns: Task<()>, } impl Session { pub fn new(project: Entity<Project>, cx: &mut Context<Self>) -> Self { let (foreground_fns_tx, mut foreground_fns_rx) = mpsc::channel(128); Session { project, fs_changes: Arc::new(Mutex::new(HashMap::default())), foreground_fns_tx, _invoke_foreground_fns: cx.spawn(|this, cx| async move { while let Some(foreground_fn) = foreground_fns_rx.next().await { foreground_fn.0(this.clone(), cx.clone()); } }), } } /// Runs a Lua script in a sandboxed environment and returns the printed lines pub fn run_script( &mut self, script: String, cx: &mut Context<Self>, ) -> Task<Result<ScriptOutput>> { const SANDBOX_PREAMBLE: &str = include_str!("sandbox_preamble.lua"); // TODO Remove fs_changes let fs_changes = self.fs_changes.clone(); // TODO Honor all worktrees instead of the first one let root_dir = self .project .read(cx) .visible_worktrees(cx) .next() .map(|worktree| worktree.read(cx).abs_path()); let fs = self.project.read(cx).fs().clone(); let foreground_fns_tx = self.foreground_fns_tx.clone(); cx.background_spawn(async move { let lua = Lua::new(); lua.set_memory_limit(2 * 1024 * 1024 * 1024)?; // 2 GB let globals = lua.globals(); let stdout = Arc::new(Mutex::new(String::new())); globals.set( "sb_print", lua.create_function({ let stdout = stdout.clone(); move |_, args: MultiValue| Self::print(args, &stdout) })?, )?; globals.set( "search", lua.create_async_function({ let foreground_fns_tx = foreground_fns_tx.clone(); let fs = fs.clone(); move |lua, regex| { Self::search(lua, foreground_fns_tx.clone(), fs.clone(), regex) } })?, )?; globals.set( "sb_io_open", lua.create_function({ let fs_changes = fs_changes.clone(); let root_dir = root_dir.clone(); move |lua, (path_str, mode)| { Self::io_open(&lua, &fs_changes, root_dir.as_ref(), path_str, mode) } })?, )?; globals.set("user_script", script)?; lua.load(SANDBOX_PREAMBLE).exec_async().await?; // Drop Lua instance to decrement reference count. drop(lua); let stdout = Arc::try_unwrap(stdout) .expect("no more references to stdout") .into_inner(); Ok(ScriptOutput { stdout }) }) } /// Sandboxed print() function in Lua. fn print(args: MultiValue, stdout: &Mutex<String>) -> mlua::Result<()> { for (index, arg) in args.into_iter().enumerate() { // Lua's `print()` prints tab characters between each argument. if index > 0 { stdout.lock().push('\t'); } // If the argument's to_string() fails, have the whole function call fail. stdout.lock().push_str(&arg.to_string()?); } stdout.lock().push('\n'); Ok(()) } /// Sandboxed io.open() function in Lua. fn io_open( lua: &Lua, fs_changes: &Arc<Mutex<HashMap<PathBuf, Vec<u8>>>>, root_dir: Option<&Arc<Path>>, path_str: String, mode: Option<String>, ) -> mlua::Result<(Option<Table>, String)> { let root_dir = root_dir .ok_or_else(|| mlua::Error::runtime("cannot open file without a root directory"))?; let mode = mode.unwrap_or_else(|| "r".to_string()); // Parse the mode string to determine read/write permissions let read_perm = mode.contains('r'); let write_perm = mode.contains('w') || mode.contains('a') || mode.contains('+'); let append = mode.contains('a'); let truncate = mode.contains('w'); // This will be the Lua value returned from the `open` function. let file = lua.create_table()?; // Store file metadata in the file file.set("__path", path_str.clone())?; file.set("__mode", mode.clone())?; file.set("__read_perm", read_perm)?; file.set("__write_perm", write_perm)?; // Sandbox the path; it must be within root_dir let path: PathBuf = { let rust_path = Path::new(&path_str); // Get absolute path if rust_path.is_absolute() { // Check if path starts with root_dir prefix without resolving symlinks if !rust_path.starts_with(&root_dir) { return Ok(( None, format!( "Error: Absolute path {} is outside the current working directory", path_str ), )); } rust_path.to_path_buf() } else { // Make relative path absolute relative to cwd root_dir.join(rust_path) } }; // close method let close_fn = { let fs_changes = fs_changes.clone(); lua.create_function(move |_lua, file_userdata: mlua::Table| { let write_perm = file_userdata.get::<bool>("__write_perm")?; let path = file_userdata.get::<String>("__path")?; if write_perm { // When closing a writable file, record the content let content = file_userdata.get::<mlua::AnyUserData>("__content")?; let content_ref = content.borrow::<FileContent>()?; let content_vec = content_ref.0.borrow(); // Don't actually write to disk; instead, just update fs_changes. let path_buf = PathBuf::from(&path); fs_changes .lock() .insert(path_buf.clone(), content_vec.clone()); } Ok(true) })? }; file.set("close", close_fn)?; // If it's a directory, give it a custom read() and return early. if path.is_dir() { // TODO handle the case where we changed it in the in-memory fs // Create a special directory handle file.set("__is_directory", true)?; // Store directory entries let entries = match std::fs::read_dir(&path) { Ok(entries) => { let mut entry_names = Vec::new(); for entry in entries.flatten() { entry_names.push(entry.file_name().to_string_lossy().into_owned()); } entry_names } Err(e) => return Ok((None, format!("Error reading directory: {}", e))), }; // Save the list of entries file.set("__dir_entries", entries)?; file.set("__dir_position", 0usize)?; // Create a directory-specific read function let read_fn = lua.create_function(|_lua, file_userdata: mlua::Table| { let position = file_userdata.get::<usize>("__dir_position")?; let entries = file_userdata.get::<Vec<String>>("__dir_entries")?; if position >= entries.len() { return Ok(None); // No more entries } let entry = entries[position].clone(); file_userdata.set("__dir_position", position + 1)?; Ok(Some(entry)) })?; file.set("read", read_fn)?; // If we got this far, the directory was opened successfully return Ok((Some(file), String::new())); } let fs_changes_map = fs_changes.lock(); let is_in_changes = fs_changes_map.contains_key(&path); let file_exists = is_in_changes || path.exists(); let mut file_content = Vec::new(); if file_exists && !truncate { if is_in_changes { file_content = fs_changes_map.get(&path).unwrap().clone(); } else { // Try to read existing content if file exists and we're not truncating match std::fs::read(&path) { Ok(content) => file_content = content, Err(e) => return Ok((None, format!("Error reading file: {}", e))), } } } drop(fs_changes_map); // Unlock the fs_changes mutex. // If in append mode, position should be at the end let position = if append && file_exists { file_content.len() } else { 0 }; file.set("__position", position)?; file.set( "__content", lua.create_userdata(FileContent(RefCell::new(file_content)))?, )?; // Create file methods // read method let read_fn = { lua.create_function( |_lua, (file_userdata, format): (mlua::Table, Option<mlua::Value>)| { let read_perm = file_userdata.get::<bool>("__read_perm")?; if !read_perm { return Err(mlua::Error::runtime("File not open for reading")); } let content = file_userdata.get::<mlua::AnyUserData>("__content")?; let mut position = file_userdata.get::<usize>("__position")?; let content_ref = content.borrow::<FileContent>()?; let content_vec = content_ref.0.borrow(); if position >= content_vec.len() { return Ok(None); // EOF } match format { Some(mlua::Value::String(s)) => { let lossy_string = s.to_string_lossy(); let format_str: &str = lossy_string.as_ref(); // Only consider the first 2 bytes, since it's common to pass e.g. "*all" instead of "*a" match &format_str[0..2] { "*a" => { // Read entire file from current position let result = String::from_utf8_lossy(&content_vec[position..]) .to_string(); position = content_vec.len(); file_userdata.set("__position", position)?; Ok(Some(result)) } "*l" => { // Read next line let mut line = Vec::new(); let mut found_newline = false; while position < content_vec.len() { let byte = content_vec[position]; position += 1; if byte == b'\n' { found_newline = true; break; } // Skip \r in \r\n sequence but add it if it's alone if byte == b'\r' { if position < content_vec.len() && content_vec[position] == b'\n' { position += 1; found_newline = true; break; } } line.push(byte); } file_userdata.set("__position", position)?; if !found_newline && line.is_empty() && position >= content_vec.len() { return Ok(None); // EOF } let result = String::from_utf8_lossy(&line).to_string(); Ok(Some(result)) } "*n" => { // Try to parse as a number (number of bytes to read) match format_str.parse::<usize>() { Ok(n) => { let end = std::cmp::min(position + n, content_vec.len()); let bytes = &content_vec[position..end]; let result = String::from_utf8_lossy(bytes).to_string(); position = end; file_userdata.set("__position", position)?; Ok(Some(result)) } Err(_) => Err(mlua::Error::runtime(format!( "Invalid format: {}", format_str ))), } } "*L" => { // Read next line keeping the end of line let mut line = Vec::new(); while position < content_vec.len() { let byte = content_vec[position]; position += 1; line.push(byte); if byte == b'\n' { break; } // If we encounter a \r, add it and check if the next is \n if byte == b'\r' && position < content_vec.len() && content_vec[position] == b'\n' { line.push(content_vec[position]); position += 1; break; } } file_userdata.set("__position", position)?; if line.is_empty() && position >= content_vec.len() { return Ok(None); // EOF } let result = String::from_utf8_lossy(&line).to_string(); Ok(Some(result)) } _ => Err(mlua::Error::runtime(format!( "Unsupported format: {}", format_str ))), } } Some(mlua::Value::Number(n)) => { // Read n bytes let n = n as usize; let end = std::cmp::min(position + n, content_vec.len()); let bytes = &content_vec[position..end]; let result = String::from_utf8_lossy(bytes).to_string(); position = end; file_userdata.set("__position", position)?; Ok(Some(result)) } Some(_) => Err(mlua::Error::runtime("Invalid format")), None => { // Default is to read a line let mut line = Vec::new(); let mut found_newline = false; while position < content_vec.len() { let byte = content_vec[position]; position += 1; if byte == b'\n' { found_newline = true; break; } // Handle \r\n if byte == b'\r' { if position < content_vec.len() && content_vec[position] == b'\n' { position += 1; found_newline = true; break; } } line.push(byte); } file_userdata.set("__position", position)?; if !found_newline && line.is_empty() && position >= content_vec.len() { return Ok(None); // EOF } let result = String::from_utf8_lossy(&line).to_string(); Ok(Some(result)) } } }, )? }; file.set("read", read_fn)?; // write method let write_fn = { let fs_changes = fs_changes.clone(); lua.create_function(move |_lua, (file_userdata, text): (mlua::Table, String)| { let write_perm = file_userdata.get::<bool>("__write_perm")?; if !write_perm { return Err(mlua::Error::runtime("File not open for writing")); } let content = file_userdata.get::<mlua::AnyUserData>("__content")?; let position = file_userdata.get::<usize>("__position")?; let content_ref = content.borrow::<FileContent>()?; let mut content_vec = content_ref.0.borrow_mut(); let bytes = text.as_bytes(); // Ensure the vector has enough capacity if position + bytes.len() > content_vec.len() { content_vec.resize(position + bytes.len(), 0); } // Write the bytes for (i, &byte) in bytes.iter().enumerate() { content_vec[position + i] = byte; } // Update position let new_position = position + bytes.len(); file_userdata.set("__position", new_position)?; // Update fs_changes let path = file_userdata.get::<String>("__path")?; let path_buf = PathBuf::from(path);