Spiegel von
https://github.com/dani-garcia/vaultwarden.git
synchronisiert 2024-11-05 02:28:00 +01:00
426 Zeilen
12 KiB
Rust
426 Zeilen
12 KiB
Rust
use std::{
|
|
net::SocketAddr,
|
|
sync::{
|
|
atomic::{AtomicBool, Ordering},
|
|
Arc,
|
|
},
|
|
time::Duration,
|
|
};
|
|
|
|
use chrono::NaiveDateTime;
|
|
use futures::{SinkExt, StreamExt};
|
|
use rmpv::Value;
|
|
use rocket::{serde::json::Json, Route};
|
|
use serde_json::Value as JsonValue;
|
|
use tokio::{
|
|
net::{TcpListener, TcpStream},
|
|
sync::mpsc::Sender,
|
|
};
|
|
use tokio_tungstenite::{
|
|
accept_hdr_async,
|
|
tungstenite::{handshake, Message},
|
|
};
|
|
|
|
use crate::{
|
|
api::EmptyResult,
|
|
auth::Headers,
|
|
db::models::{Cipher, Folder, Send, User},
|
|
Error, CONFIG,
|
|
};
|
|
|
|
pub fn routes() -> Vec<Route> {
|
|
routes![negotiate, websockets_err]
|
|
}
|
|
|
|
#[get("/hub")]
|
|
fn websockets_err() -> EmptyResult {
|
|
static SHOW_WEBSOCKETS_MSG: AtomicBool = AtomicBool::new(true);
|
|
|
|
if CONFIG.websocket_enabled()
|
|
&& SHOW_WEBSOCKETS_MSG.compare_exchange(true, false, Ordering::Relaxed, Ordering::Relaxed).is_ok()
|
|
{
|
|
err!(
|
|
"
|
|
###########################################################
|
|
'/notifications/hub' should be proxied to the websocket server or notifications won't work.
|
|
Go to the Wiki for more info, or disable WebSockets setting WEBSOCKET_ENABLED=false.
|
|
###########################################################################################\n"
|
|
)
|
|
} else {
|
|
Err(Error::empty())
|
|
}
|
|
}
|
|
|
|
#[post("/hub/negotiate")]
|
|
fn negotiate(_headers: Headers) -> Json<JsonValue> {
|
|
use crate::crypto;
|
|
use data_encoding::BASE64URL;
|
|
|
|
let conn_id = BASE64URL.encode(&crypto::get_random(vec![0u8; 16]));
|
|
let mut available_transports: Vec<JsonValue> = Vec::new();
|
|
|
|
if CONFIG.websocket_enabled() {
|
|
available_transports.push(json!({"transport":"WebSockets", "transferFormats":["Text","Binary"]}));
|
|
}
|
|
|
|
// TODO: Implement transports
|
|
// Rocket WS support: https://github.com/SergioBenitez/Rocket/issues/90
|
|
// Rocket SSE support: https://github.com/SergioBenitez/Rocket/issues/33
|
|
// {"transport":"ServerSentEvents", "transferFormats":["Text"]},
|
|
// {"transport":"LongPolling", "transferFormats":["Text","Binary"]}
|
|
Json(json!({
|
|
"connectionId": conn_id,
|
|
"availableTransports": available_transports
|
|
}))
|
|
}
|
|
|
|
//
|
|
// Websockets server
|
|
//
|
|
|
|
fn serialize(val: Value) -> Vec<u8> {
|
|
use rmpv::encode::write_value;
|
|
|
|
let mut buf = Vec::new();
|
|
write_value(&mut buf, &val).expect("Error encoding MsgPack");
|
|
|
|
// Add size bytes at the start
|
|
// Extracted from BinaryMessageFormat.js
|
|
let mut size: usize = buf.len();
|
|
let mut len_buf: Vec<u8> = Vec::new();
|
|
|
|
loop {
|
|
let mut size_part = size & 0x7f;
|
|
size >>= 7;
|
|
|
|
if size > 0 {
|
|
size_part |= 0x80;
|
|
}
|
|
|
|
len_buf.push(size_part as u8);
|
|
|
|
if size == 0 {
|
|
break;
|
|
}
|
|
}
|
|
|
|
len_buf.append(&mut buf);
|
|
len_buf
|
|
}
|
|
|
|
fn serialize_date(date: NaiveDateTime) -> Value {
|
|
let seconds: i64 = date.timestamp();
|
|
let nanos: i64 = date.timestamp_subsec_nanos().into();
|
|
let timestamp = nanos << 34 | seconds;
|
|
|
|
let bs = timestamp.to_be_bytes();
|
|
|
|
// -1 is Timestamp
|
|
// https://github.com/msgpack/msgpack/blob/master/spec.md#timestamp-extension-type
|
|
Value::Ext(-1, bs.to_vec())
|
|
}
|
|
|
|
fn convert_option<T: Into<Value>>(option: Option<T>) -> Value {
|
|
match option {
|
|
Some(a) => a.into(),
|
|
None => Value::Nil,
|
|
}
|
|
}
|
|
|
|
const RECORD_SEPARATOR: u8 = 0x1e;
|
|
const INITIAL_RESPONSE: [u8; 3] = [0x7b, 0x7d, RECORD_SEPARATOR]; // {, }, <RS>
|
|
|
|
#[derive(Deserialize, Copy, Clone, Eq, PartialEq)]
|
|
struct InitialMessage<'a> {
|
|
protocol: &'a str,
|
|
version: i32,
|
|
}
|
|
|
|
static INITIAL_MESSAGE: InitialMessage<'static> = InitialMessage {
|
|
protocol: "messagepack",
|
|
version: 1,
|
|
};
|
|
|
|
// We attach the UUID to the sender so we can differentiate them when we need to remove them from the Vec
|
|
type UserSenders = (uuid::Uuid, Sender<Message>);
|
|
#[derive(Clone)]
|
|
pub struct WebSocketUsers {
|
|
map: Arc<dashmap::DashMap<String, Vec<UserSenders>>>,
|
|
}
|
|
|
|
impl WebSocketUsers {
|
|
async fn send_update(&self, user_uuid: &str, data: &[u8]) {
|
|
if let Some(user) = self.map.get(user_uuid).map(|v| v.clone()) {
|
|
for (_, sender) in user.iter() {
|
|
if sender.send(Message::binary(data)).await.is_err() {
|
|
// TODO: Delete from map here too?
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// NOTE: The last modified date needs to be updated before calling these methods
|
|
pub async fn send_user_update(&self, ut: UpdateType, user: &User) {
|
|
let data = create_update(
|
|
vec![("UserId".into(), user.uuid.clone().into()), ("Date".into(), serialize_date(user.updated_at))],
|
|
ut,
|
|
);
|
|
|
|
self.send_update(&user.uuid, &data).await;
|
|
}
|
|
|
|
pub async fn send_folder_update(&self, ut: UpdateType, folder: &Folder) {
|
|
let data = create_update(
|
|
vec![
|
|
("Id".into(), folder.uuid.clone().into()),
|
|
("UserId".into(), folder.user_uuid.clone().into()),
|
|
("RevisionDate".into(), serialize_date(folder.updated_at)),
|
|
],
|
|
ut,
|
|
);
|
|
|
|
self.send_update(&folder.user_uuid, &data).await;
|
|
}
|
|
|
|
pub async fn send_cipher_update(&self, ut: UpdateType, cipher: &Cipher, user_uuids: &[String]) {
|
|
let user_uuid = convert_option(cipher.user_uuid.clone());
|
|
let org_uuid = convert_option(cipher.organization_uuid.clone());
|
|
|
|
let data = create_update(
|
|
vec![
|
|
("Id".into(), cipher.uuid.clone().into()),
|
|
("UserId".into(), user_uuid),
|
|
("OrganizationId".into(), org_uuid),
|
|
("CollectionIds".into(), Value::Nil),
|
|
("RevisionDate".into(), serialize_date(cipher.updated_at)),
|
|
],
|
|
ut,
|
|
);
|
|
|
|
for uuid in user_uuids {
|
|
self.send_update(uuid, &data).await;
|
|
}
|
|
}
|
|
|
|
pub async fn send_send_update(&self, ut: UpdateType, send: &Send, user_uuids: &[String]) {
|
|
let user_uuid = convert_option(send.user_uuid.clone());
|
|
|
|
let data = create_update(
|
|
vec![
|
|
("Id".into(), send.uuid.clone().into()),
|
|
("UserId".into(), user_uuid),
|
|
("RevisionDate".into(), serialize_date(send.revision_date)),
|
|
],
|
|
ut,
|
|
);
|
|
|
|
for uuid in user_uuids {
|
|
self.send_update(uuid, &data).await;
|
|
}
|
|
}
|
|
}
|
|
|
|
/* Message Structure
|
|
[
|
|
1, // MessageType.Invocation
|
|
{}, // Headers (map)
|
|
null, // InvocationId
|
|
"ReceiveMessage", // Target
|
|
[ // Arguments
|
|
{
|
|
"ContextId": "app_id",
|
|
"Type": ut as i32,
|
|
"Payload": {}
|
|
}
|
|
]
|
|
]
|
|
*/
|
|
fn create_update(payload: Vec<(Value, Value)>, ut: UpdateType) -> Vec<u8> {
|
|
use rmpv::Value as V;
|
|
|
|
let value = V::Array(vec![
|
|
1.into(),
|
|
V::Map(vec![]),
|
|
V::Nil,
|
|
"ReceiveMessage".into(),
|
|
V::Array(vec![V::Map(vec![
|
|
("ContextId".into(), "app_id".into()),
|
|
("Type".into(), (ut as i32).into()),
|
|
("Payload".into(), payload.into()),
|
|
])]),
|
|
]);
|
|
|
|
serialize(value)
|
|
}
|
|
|
|
fn create_ping() -> Vec<u8> {
|
|
serialize(Value::Array(vec![6.into()]))
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
#[derive(Eq, PartialEq)]
|
|
pub enum UpdateType {
|
|
CipherUpdate = 0,
|
|
CipherCreate = 1,
|
|
LoginDelete = 2,
|
|
FolderDelete = 3,
|
|
Ciphers = 4,
|
|
|
|
Vault = 5,
|
|
OrgKeys = 6,
|
|
FolderCreate = 7,
|
|
FolderUpdate = 8,
|
|
CipherDelete = 9,
|
|
SyncSettings = 10,
|
|
|
|
LogOut = 11,
|
|
|
|
SyncSendCreate = 12,
|
|
SyncSendUpdate = 13,
|
|
SyncSendDelete = 14,
|
|
|
|
None = 100,
|
|
}
|
|
|
|
pub type Notify<'a> = &'a rocket::State<WebSocketUsers>;
|
|
|
|
pub fn start_notification_server() -> WebSocketUsers {
|
|
let users = WebSocketUsers {
|
|
map: Arc::new(dashmap::DashMap::new()),
|
|
};
|
|
|
|
if CONFIG.websocket_enabled() {
|
|
let users2 = users.clone();
|
|
tokio::spawn(async move {
|
|
let addr = (CONFIG.websocket_address(), CONFIG.websocket_port());
|
|
info!("Starting WebSockets server on {}:{}", addr.0, addr.1);
|
|
let listener = TcpListener::bind(addr).await.expect("Can't listen on websocket port");
|
|
|
|
let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>();
|
|
CONFIG.set_ws_shutdown_handle(shutdown_tx);
|
|
|
|
loop {
|
|
tokio::select! {
|
|
Ok((stream, addr)) = listener.accept() => {
|
|
tokio::spawn(handle_connection(stream, users2.clone(), addr));
|
|
}
|
|
|
|
_ = &mut shutdown_rx => {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
info!("Shutting down WebSockets server!")
|
|
});
|
|
}
|
|
|
|
users
|
|
}
|
|
|
|
async fn handle_connection(stream: TcpStream, users: WebSocketUsers, addr: SocketAddr) -> Result<(), Error> {
|
|
let mut user_uuid: Option<String> = None;
|
|
|
|
info!("Accepting WS connection from {addr}");
|
|
|
|
// Accept connection, do initial handshake, validate auth token and get the user ID
|
|
use handshake::server::{Request, Response};
|
|
let mut stream = accept_hdr_async(stream, |req: &Request, res: Response| {
|
|
if let Some(token) = get_request_token(req) {
|
|
if let Ok(claims) = crate::auth::decode_login(&token) {
|
|
user_uuid = Some(claims.sub);
|
|
return Ok(res);
|
|
}
|
|
}
|
|
Err(Response::builder().status(401).body(None).unwrap())
|
|
})
|
|
.await?;
|
|
|
|
let user_uuid = user_uuid.expect("User UUID should be set after the handshake");
|
|
|
|
// Add a channel to send messages to this client to the map
|
|
let entry_uuid = uuid::Uuid::new_v4();
|
|
let (tx, mut rx) = tokio::sync::mpsc::channel(100);
|
|
users.map.entry(user_uuid.clone()).or_default().push((entry_uuid, tx));
|
|
|
|
let mut interval = tokio::time::interval(Duration::from_secs(15));
|
|
loop {
|
|
tokio::select! {
|
|
res = stream.next() => {
|
|
match res {
|
|
Some(Ok(message)) => {
|
|
// Respond to any pings
|
|
if let Message::Ping(ping) = message {
|
|
if stream.send(Message::Pong(ping)).await.is_err() {
|
|
break;
|
|
}
|
|
continue;
|
|
} else if let Message::Pong(_) = message {
|
|
/* Ignored */
|
|
continue;
|
|
}
|
|
|
|
// We should receive an initial message with the protocol and version, and we will reply to it
|
|
if let Message::Text(ref message) = message {
|
|
let msg = message.strip_suffix(RECORD_SEPARATOR as char).unwrap_or(message);
|
|
|
|
if serde_json::from_str(msg).ok() == Some(INITIAL_MESSAGE) {
|
|
stream.send(Message::binary(INITIAL_RESPONSE)).await?;
|
|
continue;
|
|
}
|
|
}
|
|
|
|
// Just echo anything else the client sends
|
|
if stream.send(message).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
_ => break,
|
|
}
|
|
}
|
|
|
|
res = rx.recv() => {
|
|
match res {
|
|
Some(res) => {
|
|
if stream.send(res).await.is_err() {
|
|
break;
|
|
}
|
|
},
|
|
None => break,
|
|
}
|
|
}
|
|
|
|
_= interval.tick() => {
|
|
if stream.send(Message::Ping(create_ping())).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
info!("Closing WS connection from {addr}");
|
|
|
|
// Delete from map
|
|
users.map.entry(user_uuid).or_default().retain(|(uuid, _)| uuid != &entry_uuid);
|
|
Ok(())
|
|
}
|
|
|
|
fn get_request_token(req: &handshake::server::Request) -> Option<String> {
|
|
const ACCESS_TOKEN_KEY: &str = "access_token=";
|
|
|
|
if let Some(Ok(auth)) = req.headers().get("Authorization").map(|a| a.to_str()) {
|
|
if let Some(token_part) = auth.strip_prefix("Bearer ") {
|
|
return Some(token_part.to_owned());
|
|
}
|
|
}
|
|
|
|
if let Some(params) = req.uri().query() {
|
|
let params_iter = params.split('&').take(1);
|
|
for val in params_iter {
|
|
if let Some(stripped) = val.strip_prefix(ACCESS_TOKEN_KEY) {
|
|
return Some(stripped.to_owned());
|
|
}
|
|
}
|
|
}
|
|
None
|
|
}
|