use crate::{
agents::Agent,
errors::{AtomicError, AtomicResult},
sync::protocol,
};
use futures::{SinkExt, StreamExt};
use tokio::sync::{broadcast, mpsc};
use tokio_tungstenite::{connect_async, tungstenite::Message};
#[derive(Clone, Debug)]
pub enum WsMessage {
LoroSyncUpdate { subject: String, update: Vec<u8> },
LoroEphemeralUpdate { subject: String, update: Vec<u8> },
PresenceUpdate { subject: String, update: Vec<u8> },
Authenticated,
Keepalive,
BlobResponse { hash: [u8; 32], bytes: Vec<u8> },
Update {
subject: String,
loro_bytes: Vec<u8>,
commit_id: Option<String>,
is_snapshot: bool,
is_push: bool,
},
Destroy { subject: String },
CommitOk {
request_id: u16,
commit_id: String,
commit_json: Option<String>,
},
Challenge { nonce: String },
SyncOk { drive: String },
SyncResend { drive: String },
SyncDiff {
drive: String,
pull: Vec<String>,
push: Vec<String>,
remove: Vec<String>,
},
SyncPush {
drive: String,
entries: Vec<(String, Vec<u8>)>,
last: bool,
},
BlobRequest { hash: [u8; 32] },
Error {
request_id: u16,
code: u16,
message: String,
},
Unrecognized(String),
}
pub struct WsClient {
tx: mpsc::Sender<Message>,
broadcast_tx: broadcast::Sender<WsMessage>,
origin: String,
server_capabilities: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
challenge: std::sync::Arc<std::sync::Mutex<Option<String>>>,
}
const CHALLENGE_WAIT: std::time::Duration = std::time::Duration::from_millis(300);
fn http_origin_of_ws_url(url: &str) -> String {
let Ok(parsed) = url::Url::parse(url) else {
return url.to_string();
};
let scheme = match parsed.scheme() {
"wss" | "https" => "https",
_ => "http",
};
let Some(host) = parsed.host_str() else {
return url.to_string();
};
match parsed.port() {
Some(port) => format!("{scheme}://{host}:{port}"),
None => format!("{scheme}://{host}"),
}
}
impl WsClient {
pub async fn connect(url: &str) -> AtomicResult<Self> {
let (ws_stream, _response) = connect_async(url)
.await
.map_err(|e| format!("WebSocket connection failed to {}: {}", url, e))?;
let origin = http_origin_of_ws_url(url);
let server_capabilities = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
let caps_for_reader = server_capabilities.clone();
let challenge = std::sync::Arc::new(std::sync::Mutex::new(None));
let challenge_for_reader = challenge.clone();
let (mut write, mut read) = ws_stream.split();
let (tx, mut rx) = mpsc::channel::<Message>(64);
let (broadcast_tx, _) = broadcast::channel::<WsMessage>(256);
let broadcast_tx_clone = broadcast_tx.clone();
tokio::spawn(async move {
while let Some(msg) = rx.recv().await {
if write.send(msg).await.is_err() {
break;
}
}
});
tokio::spawn(async move {
while let Some(Ok(msg)) = read.next().await {
let parsed = match msg {
Message::Text(text) => Some(parse_server_message(&text)),
Message::Binary(bin) => {
if bin.first() == Some(&protocol::tag::AUTH_OK) {
if let Ok(mut caps) = caps_for_reader.lock() {
*caps = protocol::decode_auth_ok(&bin[1..]);
}
}
if bin.first() == Some(&protocol::tag::CHALLENGE) {
if let (Some(nonce), Ok(mut slot)) = (
protocol::decode_challenge(&bin[1..]),
challenge_for_reader.lock(),
) {
*slot = Some(nonce.to_string());
}
}
parse_binary_message(&bin)
}
_ => None,
};
if let Some(parsed) = parsed {
let _ = broadcast_tx_clone.send(parsed);
}
}
});
let client = Self {
tx,
broadcast_tx,
origin,
server_capabilities,
challenge,
};
client
.send_binary(protocol::encode_hello_with_caps(
"atomic_lib WsClient",
protocol::CLIENT_CAPABILITIES,
))
.await?;
Ok(client)
}
pub fn challenge_nonce(&self) -> Option<String> {
self.challenge.lock().ok().and_then(|c| c.clone())
}
async fn await_challenge(&self) -> Option<String> {
if let Some(nonce) = self.challenge_nonce() {
return Some(nonce);
}
let mut rx = self.subscribe();
tokio::time::timeout(CHALLENGE_WAIT, async {
while let Ok(msg) = rx.recv().await {
if let WsMessage::Challenge { nonce } = msg {
return Some(nonce);
}
}
None
})
.await
.ok()
.flatten()
.or_else(|| self.challenge_nonce())
}
pub async fn auth_subject(&self) -> String {
match self.await_challenge().await {
Some(nonce) => format!("{}#{}", self.origin, nonce),
None => self.origin.clone(),
}
}
pub fn origin(&self) -> &str {
&self.origin
}
pub fn server_capabilities(&self) -> Vec<String> {
self.server_capabilities
.lock()
.map(|c| c.clone())
.unwrap_or_default()
}
pub fn subscribe(&self) -> broadcast::Receiver<WsMessage> {
self.broadcast_tx.subscribe()
}
pub async fn authenticate(&self, agent: &Agent) -> AtomicResult<()> {
let subject = self.auth_subject().await;
let frame = protocol::encode_auth(agent, &subject)?;
self.authenticate_with_frame(frame).await
}
pub async fn authenticate_with_frame(&self, frame: Vec<u8>) -> AtomicResult<()> {
let mut rx = self.subscribe();
self.send_binary(frame).await?;
let timeout = tokio::time::timeout(std::time::Duration::from_secs(5), async {
while let Ok(msg) = rx.recv().await {
match msg {
WsMessage::Authenticated => return Ok(()),
WsMessage::Error {
request_id: 0,
message,
..
} => {
return Err(AtomicError::from(format!("Auth failed: {}", message)));
}
_ => continue,
}
}
Err(AtomicError::from("WebSocket closed during authentication"))
});
timeout
.await
.map_err(|_| AtomicError::from("Authentication timed out"))?
}
pub async fn subscribe_resource(&self, subject: &str) -> AtomicResult<()> {
self.send_binary(protocol::encode_sub(subject)).await
}
pub async fn subscribe_loro_sync(&self, subject: &str) -> AtomicResult<()> {
self.send_raw(&format!(
"LORO_SYNC_SUBSCRIBE {}",
serde_json::json!({ "subject": subject })
))
.await
}
async fn send_ephemeral(&self, kind: u8, subject: &str, update: &[u8]) -> AtomicResult<()> {
self.send_binary(protocol::encode_ephemeral(kind, subject, "", update))
.await
}
pub async fn send_loro_sync_update(&self, subject: &str, update: &[u8]) -> AtomicResult<()> {
self.send_ephemeral(protocol::ephemeral_kind::DOC, subject, update)
.await
}
pub async fn send_loro_ephemeral_update(
&self,
subject: &str,
update: &[u8],
) -> AtomicResult<()> {
self.send_ephemeral(protocol::ephemeral_kind::LORO, subject, update)
.await
}
pub async fn subscribe_presence(&self, drive: &str) -> AtomicResult<()> {
self.send_raw(&format!(
"PRESENCE_SUBSCRIBE {}",
serde_json::json!({ "subject": drive })
))
.await
}
pub async fn unsubscribe_presence(&self, drive: &str) -> AtomicResult<()> {
self.send_raw(&format!(
"PRESENCE_UNSUBSCRIBE {}",
serde_json::json!({ "subject": drive })
))
.await
}
pub async fn send_presence_update(&self, drive: &str, update: &[u8]) -> AtomicResult<()> {
self.send_ephemeral(protocol::ephemeral_kind::PRESENCE, drive, update)
.await
}
pub async fn fetch_blob(&self, hash: &[u8; 32]) -> AtomicResult<Vec<u8>> {
let mut rx = self.subscribe();
self.send_binary(protocol::encode_blob_request(hash))
.await?;
let timeout = tokio::time::timeout(std::time::Duration::from_secs(10), async {
while let Ok(msg) = rx.recv().await {
match msg {
WsMessage::BlobResponse {
hash: rcv_hash,
bytes,
} if rcv_hash == *hash => return Ok(bytes),
WsMessage::Error { message, .. } => {
return Err(AtomicError::from(format!("Blob fetch error: {}", message)));
}
_ => continue,
}
}
Err(AtomicError::from("WebSocket closed during blob fetch"))
});
timeout
.await
.map_err(|_| AtomicError::from("Timeout fetching blob"))?
}
pub async fn send_raw(&self, msg: &str) -> AtomicResult<()> {
self.tx
.send(Message::Text(msg.to_string().into()))
.await
.map_err(|e| format!("Failed to send WebSocket message: {}", e).into())
}
pub async fn send_binary(&self, bytes: Vec<u8>) -> AtomicResult<()> {
self.tx
.send(Message::Binary(bytes.into()))
.await
.map_err(|e| format!("Failed to send WebSocket binary: {}", e).into())
}
pub async fn subscribe_drive(&self, drive_subject: &str) -> AtomicResult<()> {
self.send_binary(protocol::encode_sub(drive_subject)).await
}
pub async fn unsubscribe_drive(&self, drive_subject: &str) -> AtomicResult<()> {
self.send_binary(protocol::encode_unsub(drive_subject))
.await
}
pub async fn send_keepalive(&self) -> AtomicResult<()> {
self.send_binary(protocol::encode_keepalive()).await
}
pub async fn post_commit(&self, request_id: u16, commit_json: &str) -> AtomicResult<String> {
let mut rx = self.subscribe();
self.send_binary(protocol::encode_commit(request_id, commit_json))
.await?;
let timeout = tokio::time::timeout(std::time::Duration::from_secs(30), async {
while let Ok(msg) = rx.recv().await {
match msg {
WsMessage::CommitOk {
request_id: rid,
commit_id,
..
} if rid == request_id => return Ok(commit_id),
WsMessage::Error {
request_id: rid,
code,
message,
} if rid == request_id => {
return Err(AtomicError::from(format!(
"COMMIT failed (code {code}): {message}"
)));
}
_ => continue,
}
}
Err(AtomicError::from(
"WebSocket closed while waiting for COMMIT_OK",
))
});
timeout
.await
.map_err(|_| AtomicError::from("COMMIT timed out"))?
}
}
fn parse_server_message(text: &str) -> WsMessage {
WsMessage::Unrecognized(text.to_string())
}
fn parse_binary_message(bin: &[u8]) -> Option<WsMessage> {
use protocol::tag;
let tag = *bin.first()?;
match tag {
tag::AUTH_OK => Some(WsMessage::Authenticated),
tag::KEEPALIVE => Some(WsMessage::Keepalive),
tag::ERROR => Some(match protocol::decode_error(&bin[1..]) {
Some(e) => WsMessage::Error {
request_id: e.request_id,
code: e.code,
message: e.message,
},
None => WsMessage::Error {
request_id: 0,
code: protocol::error_code::UNKNOWN,
message: "Malformed ERROR frame".into(),
},
}),
tag::BLOB_RESPONSE => {
let resp = protocol::decode_blob_response(&bin[1..])?;
Some(WsMessage::BlobResponse {
hash: resp.hash,
bytes: resp.bytes,
})
}
tag::UPDATE => decode_update_frame(&bin[1..]),
tag::DESTROY => {
if bin.len() < 3 {
return None;
}
let subject = std::str::from_utf8(&bin[3..]).ok()?.to_string();
Some(WsMessage::Destroy { subject })
}
tag::COMMIT_OK => {
let decoded = protocol::decode_commit_ok(&bin[1..])?;
Some(WsMessage::CommitOk {
request_id: decoded.request_id,
commit_id: decoded.commit_id,
commit_json: decoded.commit_json,
})
}
tag::CHALLENGE => Some(WsMessage::Challenge {
nonce: protocol::decode_challenge(&bin[1..])?.to_string(),
}),
tag::SYNC_RESEND => Some(WsMessage::SyncResend {
drive: protocol::decode_sync_resend(&bin[1..])?.to_string(),
}),
tag::EPHEMERAL => {
let decoded = protocol::decode_ephemeral(&bin[1..])?;
let subject = decoded.drive;
let update = decoded.payload;
Some(match decoded.kind {
protocol::ephemeral_kind::DOC => WsMessage::LoroSyncUpdate { subject, update },
protocol::ephemeral_kind::LORO => {
WsMessage::LoroEphemeralUpdate { subject, update }
}
protocol::ephemeral_kind::PRESENCE => WsMessage::PresenceUpdate { subject, update },
_ => return None,
})
}
tag::SYNC_OK => {
let data = &bin[1..];
if data.len() < 2 {
return None;
}
let drive_len = u16::from_be_bytes([data[0], data[1]]) as usize;
let drive = std::str::from_utf8(data.get(2..2 + drive_len)?)
.ok()?
.to_string();
Some(WsMessage::SyncOk { drive })
}
tag::SYNC_DIFF => {
let diff = protocol::decode_sync_diff(&bin[1..])?;
Some(WsMessage::SyncDiff {
drive: diff.drive,
pull: diff.pull,
push: diff.push,
remove: diff.remove,
})
}
tag::SYNC_PUSH => {
let push = protocol::decode_sync_push(&bin[1..])?;
Some(WsMessage::SyncPush {
drive: push.drive,
entries: push
.entries
.into_iter()
.map(|e| (e.subject, e.loro_bytes))
.collect(),
last: push.last,
})
}
tag::BLOB_REQUEST => {
let hash = protocol::decode_blob_request(&bin[1..])?;
Some(WsMessage::BlobRequest { hash })
}
_ => None,
}
}
fn decode_update_frame(payload: &[u8]) -> Option<WsMessage> {
use protocol::flags;
let decoded = protocol::decode_update(payload)?;
Some(WsMessage::Update {
subject: decoded.subject,
loro_bytes: decoded.loro_bytes,
commit_id: decoded.commit_id,
is_snapshot: decoded.flag_bits & flags::SNAPSHOT != 0,
is_push: decoded.flag_bits & flags::PUSH != 0,
})
}