pub mod profiles;
pub mod pty;
#[cfg(test)]
mod tests;
use crate::serve::ServeState;
use crate::serve::error::{ApiError, ApiResult};
use axum::extract::State;
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
use axum::http::{HeaderMap, StatusCode};
use axum::response::{IntoResponse, Response};
use pty::{Pty, Spawn};
use serde::Deserialize;
use std::io::{Read, Write};
use std::path::PathBuf;
use std::sync::Arc;
#[derive(Debug, Clone)]
pub struct TerminalConfig {
pub token: String,
}
impl TerminalConfig {
pub fn new() -> TerminalConfig {
let mut bytes = [0u8; 32];
getrandom::fill(&mut bytes).expect("the system random source");
TerminalConfig {
token: bytes.iter().map(|b| format!("{b:02x}")).collect(),
}
}
}
#[derive(Debug, Deserialize)]
pub struct Open {
pub token: String,
#[serde(default)]
pub cols: u16,
#[serde(default)]
pub rows: u16,
#[serde(default)]
pub command: String,
}
#[derive(Debug, Deserialize)]
#[serde(tag = "type", rename_all = "kebab-case")]
pub enum Control {
Resize { cols: u16, rows: u16 },
}
pub fn admit(
configured: Option<&TerminalConfig>,
origin: Option<&str>,
expected_origins: &[String],
token: &str,
) -> Result<(), String> {
let Some(cfg) = configured else {
return Err(
"this store is served without a terminal. Start it with `cyberbrain serve \
--terminal` if you want one."
.to_string(),
);
};
if let Some(origin) = origin
&& !expected_origins.iter().any(|o| o == origin)
{
return Err(format!("a terminal cannot be opened from {origin}"));
}
if !constant_time_eq(cfg.token.as_bytes(), token.as_bytes()) {
return Err("that is not this session's terminal token".to_string());
}
Ok(())
}
fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
if a.len() != b.len() {
return false;
}
a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
}
const TOKEN_HEADER: &str = "x-cyberbrain-terminal-token";
fn token_of(headers: &HeaderMap) -> String {
headers
.get(TOKEN_HEADER)
.and_then(|v| v.to_str().ok())
.unwrap_or_default()
.to_string()
}
pub async fn list_profiles(
State(st): State<Arc<ServeState>>,
headers: HeaderMap,
) -> ApiResult<axum::Json<serde_json::Value>> {
guard(&st, &headers)?;
let saved = profiles::load().map_err(|e| ApiError::internal(e.to_string()))?;
Ok(axum::Json(serde_json::json!({
"path": profiles::path().map(|p| cyberbrain_core::Slash(&p).to_string()),
"profiles": saved.profiles,
})))
}
pub async fn put_profiles(
State(st): State<Arc<ServeState>>,
headers: HeaderMap,
axum::Json(saved): axum::Json<profiles::Profiles>,
) -> ApiResult<axum::Json<serde_json::Value>> {
guard(&st, &headers)?;
let path = profiles::save(&saved).map_err(|e| ApiError::bad_request(e.to_string()))?;
Ok(axum::Json(serde_json::json!({
"path": cyberbrain_core::Slash(&path).to_string(),
"profiles": saved.profiles,
})))
}
fn guard(st: &ServeState, headers: &HeaderMap) -> Result<(), ApiError> {
match st.terminal.as_ref() {
None => Err(refusal(
"this store is served without a terminal; `cyberbrain serve --terminal` enables it",
)),
Some(cfg) if constant_time_eq(cfg.token.as_bytes(), token_of(headers).as_bytes()) => Ok(()),
Some(_) => Err(refusal("that is not this session's terminal token")),
}
}
fn refusal(message: &str) -> ApiError {
ApiError::new(StatusCode::FORBIDDEN, "policy-refusal", message.to_string())
}
pub async fn open(
State(st): State<Arc<ServeState>>,
headers: HeaderMap,
ws: WebSocketUpgrade,
) -> Response {
if st.terminal.is_none() {
return (
StatusCode::NOT_FOUND,
"this store is served without a terminal; `cyberbrain serve --terminal` enables it",
)
.into_response();
}
let origin = headers
.get(axum::http::header::ORIGIN)
.and_then(|v| v.to_str().ok())
.map(str::to_string);
ws.on_upgrade(move |socket| session(st, origin, socket))
}
async fn session(st: Arc<ServeState>, origin: Option<String>, mut socket: WebSocket) {
let Some(Ok(first)) = socket.recv().await else {
return;
};
let text = match first {
Message::Text(t) => t.to_string(),
_ => {
let _ = say(&mut socket, "the first frame has to be the open message").await;
return;
}
};
let Ok(req) = serde_json::from_str::<Open>(&text) else {
let _ = say(&mut socket, "the open message did not parse").await;
return;
};
if let Err(why) = admit(
st.terminal.as_ref(),
origin.as_deref(),
&st.origins,
&req.token,
) {
let _ = say(&mut socket, &why).await;
return;
}
let cwd: PathBuf = st
.app
.root()
.parent()
.unwrap_or_else(|| st.app.root())
.to_path_buf();
let argv = match crate::serve::command::tokenise(&req.command) {
Ok(argv) => argv.unwrap_or_default(),
Err(why) => {
let _ = say(&mut socket, &why).await;
return;
}
};
let mut pty = match Pty::spawn(Spawn {
command: &argv,
cwd: &cwd,
cols: req.cols.max(1),
rows: req.rows.max(1),
}) {
Ok(p) => p,
Err(e) => {
let _ = say(&mut socket, &format!("the terminal could not start: {e}")).await;
return;
}
};
let (Ok(mut reader), Ok(writer)) = (pty.reader(), pty.writer()) else {
let _ = say(&mut socket, "the terminal could not be read or written").await;
pty.kill();
return;
};
let (out_tx, mut out_rx) = tokio::sync::mpsc::channel::<Vec<u8>>(64);
std::thread::spawn(move || {
let mut buf = [0u8; 8192];
loop {
match reader.read(&mut buf) {
Ok(0) | Err(_) => break,
Ok(n) => {
if out_tx.blocking_send(buf[..n].to_vec()).is_err() {
break;
}
}
}
}
});
let in_tx = spawn_writer(writer);
loop {
tokio::select! {
out = out_rx.recv() => match out {
Some(bytes) => {
if socket.send(Message::Binary(bytes.into())).await.is_err() {
break;
}
}
None => break,
},
incoming = socket.recv() => match incoming {
Some(Ok(Message::Binary(bytes))) => {
if in_tx.try_send(bytes.to_vec()).is_err() && in_tx.is_closed() {
break;
}
}
Some(Ok(Message::Text(t))) => {
if let Ok(Control::Resize { cols, rows }) = serde_json::from_str(&t) {
let _ = pty.resize(cols.max(1), rows.max(1));
}
}
Some(Ok(_)) => {}
Some(Err(_)) | None => break,
},
}
}
if let Some(code) = pty.exited() {
let _ = socket
.send(Message::Text(
serde_json::json!({ "type": "exit", "code": code })
.to_string()
.into(),
))
.await;
}
tokio::task::spawn_blocking(move || pty.kill());
drop(socket);
}
fn spawn_writer(mut writer: Box<dyn Write + Send>) -> tokio::sync::mpsc::Sender<Vec<u8>> {
let (tx, mut rx) = tokio::sync::mpsc::channel::<Vec<u8>>(64);
std::thread::spawn(move || {
while let Some(bytes) = rx.blocking_recv() {
if writer.write_all(&bytes).is_err() {
break;
}
let _ = writer.flush();
}
});
tx
}
async fn say(socket: &mut WebSocket, message: &str) -> Result<(), axum::Error> {
socket
.send(Message::Text(
serde_json::json!({ "type": "error", "message": message })
.to_string()
.into(),
))
.await
}