Commit bc905bad by DigHuang Committed by GitHub

feat(ide-agent): add filesystem exec and path operations support (#7191)

parent 42b661ca
......@@ -5,10 +5,10 @@ edition = "2024"
rust-version = "1.95"
[dependencies]
tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "net", "time", "sync", "fs", "io-util"] }
tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "net", "time", "sync", "fs", "io-util", "process"] }
tokio-tungstenite = { version = "0.29.0", default-features = false, features = ["handshake"] }
portable-pty = "0.9.0"
serde = { version = "1.0", features = ["derive"] }
serde = { version = "1.0.228", features = ["derive"] }
serde_json = "1.0.150"
base64 = "0.22.1"
futures-util = "0.3.32"
......
......@@ -2,11 +2,17 @@ use std::sync::Arc;
use tokio::net::TcpStream;
use crate::fs::handle_fs_session;
use crate::fs::{FsPermission, handle_fs_session};
const MAX_WS_MESSAGE_SIZE: usize = 16 * 1024 * 1024;
const MAX_WS_FRAME_SIZE: usize = 4 * 1024 * 1024;
use crate::terminal::handle_terminal_session;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Channel {
Fs,
Terminal,
}
fn ws_config() -> tokio_tungstenite::tungstenite::protocol::WebSocketConfig {
tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default()
.max_message_size(Some(MAX_WS_MESSAGE_SIZE))
......@@ -40,23 +46,26 @@ fn build_unauthorized_response(
resp
}
fn parse_channel(path: &str) -> Option<Channel> {
match path {
"/fs" => Some(Channel::Fs),
"/terminal" => Some(Channel::Terminal),
_ => None,
}
}
#[allow(clippy::result_large_err)]
pub async fn handle_connection(stream: TcpStream, expected_password: Arc<String>) {
let mut path = String::new();
let mut fs_permission = "write".to_string();
let mut channel = None;
let mut fs_permission = FsPermission::Read;
let ws_stream = match tokio_tungstenite::accept_hdr_async_with_config(
stream,
|req: &tokio_tungstenite::tungstenite::handshake::server::Request,
response: tokio_tungstenite::tungstenite::handshake::server::Response| {
path = req.uri().path().to_string();
let expected_channel = if path.ends_with("/terminal") {
"terminal"
} else if path.ends_with("/fs") {
"fs"
} else {
return Err(build_unauthorized_response("Unknown websocket path"));
};
let parsed_channel = parse_channel(req.uri().path())
.ok_or_else(|| build_unauthorized_response("Unknown websocket path"))?;
channel = Some(parsed_channel);
let token_opt = extract_token(req);
......@@ -66,13 +75,13 @@ pub async fn handle_connection(stream: TcpStream, expected_password: Arc<String>
));
}
fs_permission =
extract_query_value(req, "permission").unwrap_or_else(|| "read".to_string());
if fs_permission != "read" && fs_permission != "write" {
return Err(build_unauthorized_response("Invalid fs permission"));
}
fs_permission = match extract_query_value(req, "permission") {
Some(value) => FsPermission::parse(&value)
.ok_or_else(|| build_unauthorized_response("Invalid fs permission"))?,
None => FsPermission::Read,
};
if expected_channel == "terminal" && fs_permission != "write" {
if parsed_channel == Channel::Terminal && fs_permission != FsPermission::Write {
return Err(build_unauthorized_response(
"Terminal connection requires write permission",
));
......@@ -91,11 +100,9 @@ pub async fn handle_connection(stream: TcpStream, expected_password: Arc<String>
}
};
if path.ends_with("/terminal") {
handle_terminal_session(ws_stream).await;
} else if path.ends_with("/fs") {
handle_fs_session(ws_stream, fs_permission).await;
} else {
eprintln!("Unknown request path for websocket: {}", path);
match channel {
Some(Channel::Terminal) => handle_terminal_session(ws_stream).await,
Some(Channel::Fs) => handle_fs_session(ws_stream, fs_permission).await,
None => eprintln!("Unknown request path after websocket handshake"),
}
}
......@@ -22,7 +22,11 @@ async fn main() {
);
if !workspace.exists() {
let _ = tokio::fs::create_dir_all(workspace).await;
tokio::fs::create_dir_all(workspace)
.await
.unwrap_or_else(|err| {
panic!("Failed to create workspace root {:?}: {}", workspace, err)
});
}
let password = Arc::new(
......
......@@ -33,15 +33,21 @@ pub fn load_or_create_ide_agent_password() -> Result<String, String> {
let password_path = get_password_path();
let password_path = password_path.as_path();
match std::fs::read_to_string(password_path) {
let read_existing_password = || match std::fs::read_to_string(password_path) {
Ok(content) => {
let password = content.trim().to_string();
if !password.is_empty() {
return Ok(password);
Ok(Some(password))
} else {
Ok(None)
}
}
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
Err(err) => return Err(format!("Failed to read IDE Agent password file: {}", err)),
Err(err) if err.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(err) => Err(format!("Failed to read IDE Agent password file: {}", err)),
};
if let Some(password) = read_existing_password()? {
return Ok(password);
}
let parent = password_path
......@@ -52,14 +58,19 @@ pub fn load_or_create_ide_agent_password() -> Result<String, String> {
let password = generate_ide_agent_password();
let mut options = OpenOptions::new();
options.write(true).create(true).truncate(true);
options.write(true).create_new(true);
#[cfg(unix)]
{
options.mode(0o600);
}
let mut file = options
.open(password_path)
.map_err(|err| format!("Failed to create IDE Agent password file: {}", err))?;
let mut file = match options.open(password_path) {
Ok(file) => file,
Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => {
return read_existing_password()?
.ok_or_else(|| "IDE Agent password file exists but is empty".to_string());
}
Err(err) => return Err(format!("Failed to create IDE Agent password file: {}", err)),
};
file.write_all(format!("{}\n", password).as_bytes())
.map_err(|err| format!("Failed to write IDE Agent password file: {}", err))?;
......
......@@ -189,40 +189,41 @@ pub fn init_test_workspace() -> &'static Path {
mod tests {
use super::*;
use std::fs;
use std::path::PathBuf;
fn assert_path_ok(result: Result<PathBuf, String>, expected_suffix: &str) -> PathBuf {
let path = result.expect("path should be valid");
assert!(path.ends_with(expected_suffix));
path
}
fn assert_path_err(result: Result<PathBuf, String>) {
assert!(result.is_err());
}
#[tokio::test]
async fn test_sanitize_path_success() {
let temp_workspace = init_test_workspace();
// 写入一个虚拟测试文件,以确保 validate_path 能正常通过
let test_file = temp_workspace.join("dummy.txt");
fs::write(&test_file, "dummy").unwrap();
let res = sanitize_path("dummy.txt").await;
assert!(res.is_ok());
let path = res.unwrap();
assert!(path.ends_with("dummy.txt"));
let res = sanitize_path("./dummy.txt").await;
assert!(res.is_ok());
let path = res.unwrap();
assert!(path.ends_with("dummy.txt"));
assert_path_ok(sanitize_path("dummy.txt").await, "dummy.txt");
assert_path_ok(sanitize_path("./dummy.txt").await, "dummy.txt");
}
#[tokio::test]
async fn test_sanitize_path_absolute_denied() {
let _temp_workspace = init_test_workspace();
let res = sanitize_path("/dummy.txt").await;
assert!(res.is_err());
assert_path_err(sanitize_path("/dummy.txt").await);
}
#[tokio::test]
async fn test_sanitize_path_traversal_denied() {
let _temp_workspace = init_test_workspace();
let res = sanitize_path("../../../etc/passwd").await;
assert!(res.is_err());
assert_path_err(sanitize_path("../../../etc/passwd").await);
}
#[tokio::test]
......@@ -231,17 +232,17 @@ mod tests {
let target = temp_workspace.join("nested_missing_parent");
let _ = fs::remove_dir_all(&target);
let res = sanitize_create_path("nested_missing_parent/a/file.txt").await;
assert!(res.is_ok());
assert!(res.unwrap().ends_with("nested_missing_parent/a/file.txt"));
assert_path_ok(
sanitize_create_path("nested_missing_parent/a/file.txt").await,
"nested_missing_parent/a/file.txt",
);
}
#[tokio::test]
async fn test_sanitize_create_path_traversal_denied() {
let _temp_workspace = init_test_workspace();
let res = sanitize_create_path("nested/../../../etc/passwd").await;
assert!(res.is_err());
assert_path_err(sanitize_create_path("nested/../../../etc/passwd").await);
}
#[cfg(unix)]
......@@ -253,9 +254,10 @@ mod tests {
let _ = fs::remove_file(&link_path);
std::os::unix::fs::symlink(outside.path(), &link_path).unwrap();
let res = sanitize_existing_workspace_entry_path("move_source_link")
.await
.unwrap();
let res = assert_path_ok(
sanitize_existing_workspace_entry_path("move_source_link").await,
"move_source_link",
);
assert_eq!(res.file_name(), link_path.file_name());
assert_ne!(res, outside.path());
......@@ -270,7 +272,6 @@ mod tests {
let _ = fs::remove_file(&link_path);
std::os::unix::fs::symlink(outside.path(), &link_path).unwrap();
let res = sanitize_create_path("outside_link_for_create/file.txt").await;
assert!(res.is_err());
assert_path_err(sanitize_create_path("outside_link_for_create/file.txt").await);
}
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment