diff --git a/Cargo.lock b/Cargo.lock index 97bf679..e810b47 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -168,6 +168,7 @@ dependencies = [ "serde", "serde_json", "serde_path_to_error", + "serde_urlencoded", "sync_wrapper", "tokio", "tower", @@ -494,6 +495,15 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "form_urlencoded" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a62bc1cf6f830c2ec14a513a9fb124d0a213a629668a4186f329db21fe045652" +dependencies = [ + "percent-encoding", +] + [[package]] name = "futures" version = "0.3.28" @@ -1444,6 +1454,18 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + [[package]] name = "sharded-slab" version = "0.1.7" diff --git a/Cargo.toml b/Cargo.toml index 9e6f3b9..931d38c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ axum = { version = "0.6.20", default-features = false, features = [ "http1", "http2", "json", + "query", "tokio", ] } base64 = "0.21.5" diff --git a/docs/api.yaml b/docs/api.yaml index d61a43a..98fe93c 100644 --- a/docs/api.yaml +++ b/docs/api.yaml @@ -29,10 +29,33 @@ paths: get: tags: - Message - summary: Pop a message from the inbound message queue + summary: Get a message from the inbound message queue description: | - Pop a message from the inbound message queue. The message is removed from the queue and won't be shown again + Get a message from the inbound message queue. By default, the message is removed from the queue and won't be shown again. + If the peek query parameter is set to true, the message will be peeked, and the next call to this endpoint will show the same message. + This method returns immediately by default: a message is returned if one is ready, and if there isn't nothing is returned. If the timeout + query parameter is set, this call won't return for the given amount of seconds, unless a message is received operationId: popMessage + parameters: + - in: query + name: peek + required: false + schema: + type: boolean + description: Whether to peek the message or not. If this is true, the message won't be removed from the inbound queue when it is read + example: true + - in: query + name: timeout + required: false + schema: + type: integer + format: int64 + minimum: 0 + description: | + Amount of seconds to wait for a message to arrive if one is not available. Setting this to 0 is valid and will return + a message if present, or return immediately if there isn't + example: 60 + responses: '200': description: Message retrieved @@ -63,25 +86,6 @@ paths: schema: $ref: '#/components/schemas/PushMessageResponse' - '/api/v1/messages/peek': - get: - tags: - - Message - summary: Peek a message from the inbound queue - description: | - Peek a message from the inbound message queue. The message is removed from the queue and will be shown - again on the next call again to peek or pop - operationId: peekMessage - responses: - '200': - description: Message peeked successfully - content: - application/json: - schema: - $ref: '#/components/schemas/InboundMessage' - '204': - description: No message ready - '/api/v1/messages/status/:id': get: tags: diff --git a/src/api.rs b/src/api.rs index 6ba4f2b..e59c411 100644 --- a/src/api.rs +++ b/src/api.rs @@ -4,7 +4,7 @@ use std::{ }; use axum::{ - extract::{Path, State}, + extract::{Path, Query, State}, http::StatusCode, routing::get, Json, Router, @@ -74,8 +74,7 @@ impl Http { pub fn spawn(message_stack: MessageStack, listen_addr: &SocketAddr) -> Self { let server_state = HttpServerState { message_stack }; let msg_routes = Router::new() - .route("/messages", get(pop_message).post(push_message)) - .route("/messages/peek", get(peek_message)) + .route("/messages", get(get_message).post(push_message)) .route("/messages/status/:id", get(message_status)) .with_state(server_state); let app = Router::new().nest("/api/v1", msg_routes); @@ -95,44 +94,52 @@ impl Http { } } -async fn peek_message( - State(state): State, -) -> Result, StatusCode> { - debug!("Attempt to peek message"); - state - .message_stack - .peek_message() - .ok_or(StatusCode::NO_CONTENT) - .map(|m| { - Json(MessageReceiveInfo { - id: m.id, - src_ip: m.src_ip, - src_pk: m.src_pk, - dst_ip: m.dst_ip, - dst_pk: m.dst_pk, - payload: m.data, - }) - }) +#[derive(Deserialize)] +struct GetMessageQuery { + peek: Option, + timeout: Option, } -async fn pop_message( +impl GetMessageQuery { + /// Did the query indicate we should peek the message instead of pop? + fn peek(&self) -> bool { + matches!(self.peek, Some(true)) + } + + /// Amount of seconds to hold and try and get values. + fn timeout_secs(&self) -> u64 { + self.timeout.unwrap_or(0) + } +} + +async fn get_message( State(state): State, + Query(query): Query, ) -> Result, StatusCode> { - debug!("Attempt to pop message"); - state - .message_stack - .pop_message() - .ok_or(StatusCode::NO_CONTENT) - .map(|m| { - Json(MessageReceiveInfo { - id: m.id, - src_ip: m.src_ip, - src_pk: m.src_pk, - dst_ip: m.dst_ip, - dst_pk: m.dst_pk, - payload: m.data, - }) + debug!( + "Attempt to get message, peek {}, timeout {} seconds", + query.peek(), + query.timeout_secs() + ); + // A timeout of 0 seconds essentially means get a message if there is one, and return + // immediatly if there isn't. This is the result of the implementation of Timeout, which does a + // poll of the internal future first, before polling the delay. + tokio::time::timeout( + Duration::from_secs(query.timeout_secs()), + state.message_stack.message(!query.peek()), + ) + .await + .or(Err(StatusCode::NO_CONTENT)) + .map(|m| { + Json(MessageReceiveInfo { + id: m.id, + src_ip: m.src_ip, + src_pk: m.src_pk, + dst_ip: m.dst_ip, + dst_pk: m.dst_pk, + payload: m.data, }) + }) } #[derive(Serialize)] diff --git a/src/message.rs b/src/message.rs index fb0a381..ced2600 100644 --- a/src/message.rs +++ b/src/message.rs @@ -19,6 +19,7 @@ use futures::{Stream, StreamExt}; use log::{debug, error, trace, warn}; use rand::Fill; use serde::{de::Visitor, Deserialize, Deserializer, Serialize}; +use tokio::sync::watch; use crate::{ crypto::{PacketBuffer, PublicKey}, @@ -79,6 +80,8 @@ pub struct MessageStack { data_plane: Arc>, inbox: Arc>, outbox: Arc>, + /// Receiver handle for inbox listeners (basically a condvar). + subscriber: watch::Receiver<()>, } struct MessageOutbox { @@ -91,6 +94,8 @@ struct MessageInbox { pending_msges: HashMap, /// Messages which have been completed. complete_msges: VecDeque, + /// Notification sender used to allert subscribed listeners. + notify: watch::Sender<()>, } struct ReceivedMessageInfo { @@ -163,10 +168,11 @@ enum TransmissionState { } impl MessageInbox { - fn new() -> Self { + fn new(notify: watch::Sender<()>) -> Self { Self { pending_msges: HashMap::new(), complete_msges: VecDeque::new(), + notify, } } } @@ -193,10 +199,12 @@ impl MessageStack { where S: Stream + Send + Unpin + 'static, { + let (notify, subscriber) = watch::channel(()); let ms = Self { data_plane: Arc::new(Mutex::new(data_plane)), - inbox: Arc::new(Mutex::new(MessageInbox::new())), + inbox: Arc::new(Mutex::new(MessageInbox::new(notify))), outbox: Arc::new(Mutex::new(MessageOutbox::new())), + subscriber, }; tokio::task::spawn( @@ -341,7 +349,7 @@ impl MessageStack { // case we consider this to be a lingering chunk which was already accepted in the // meantime (as the message is complete). // - // SAFETY: a malcious node could send a lot of empty chunks, which trigger allocations + // SAFETY: a malicious node could send a lot of empty chunks, which trigger allocations // to hold the chunk array, effectively exhausting memory. As such, we first need to // determine if the chunk is feasible. let mut inbox = self.inbox.lock().unwrap(); @@ -460,6 +468,8 @@ impl MessageStack { // Move message to be read. inbox.complete_msges.push_back(message); inbox.pending_msges.remove(&message_id); + // Notify subscribers we have a new message. + inbox.notify.send_replace(()); Some(md.into_reply().into_inner()) } else { @@ -774,35 +784,6 @@ impl MessageStack { id } - /// Peek a [`Message`] from the inbound queue. - /// - /// The next call to [`MessageStack::peek_message`] or [`MessageStack::pop_message`] will - /// return the same message. - pub fn peek_message(&self) -> Option { - let inbox = self.inbox.lock().unwrap(); - let maybe_msg = inbox.complete_msges.front().map(Clone::clone); - - if let Some(ref msg) = maybe_msg { - // If a message is popped, notify the sender - self.notify_read(msg); - } - - maybe_msg - } - - /// Read a [`Message`] from the inbound queue and delete it. - pub fn pop_message(&self) -> Option { - let mut inbox = self.inbox.lock().unwrap(); - let maybe_msg = inbox.complete_msges.pop_front(); - - if let Some(ref msg) = maybe_msg { - // If a message is popped, notify the sender - self.notify_read(msg); - } - - maybe_msg - } - /// Get information about the status of an outbound message. pub fn message_info(&self, id: MessageId) -> Option { let outbox = self.outbox.lock().unwrap(); @@ -846,6 +827,35 @@ impl MessageStack { }) } + /// A future which eventually resolves to a new (inbound message)[`ReceivedMessage`], if new messages come in. + /// + /// If pop is false, the message is not removed and the next call of this method will return + /// the same message. + pub async fn message(&self, pop: bool) -> ReceivedMessage { + // Copy the subscriber since we need mutable access to it. + let mut subscriber = self.subscriber.clone(); + + loop { + // Scope to ensure we drop the lock after we checked for a message and don't hold + // it while waiting for a new notification. + { + let mut inbox = self.inbox.lock().unwrap(); + if let Some(msg) = if pop { + inbox.complete_msges.pop_front() + } else { + inbox.complete_msges.front().map(Clone::clone) + } { + self.notify_read(&msg); + return msg; + }; + } + + // Sender can never be dropped since we hold a reference to self which contains the + // inbox. + let _ = subscriber.changed().await; + } + } + /// Notify the sender of a message that it has been read. fn notify_read(&self, msg: &ReceivedMessage) { let mut mp = MessagePacket::new(PacketBuffer::new());