Add native services and embeddable transport SDK

This commit is contained in:
DuProcess
2026-06-30 19:21:23 -04:00
parent 35d9c9c6de
commit bbf6ffd725
13 changed files with 4328 additions and 33 deletions
+658 -26
View File
@@ -1,4 +1,4 @@
use anyhow::{Context, Result, anyhow};
use anyhow::{Context, Result, anyhow, bail};
use clap::{Parser, Subcommand};
use dosh::auth::{
build_attach_ticket, build_bootstrap, encode_bootstrap, load_or_create_server_secret, now_secs,
@@ -6,6 +6,14 @@ use dosh::auth::{
};
use dosh::config::{ServerConfig, expand_tilde, load_server_config};
use dosh::crypto;
use dosh::exec_service::{
EXEC_STREAM_SENTINEL, ExecResponse, FrameDecoder as ExecFrameDecoder,
encode_response as encode_exec_response,
};
use dosh::file_transfer::{
CHUNK_SIZE, FILE_STREAM_SENTINEL, FileEntry, FileKind, FileMeta, FileRequest, FileResponse,
FrameDecoder, clean_remote_path, encode_response,
};
use dosh::native::{
EnvVar, ForwardingKind, ForwardingRequest, NativeAuthOk, NativeServerHello,
derive_native_session_key, generate_native_ephemeral, host_public_key, host_public_key_line,
@@ -20,16 +28,21 @@ use dosh::protocol::{
TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
};
use dosh::pty::{PtyHandle, PtyOutput, adopt_pty_from_fd, spawn_pty_session};
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::fs;
use std::io::{Read, Seek, SeekFrom, Write};
use std::net::{IpAddr, SocketAddr};
use std::os::unix::fs::PermissionsExt;
use std::os::unix::net::UnixStream as StdUnixStream;
use std::path::PathBuf;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, UdpSocket, UnixListener};
use tokio::process::Command as TokioCommand;
use tokio::sync::mpsc;
const STREAM_INITIAL_WINDOW: usize = 1024 * 1024;
@@ -1017,6 +1030,8 @@ async fn handle_native_user_auth(
allow_tcp_forwarding: locked.config.allow_tcp_forwarding,
allow_remote_forwarding: locked.config.allow_remote_forwarding,
allow_agent_forwarding: locked.config.allow_agent_forwarding,
allow_file_transfer: locked.config.allow_file_transfer,
allow_exec_command: locked.config.allow_exec_command,
policy_flags: Vec::new(),
server_version: env!("CARGO_PKG_VERSION").to_string(),
})
@@ -1801,6 +1816,88 @@ async fn handle_stream_open(
}
}
if open.target_host == FILE_STREAM_SENTINEL {
let (writer_tx, writer_rx) = mpsc::channel::<Vec<u8>>(1024);
register_opened_stream(
state,
&session_name,
packet.header.conn_id,
open.stream_id,
writer_tx,
)?;
send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?;
let state = Arc::clone(state);
let socket = Arc::clone(socket);
let client_id = packet.header.conn_id;
let stream_id = open.stream_id;
tokio::spawn(async move {
if let Err(err) = run_file_stream_service(
state.clone(),
socket.clone(),
client_id,
stream_id,
writer_rx,
)
.await
{
let _ = send_file_response_to_client(
&state,
&socket,
client_id,
stream_id,
FileResponse::Error {
message: err.to_string(),
},
)
.await;
}
let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
});
return Ok(());
}
if open.target_host == EXEC_STREAM_SENTINEL {
let (writer_tx, writer_rx) = mpsc::channel::<Vec<u8>>(1024);
register_opened_stream(
state,
&session_name,
packet.header.conn_id,
open.stream_id,
writer_tx,
)?;
send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?;
let state = Arc::clone(state);
let socket = Arc::clone(socket);
let client_id = packet.header.conn_id;
let stream_id = open.stream_id;
tokio::spawn(async move {
if let Err(err) = run_exec_stream_service(
state.clone(),
socket.clone(),
client_id,
stream_id,
writer_rx,
)
.await
{
let _ = send_exec_response_to_client(
&state,
&socket,
client_id,
stream_id,
ExecResponse::Error {
message: err.to_string(),
},
)
.await;
}
let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
});
return Ok(());
}
let target = format!("{}:{}", open.target_host, open.target_port);
let stream = match TcpStream::connect((open.target_host.as_str(), open.target_port)).await {
Ok(stream) => stream,
@@ -1817,30 +1914,13 @@ async fn handle_stream_open(
};
let (mut reader, mut writer) = stream.into_split();
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(1024);
{
let mut locked = state.lock().expect("server state poisoned");
let session = locked
.sessions
.get_mut(&session_name)
.ok_or_else(|| anyhow!("unknown session"))?;
let client = session
.clients
.get_mut(&packet.header.conn_id)
.ok_or_else(|| anyhow!("unknown client"))?;
client.stream_writers.insert(open.stream_id, writer_tx);
client.opened_streams.insert(open.stream_id);
client
.stream_send_credit
.insert(open.stream_id, STREAM_INITIAL_WINDOW);
client
.stream_next_send_offset
.entry(open.stream_id)
.or_insert(0);
client
.stream_next_recv_offset
.entry(open.stream_id)
.or_insert(0);
}
register_opened_stream(
state,
&session_name,
packet.header.conn_id,
open.stream_id,
writer_tx,
)?;
tokio::spawn(async move {
while let Some(bytes) = writer_rx.recv().await {
@@ -1884,6 +1964,32 @@ async fn handle_stream_open(
Ok(())
}
fn register_opened_stream(
state: &Arc<Mutex<ServerState>>,
session_name: &str,
client_id: [u8; 16],
stream_id: u64,
writer_tx: mpsc::Sender<Vec<u8>>,
) -> Result<()> {
let mut locked = state.lock().expect("server state poisoned");
let session = locked
.sessions
.get_mut(session_name)
.ok_or_else(|| anyhow!("unknown session"))?;
let client = session
.clients
.get_mut(&client_id)
.ok_or_else(|| anyhow!("unknown client"))?;
client.stream_writers.insert(stream_id, writer_tx);
client.opened_streams.insert(stream_id);
client
.stream_send_credit
.insert(stream_id, STREAM_INITIAL_WINDOW);
client.stream_next_send_offset.entry(stream_id).or_insert(0);
client.stream_next_recv_offset.entry(stream_id).or_insert(0);
Ok(())
}
async fn handle_stream_data(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
@@ -2320,6 +2426,26 @@ fn stream_open_allowed(
client: &ClientState,
open: &StreamOpen,
) -> Result<()> {
if open.target_host == FILE_STREAM_SENTINEL {
anyhow::ensure!(config.allow_file_transfer, "file transfer disabled");
let allowed = client.allowed_forwardings.iter().any(|forward| {
forward.kind == ForwardingKind::Local
&& forward.target_host.as_deref() == Some(FILE_STREAM_SENTINEL)
&& forward.target_port == Some(0)
});
anyhow::ensure!(allowed, "file transfer was not requested during auth");
return Ok(());
}
if open.target_host == EXEC_STREAM_SENTINEL {
anyhow::ensure!(config.allow_exec_command, "exec command disabled");
let allowed = client.allowed_forwardings.iter().any(|forward| {
forward.kind == ForwardingKind::Local
&& forward.target_host.as_deref() == Some(EXEC_STREAM_SENTINEL)
&& forward.target_port == Some(0)
});
anyhow::ensure!(allowed, "exec command was not requested during auth");
return Ok(());
}
anyhow::ensure!(config.allow_tcp_forwarding, "TCP forwarding disabled");
let allowed = client
.allowed_forwardings
@@ -2343,6 +2469,512 @@ fn stream_open_allowed(
Ok(())
}
struct UploadState {
final_path: PathBuf,
temp_path: Option<PathBuf>,
file: fs::File,
hasher: Sha256,
written: u64,
expected_size: u64,
modified_secs: Option<u64>,
atomic: bool,
}
async fn run_file_stream_service(
state: Arc<Mutex<ServerState>>,
socket: Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
mut rx: mpsc::Receiver<Vec<u8>>,
) -> Result<()> {
let home = std::env::var_os("HOME")
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from("."));
let mut decoder = FrameDecoder::default();
let mut upload: Option<UploadState> = None;
while let Some(bytes) = rx.recv().await {
for frame in decoder.push(&bytes)? {
let request = match dosh::file_transfer::decode_request(&frame) {
Ok(request) => request,
Err(err) => {
send_file_response_to_client(
&state,
&socket,
client_id,
stream_id,
FileResponse::Error {
message: err.to_string(),
},
)
.await?;
continue;
}
};
if let Err(err) = handle_file_request(
&state,
&socket,
client_id,
stream_id,
&home,
&mut upload,
request,
)
.await
{
send_file_response_to_client(
&state,
&socket,
client_id,
stream_id,
FileResponse::Error {
message: err.to_string(),
},
)
.await?;
}
}
}
Ok(())
}
async fn handle_file_request(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
home: &Path,
upload: &mut Option<UploadState>,
request: FileRequest,
) -> Result<()> {
match request {
FileRequest::Stat { path } => {
let path = clean_remote_path(&path, home)?;
let meta = file_meta(&path, home)?;
send_file_response_to_client(
state,
socket,
client_id,
stream_id,
FileResponse::Stat { meta },
)
.await
}
FileRequest::List { path } => {
let path = clean_remote_path(&path, home)?;
let mut entries = Vec::new();
for entry in fs::read_dir(&path).with_context(|| format!("list {}", path.display()))? {
let entry = entry?;
let name = entry.file_name().to_string_lossy().to_string();
entries.push(FileEntry {
name,
meta: file_meta(&entry.path(), home)?,
});
}
entries.sort_by(|a, b| a.name.cmp(&b.name));
send_file_response_to_client(
state,
socket,
client_id,
stream_id,
FileResponse::List { entries },
)
.await
}
FileRequest::Mkdir { path, mode } => {
let path = clean_remote_path(&path, home)?;
fs::create_dir_all(&path).with_context(|| format!("create {}", path.display()))?;
if let Some(mode) = mode {
set_mode(&path, mode)?;
}
send_file_response_to_client(state, socket, client_id, stream_id, FileResponse::Ok)
.await
}
FileRequest::Remove { path, recursive } => {
let path = clean_remote_path(&path, home)?;
let metadata =
fs::symlink_metadata(&path).with_context(|| format!("stat {}", path.display()))?;
if metadata.is_dir() {
anyhow::ensure!(
recursive,
"{} is a directory; pass -r to remove directories",
path.display()
);
fs::remove_dir_all(&path).with_context(|| format!("remove {}", path.display()))?;
} else {
fs::remove_file(&path).with_context(|| format!("remove {}", path.display()))?;
}
send_file_response_to_client(state, socket, client_id, stream_id, FileResponse::Ok)
.await
}
FileRequest::PutStart {
path,
size,
mode,
modified_secs,
overwrite,
resume,
} => {
if upload.is_some() {
bail!("upload already in progress");
}
let final_path = clean_remote_path(&path, home)?;
if let Some(parent) = final_path.parent() {
fs::create_dir_all(parent)
.with_context(|| format!("create {}", parent.display()))?;
}
if final_path.exists() && !overwrite && !resume {
bail!("destination exists: {}", final_path.display());
}
let mut hasher = Sha256::new();
let (file, temp_path, written, atomic) = if resume && final_path.exists() {
let mut file = fs::OpenOptions::new()
.read(true)
.append(true)
.open(&final_path)
.with_context(|| format!("open {}", final_path.display()))?;
let written = file.metadata()?.len().min(size);
if written > 0 {
let mut prefix = fs::File::open(&final_path)?;
hash_exact_prefix(&mut prefix, written, &mut hasher)?;
}
file.seek(SeekFrom::Start(written))?;
(file, None, written, false)
} else {
let temp_path = upload_temp_path(&final_path, stream_id);
let file = fs::OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.open(&temp_path)
.with_context(|| format!("create {}", temp_path.display()))?;
(file, Some(temp_path), 0, true)
};
*upload = Some(UploadState {
final_path,
temp_path,
file,
hasher,
written,
expected_size: size,
modified_secs,
atomic,
});
if let Some(mode) = mode {
let target = upload
.as_ref()
.and_then(|state| state.temp_path.as_ref())
.unwrap_or_else(|| &upload.as_ref().unwrap().final_path)
.clone();
set_mode(&target, mode)?;
}
send_file_response_to_client(
state,
socket,
client_id,
stream_id,
FileResponse::Resume { offset: written },
)
.await
}
FileRequest::PutChunk { offset, bytes } => {
let Some(active) = upload.as_mut() else {
bail!("no upload in progress");
};
anyhow::ensure!(
offset == active.written,
"upload offset mismatch: got {offset}, expected {}",
active.written
);
active.file.write_all(&bytes)?;
active.hasher.update(&bytes);
active.written = active.written.saturating_add(bytes.len() as u64);
Ok(())
}
FileRequest::PutFinish { sha256 } => {
let Some(mut active) = upload.take() else {
bail!("no upload in progress");
};
active.file.flush()?;
active.file.sync_all()?;
anyhow::ensure!(
active.written == active.expected_size,
"upload size mismatch: got {}, expected {}",
active.written,
active.expected_size
);
let digest: [u8; 32] = active.hasher.finalize().into();
if digest != sha256 {
if let Some(temp_path) = &active.temp_path {
let _ = fs::remove_file(temp_path);
}
bail!("upload checksum mismatch");
}
drop(active.file);
if active.atomic
&& let Some(temp_path) = active.temp_path.take()
{
fs::rename(&temp_path, &active.final_path).with_context(|| {
format!(
"rename {} to {}",
temp_path.display(),
active.final_path.display()
)
})?;
}
if let Some(modified_secs) = active.modified_secs {
filetime::set_file_mtime(
&active.final_path,
filetime::FileTime::from_unix_time(modified_secs as i64, 0),
)
.with_context(|| format!("set mtime {}", active.final_path.display()))?;
}
send_file_response_to_client(
state,
socket,
client_id,
stream_id,
FileResponse::Done {
sha256: digest,
bytes: active.written,
},
)
.await
}
FileRequest::Get { path, offset } => {
let path = clean_remote_path(&path, home)?;
let meta = file_meta(&path, home)?;
anyhow::ensure!(
meta.kind == FileKind::File,
"remote path is not a file: {}",
meta.path
);
let mut file =
fs::File::open(&path).with_context(|| format!("open {}", path.display()))?;
file.seek(SeekFrom::Start(offset))?;
send_file_response_to_client(
state,
socket,
client_id,
stream_id,
FileResponse::Start {
meta: meta.clone(),
offset,
},
)
.await?;
let mut hasher = Sha256::new();
if offset > 0 {
let mut prefix = fs::File::open(&path)?;
hash_exact_prefix(&mut prefix, offset, &mut hasher)?;
}
let mut sent = offset;
let mut buf = vec![0u8; CHUNK_SIZE];
loop {
let n = file.read(&mut buf)?;
if n == 0 {
break;
}
hasher.update(&buf[..n]);
send_file_response_to_client(
state,
socket,
client_id,
stream_id,
FileResponse::Chunk {
offset: sent,
bytes: buf[..n].to_vec(),
},
)
.await?;
sent = sent.saturating_add(n as u64);
}
let digest: [u8; 32] = hasher.finalize().into();
send_file_response_to_client(
state,
socket,
client_id,
stream_id,
FileResponse::Done {
sha256: digest,
bytes: sent,
},
)
.await
}
}
}
async fn send_file_response_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
response: FileResponse,
) -> Result<()> {
send_stream_data_to_client(
state,
socket,
client_id,
stream_id,
encode_response(&response)?,
)
.await
}
async fn run_exec_stream_service(
state: Arc<Mutex<ServerState>>,
socket: Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
mut rx: mpsc::Receiver<Vec<u8>>,
) -> Result<()> {
let mut decoder = ExecFrameDecoder::default();
while let Some(bytes) = rx.recv().await {
for frame in decoder.push(&bytes)? {
let request = match dosh::exec_service::decode_request(&frame) {
Ok(request) => request,
Err(err) => {
send_exec_response_to_client(
&state,
&socket,
client_id,
stream_id,
ExecResponse::Error {
message: err.to_string(),
},
)
.await?;
continue;
}
};
let output = TokioCommand::new("sh")
.arg("-lc")
.arg(&request.command)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.output()
.await
.with_context(|| format!("run command {:?}", request.command))?;
for chunk in output.stdout.chunks(CHUNK_SIZE) {
send_exec_response_to_client(
&state,
&socket,
client_id,
stream_id,
ExecResponse::Stdout(chunk.to_vec()),
)
.await?;
}
for chunk in output.stderr.chunks(CHUNK_SIZE) {
send_exec_response_to_client(
&state,
&socket,
client_id,
stream_id,
ExecResponse::Stderr(chunk.to_vec()),
)
.await?;
}
send_exec_response_to_client(
&state,
&socket,
client_id,
stream_id,
ExecResponse::Exit {
code: output.status.code(),
},
)
.await?;
}
}
Ok(())
}
async fn send_exec_response_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
response: ExecResponse,
) -> Result<()> {
send_stream_data_to_client(
state,
socket,
client_id,
stream_id,
encode_exec_response(&response)?,
)
.await
}
fn file_meta(path: &Path, home: &Path) -> Result<FileMeta> {
let metadata =
fs::symlink_metadata(path).with_context(|| format!("stat {}", path.display()))?;
let file_type = metadata.file_type();
let kind = if file_type.is_file() {
FileKind::File
} else if file_type.is_dir() {
FileKind::Directory
} else if file_type.is_symlink() {
FileKind::Symlink
} else {
FileKind::Other
};
let modified_secs = metadata
.modified()
.ok()
.and_then(|mtime| mtime.duration_since(std::time::UNIX_EPOCH).ok())
.map(|duration| duration.as_secs());
let display_path = path
.strip_prefix(home)
.ok()
.and_then(|relative| relative.to_str())
.filter(|relative| !relative.is_empty())
.map(|relative| format!("~/{relative}"))
.unwrap_or_else(|| path.display().to_string());
Ok(FileMeta {
path: display_path,
kind,
len: metadata.len(),
mode: Some(metadata.permissions().mode() & 0o7777),
modified_secs,
})
}
fn set_mode(path: &Path, mode: u32) -> Result<()> {
let mut permissions = fs::metadata(path)?.permissions();
permissions.set_mode(mode & 0o7777);
fs::set_permissions(path, permissions).with_context(|| format!("chmod {}", path.display()))
}
fn upload_temp_path(final_path: &Path, stream_id: u64) -> PathBuf {
let name = final_path
.file_name()
.and_then(|name| name.to_str())
.unwrap_or("upload");
final_path.with_file_name(format!(
".{name}.dosh-upload-{}-{stream_id}",
std::process::id()
))
}
fn hash_exact_prefix(file: &mut fs::File, bytes: u64, hasher: &mut Sha256) -> Result<()> {
file.seek(SeekFrom::Start(0))?;
let mut remaining = bytes;
let mut buf = vec![0u8; CHUNK_SIZE];
while remaining > 0 {
let want = remaining.min(buf.len() as u64) as usize;
let n = file.read(&mut buf[..want])?;
if n == 0 {
bail!("file ended while hashing resume prefix");
}
hasher.update(&buf[..n]);
remaining -= n as u64;
}
Ok(())
}
async fn send_stream_open_ok(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,