775 lines
32 KiB
Rust
775 lines
32 KiB
Rust
/*
|
|
Eagle - A simple library for RPC in Rust
|
|
Copyright (c) 2024 KodiCraft
|
|
|
|
This program is free software: you can redistribute it and/or modify
|
|
it under the terms of the GNU Affero General Public License as
|
|
published by the Free Software Foundation, either version 3 of the
|
|
License, or (at your option) any later version.
|
|
|
|
This program is distributed in the hope that it will be useful,
|
|
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
GNU Affero General Public License for more details.
|
|
|
|
You should have received a copy of the GNU Affero General Public License
|
|
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
*/
|
|
|
|
//! # Eagle - A simple library for RPC in Rust
|
|
//!
|
|
//! <div class="warning">Eagle is still in early development. This documentation is subject to change and may not be entirely accurate.</div>
|
|
//!
|
|
//! Eagle is a library for building RPC protocols in Rust. It uses a macro
|
|
//! to transform your protocol definition into the necessary code to allow
|
|
//! communication between a server and a client.
|
|
//!
|
|
//! Eagle uses [`tokio`](https://tokio.rs) as its async runtime and
|
|
//! [`ron`](https://crates.io/crates/ron) for serialization.
|
|
//!
|
|
//! ## Usage
|
|
//! `eagle` is designed to be used to create your own protocol crate. We
|
|
//! recommend creating a new cargo workspace for your project with a shared
|
|
//! crate which will contain your protocol definition and individual crates
|
|
//! for the server and client.
|
|
//!
|
|
//! In your shared crate, you can define your protocol using the [`Protocol`]
|
|
//! derive macro. This will generate the necessary code for the server and
|
|
//! client to communicate.
|
|
//!
|
|
//! ```rust
|
|
//! use eagle::Protocol;
|
|
//!
|
|
//! #[derive(Protocol)]
|
|
//! pub enum Example {
|
|
//! Add((i32, i32), i32),
|
|
//! Length(String, usize),
|
|
//! /* ... */
|
|
//! }
|
|
//! ```
|
|
//!
|
|
//! The [`Protocol`] derive macro will generate all the necessary code, including
|
|
//! your handler trait, your server struct, and your client struct.
|
|
//!
|
|
//! On your server, you will need to implement the handler trait. This trait
|
|
//! describes the functions that the client can request from the server.
|
|
//!
|
|
//! ```rust
|
|
//! # use eagle::Protocol;
|
|
//! # #[derive(Protocol)]
|
|
//! # pub enum Example {
|
|
//! # Add((i32, i32), i32),
|
|
//! # Length(String, usize),
|
|
//! # /* ... */
|
|
//! # }
|
|
//! #
|
|
//! #[derive(Clone)]
|
|
//! pub struct Handler;
|
|
//! impl ExampleServerHandler for Handler {
|
|
//! async fn add(&mut self, a: i32, b: i32) -> i32 {
|
|
//! a + b
|
|
//! }
|
|
//! async fn length(&mut self, s: String) -> usize {
|
|
//! s.len()
|
|
//! }
|
|
//! /* ... */
|
|
//! }
|
|
//! ```
|
|
//!
|
|
//! To start the server, you simply need to use the generated server struct and
|
|
//! pass it your handler.
|
|
//!
|
|
//! ```no_run
|
|
//! # use eagle::Protocol;
|
|
//! # #[derive(Protocol)]
|
|
//! # pub enum Example {
|
|
//! # Add((i32, i32), i32),
|
|
//! # Length(String, usize),
|
|
//! # /* ... */
|
|
//! # }
|
|
//! #
|
|
//! # #[derive(Clone)]
|
|
//! # pub struct Handler;
|
|
//! # impl ExampleServerHandler for Handler {
|
|
//! # async fn add(&mut self, a: i32, b: i32) -> i32 {
|
|
//! # a + b
|
|
//! # }
|
|
//! # async fn length(&mut self, s: String) -> usize {
|
|
//! # s.len()
|
|
//! # }
|
|
//! # }
|
|
//! #
|
|
//! # tokio_test::block_on(async {
|
|
//! let handler = Handler;
|
|
//! let address = "127.0.0.1:12345"; // Or, if using the 'unix' feature, "/tmp/eagle.sock"
|
|
//! let server_task = tokio::spawn(ExampleServer::bind(handler, address));
|
|
//! # });
|
|
//! ```
|
|
//!
|
|
//! Please note the usage of `tokio::spawn`. This is because the `bind` function
|
|
//! will not return until the server is closed. You can use the `abort` method
|
|
//! on the task to close the server.
|
|
//!
|
|
//! On the client side, you can simply use the generated client struct to connect
|
|
//! to the server and begin sending queries.
|
|
//!
|
|
//! ```no_run
|
|
//! # use eagle::Protocol;
|
|
//! # #[derive(Protocol)]
|
|
//! # pub enum Example {
|
|
//! # Add((i32, i32), i32),
|
|
//! # Length(String, usize),
|
|
//! # /* ... */
|
|
//! # }
|
|
//! #
|
|
//! # #[derive(Clone)]
|
|
//! # pub struct Handler;
|
|
//! # impl ExampleServerHandler for Handler {
|
|
//! # async fn add(&mut self, a: i32, b: i32) -> i32 {
|
|
//! # a + b
|
|
//! # }
|
|
//! # async fn length(&mut self, s: String) -> usize {
|
|
//! # s.len()
|
|
//! # }
|
|
//! # }
|
|
//! #
|
|
//! # tokio_test::block_on(async {
|
|
//! # let handler = Handler;
|
|
//! let address = "127.0.0.1:12345"; // Or, if using the 'unix' feature, "/tmp/eagle.sock"
|
|
//! # let server_task = tokio::spawn(ExampleServer::bind(handler, address));
|
|
//! let client = ExampleClient::connect(address).await.unwrap();
|
|
//! # // Wait for the server to start, the developer is responsible for this in production
|
|
//! # tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
|
//! assert_eq!(client.add(2, 5).await.unwrap(), 7);
|
|
//! # });
|
|
//! ```
|
|
//!
|
|
//! The client can be closed by calling the `close` method on the client struct.
|
|
//! This will abort the connection.
|
|
|
|
#![warn(missing_docs)]
|
|
use proc_macro::TokenStream;
|
|
use quote::{format_ident, quote};
|
|
use syn::{parse2, spanned::Spanned, DeriveInput, Field, Ident};
|
|
|
|
#[cfg(all(feature = "tcp", feature = "unix"))]
|
|
compile_error!("You can only enable one of the 'tcp' or 'unix' features");
|
|
#[cfg(all(not(feature = "tcp"), not(feature = "unix")))]
|
|
compile_error!("You must enable either the 'tcp' or 'unix' feature");
|
|
#[cfg(all(feature = "unix", not(unix)))]
|
|
compile_error!("The 'unix' feature requires compiling for a unix target");
|
|
|
|
/// Generate all the necessary RPC code for a protocol from an enum describing it.
|
|
///
|
|
/// This macro will generate various enums and structs to enable communication
|
|
/// between a server and a client. The following items will be generated, where {}
|
|
/// is the name of the protocol enum:
|
|
/// - `{}ServerHandler` - A trait that the server must implement to handle queries
|
|
/// - `{}Server` - A struct that the server uses to communicate with clients
|
|
/// - `{}Client` - A struct that the client uses to communicate with a server
|
|
///
|
|
/// Each variant of the passed enum represents a query that the client can send to the
|
|
/// server. The first field of each variant is the question (serverbound), the second field
|
|
/// is the answer (clientbound). You may use tuples to represent sending multiple arguments and
|
|
/// you may use the unit type `()` to represent no arguments. Only data types which implement
|
|
/// [`Clone`], [`serde::Serialize`], and [`serde::Deserialize`] can be used.
|
|
///
|
|
/// For more information on how to use the generated code, see the [crate-level documentation](index.html).
|
|
///
|
|
/// # Example
|
|
/// ```rust
|
|
/// use eagle::Protocol;
|
|
///
|
|
/// #[derive(Protocol)]
|
|
/// pub enum Example {
|
|
/// Add((i32, i32), i32),
|
|
/// Length(String, usize),
|
|
/// }
|
|
/// ```
|
|
#[proc_macro_derive(Protocol)]
|
|
pub fn derive_protocol_derive(input: TokenStream) -> TokenStream {
|
|
let expanded = derive_protocol(input.into());
|
|
TokenStream::from(expanded)
|
|
}
|
|
|
|
fn derive_protocol(input: proc_macro2::TokenStream) -> proc_macro2::TokenStream {
|
|
let input = parse2::<DeriveInput>(input).unwrap();
|
|
|
|
// TODO: These logs should be filterable in some way
|
|
#[cfg(feature = "log")]
|
|
#[allow(unused_variables)]
|
|
let debug = quote! { log::debug! };
|
|
#[cfg(feature = "log")]
|
|
#[allow(unused_variables)]
|
|
let info = quote! { log::info! };
|
|
#[cfg(feature = "log")]
|
|
#[allow(unused_variables)]
|
|
let warn = quote! { log::warn! };
|
|
#[cfg(feature = "log")]
|
|
#[allow(unused_variables)]
|
|
let error = quote! { log::error! };
|
|
#[cfg(not(feature = "log"))]
|
|
#[allow(unused_variables)]
|
|
let debug = quote! { eprintln! };
|
|
#[cfg(not(feature = "log"))]
|
|
#[allow(unused_variables)]
|
|
let info = quote! { eprintln! };
|
|
#[cfg(not(feature = "log"))]
|
|
#[allow(unused_variables)]
|
|
let warn = quote! { eprintln! };
|
|
#[cfg(not(feature = "log"))]
|
|
#[allow(unused_variables)]
|
|
let error = quote! { eprintln! };
|
|
|
|
// Must be on an enum
|
|
let enum_ = match &input.data {
|
|
syn::Data::Enum(e) => e,
|
|
_ => {
|
|
return syn::Error::new(input.span(), "Protocol can only be derived on enums")
|
|
.to_compile_error()
|
|
}
|
|
};
|
|
let name = &input.ident;
|
|
|
|
let error_enum_name = format_ident!("__{}Error", name);
|
|
let answer_enum_name = format_ident!("__{}Answer", name);
|
|
let question_enum_name = format_ident!("__{}Question", name);
|
|
let query_enum_name = format_ident!("__{}Query", name);
|
|
let queries_struct_name = format_ident!("__{}Queries", name);
|
|
let client_connection_struct_name = format_ident!("__{}Connection", name);
|
|
let server_trait_name = format_ident!("{}ServerHandler", name);
|
|
let server_connection_struct_name = format_ident!("{}Server", name);
|
|
let client_struct_name = format_ident!("{}Client", name);
|
|
|
|
let vis = &input.vis;
|
|
|
|
let mut server_trait = Vec::new();
|
|
let mut server_enum = Vec::new();
|
|
let mut client_impl = Vec::new();
|
|
let mut client_enum = Vec::new();
|
|
|
|
let mut server_handler = Vec::new();
|
|
|
|
let mut query_enum = Vec::new();
|
|
let mut query_from_question_enum = Vec::new();
|
|
let mut query_set_answer = Vec::new();
|
|
let mut query_get_answer = Vec::new();
|
|
|
|
for variant in &enum_.variants {
|
|
// Every variant must have 2 fields
|
|
// The first field is the question (serverbound), the second field is the answer (clientbound)
|
|
if variant.fields.len() != 2 {
|
|
return syn::Error::new(
|
|
variant.span(),
|
|
"Every variant on a protocol must have exactly 2 fields",
|
|
)
|
|
.to_compile_error();
|
|
}
|
|
|
|
let var_name = ident_to_snake_case(&variant.ident);
|
|
let mut variant_fields = variant.fields.iter();
|
|
let question_field = variant_fields.next().unwrap();
|
|
let question_args = field_to_args(question_field);
|
|
let question_handler_args = field_to_handler_args(question_field);
|
|
let question_tuple_args = field_to_tuple_args(question_field);
|
|
let answer_type = variant_fields.next().unwrap().ty.clone();
|
|
|
|
// The variants that either the server or the client will use
|
|
// The "server" enum contains messages the server can send, the "client" enum contains messages the client can send
|
|
server_enum.push(quote! {
|
|
#var_name(#answer_type)
|
|
});
|
|
client_enum.push(quote! {
|
|
#var_name(#question_field)
|
|
});
|
|
// There is a From implementation for the client enum to the query enum
|
|
query_from_question_enum.push(quote! {
|
|
#question_enum_name::#var_name(question) => #query_enum_name::#var_name(question, None),
|
|
});
|
|
// There is a function that must be implemented to set the answer in the query enum
|
|
query_set_answer.push(quote! {
|
|
#query_enum_name::#var_name(question, answer_opt) => match answer {
|
|
#answer_enum_name::#var_name(answer) => {
|
|
#debug("Setting answer for query {}", stringify!(#var_name));
|
|
answer_opt.replace(answer);
|
|
},
|
|
_ => panic!("The answer for this query is not the correct type."),
|
|
},
|
|
});
|
|
// There is a function that must be implemented to get the answer from the query enum
|
|
query_get_answer.push(quote! {
|
|
#query_enum_name::#var_name(_, answer) => match answer {
|
|
Some(answer) => Some(#answer_enum_name::#var_name(answer.clone())),
|
|
None => None
|
|
},
|
|
});
|
|
// There is a function that the server uses to call the appropriate function when receiving a query
|
|
server_handler.push(quote! {
|
|
#question_enum_name::#var_name(#question_tuple_args) => {
|
|
#info("Received query {}", stringify!(#var_name));
|
|
let answer = self.handler.lock().await.#var_name(#question_handler_args).await;
|
|
return #answer_enum_name::#var_name(answer);
|
|
},
|
|
});
|
|
// The function that the server needs to implement
|
|
server_trait.push(quote! {
|
|
fn #var_name(&mut self, #question_args) -> impl std::future::Future<Output = #answer_type> + Send;
|
|
});
|
|
// The function that the client uses to communicate
|
|
client_impl.push(quote! {
|
|
pub async fn #var_name(&self, #question_args) -> Result<#answer_type, #error_enum_name> {
|
|
#info("Sending query {}", stringify!(#var_name));
|
|
let nonce = self.send(#question_enum_name::#var_name(#question_tuple_args)).await?;
|
|
let answer = self.recv_until(nonce).await?;
|
|
match answer {
|
|
#answer_enum_name::#var_name(answer) => Ok(answer),
|
|
_ => panic!("The answer for this query is not the correct type."),
|
|
}
|
|
}
|
|
});
|
|
// The query enum is the same as the source enum, but the second field is always wrapped in a Option<>
|
|
query_enum.push(quote! {
|
|
#var_name(#question_field, Option<#answer_type>)
|
|
});
|
|
}
|
|
|
|
// Create an error and result type for sending messages
|
|
let error_enum = quote! {
|
|
#[derive(Debug)]
|
|
#vis enum #error_enum_name {
|
|
SendError(tokio::sync::mpsc::error::SendError<(u64, #question_enum_name)>),
|
|
Closed,
|
|
}
|
|
};
|
|
// Create enums for the types of messages the server and client will use
|
|
|
|
let answer_enum = quote! {
|
|
#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)]
|
|
#vis enum #answer_enum_name {
|
|
#(#server_enum), *,
|
|
Ready
|
|
}
|
|
};
|
|
let question_enum = quote! {
|
|
#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)]
|
|
#vis enum #question_enum_name {
|
|
#(#client_enum), *
|
|
}
|
|
};
|
|
// Create an enum to represent the queries the client has sent
|
|
let query_enum = quote! {
|
|
#[derive(Clone, Debug)]
|
|
#vis enum #query_enum_name {
|
|
#(#query_enum), *
|
|
}
|
|
impl #query_enum_name {
|
|
pub fn set_answer(&mut self, answer: #answer_enum_name) {
|
|
match self {
|
|
#(#query_set_answer)*
|
|
};
|
|
}
|
|
pub fn get_answer(&self) -> Option<#answer_enum_name> {
|
|
match self {
|
|
#(#query_get_answer)*
|
|
}
|
|
}
|
|
}
|
|
impl From<#question_enum_name> for #query_enum_name {
|
|
fn from(query: #question_enum_name) -> Self {
|
|
match query {
|
|
#(#query_from_question_enum)*
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
#[cfg(feature = "tcp")]
|
|
let stream_type = quote! { tokio::net::TcpStream };
|
|
#[cfg(feature = "tcp")]
|
|
let stream_addr_trait = quote! { tokio::net::ToSocketAddrs };
|
|
#[cfg(feature = "tcp")]
|
|
let listener_type = quote! { tokio::net::TcpListener };
|
|
#[cfg(feature = "unix")]
|
|
let stream_type = quote! { tokio::net::UnixStream };
|
|
#[cfg(feature = "unix")]
|
|
let stream_addr_trait = quote! { std::convert::AsRef<std::path::Path> };
|
|
#[cfg(feature = "unix")]
|
|
let listener_type = quote! { tokio::net::UnixListener };
|
|
|
|
// Create a trait which the server will have to implement
|
|
let server_trait = quote! {
|
|
#vis trait #server_trait_name {
|
|
#(#server_trait)*
|
|
}
|
|
};
|
|
|
|
// Create a struct to implement the communication between the server and the client
|
|
#[cfg(feature = "tcp")]
|
|
let listener_statement = quote! {
|
|
let listener = #listener_type::bind(addr).await?;
|
|
};
|
|
#[cfg(feature = "unix")]
|
|
let listener_statement = quote! {
|
|
let listener = #listener_type::bind(addr.as_ref())?;
|
|
};
|
|
let sc_struct = quote! {
|
|
#[derive(Clone)]
|
|
#vis struct #server_connection_struct_name<H: #server_trait_name + ::std::marker::Send + Clone + 'static> {
|
|
handler: ::std::sync::Arc<tokio::sync::Mutex<H>>,
|
|
tasks: ::std::sync::Arc<tokio::sync::Mutex<Vec<tokio::task::JoinHandle<()>>>>,
|
|
}
|
|
impl<H: #server_trait_name + ::std::marker::Send + Clone + 'static> #server_connection_struct_name<H> {
|
|
pub async fn bind<A: #stream_addr_trait + ::std::marker::Send + std::fmt::Display + 'static>(handler: H, addr: A) -> Self {
|
|
#info("Binding server to address {}", addr);
|
|
let handler = ::std::sync::Arc::new(tokio::sync::Mutex::new(handler));
|
|
let tasks = ::std::sync::Arc::new(tokio::sync::Mutex::new(Vec::new()));
|
|
let sc = Self {
|
|
handler,
|
|
tasks,
|
|
};
|
|
let sc_clone = sc.clone();
|
|
let acc_task = tokio::spawn(async move {
|
|
sc_clone.accept_connections(addr).await.expect("Failed to accept connections!");
|
|
});
|
|
sc.tasks.lock().await.push(acc_task);
|
|
sc
|
|
}
|
|
|
|
pub async fn accept_connections<A: #stream_addr_trait>(
|
|
&self,
|
|
addr: A,
|
|
) -> Result<(), std::io::Error> {
|
|
#listener_statement
|
|
loop {
|
|
let (stream, _) = listener.accept().await?;
|
|
#info("Accepted connection from {:?}", stream.peer_addr()?);
|
|
let self_clone = self.clone();
|
|
let run_task = tokio::spawn(async move {
|
|
self_clone.run(stream).await;
|
|
});
|
|
self.tasks.lock().await.push(run_task);
|
|
}
|
|
}
|
|
|
|
async fn handle(&self, question: #question_enum_name) -> #answer_enum_name {
|
|
match question {
|
|
#(#server_handler)*
|
|
}
|
|
}
|
|
|
|
async fn send(&self, stream: &mut #stream_type, nonce: u64, answer: #answer_enum_name) {
|
|
use tokio::io::AsyncWriteExt;
|
|
let serialized = ron::ser::to_string(&(nonce, answer)).expect("Failed to serialize response!");
|
|
let len = serialized.len() as u32;
|
|
#debug("Sending `{}`", serialized);
|
|
stream.write_all(&len.to_le_bytes()).await.expect("Failed to write length!");
|
|
stream.write_all(serialized.as_bytes()).await.expect("Failed to write response!");
|
|
}
|
|
|
|
async fn run(&self, mut stream: #stream_type) {
|
|
use tokio::io::AsyncWriteExt;
|
|
use tokio::io::AsyncReadExt;
|
|
let mut buf = Vec::with_capacity(1024);
|
|
self.send(&mut stream, 0, #answer_enum_name::Ready).await;
|
|
loop {
|
|
tokio::select! {
|
|
Ok(_) = stream.readable() => {
|
|
let mut read_buf = [0; 1024];
|
|
match stream.try_read(&mut read_buf) {
|
|
Ok(0) => break, // Stream closed
|
|
Ok(n) => {
|
|
#debug("Received {} bytes (server)", n);
|
|
buf.extend_from_slice(&read_buf[..n]);
|
|
while buf.len() >= 4 {
|
|
let len = u32::from_le_bytes(buf[..4].try_into().expect("Failed to convert bytes to u32"));
|
|
if buf.len() >= (4 + len as usize) {
|
|
let serialized = std::str::from_utf8(&buf[4..(4 + len as usize)]).expect("Failed to convert bytes to string");
|
|
let (nonce, question): (u64, #question_enum_name) = ron::de::from_str(serialized).expect("Failed to deserialize query!");
|
|
let answer = self.handle(question).await;
|
|
self.send(&mut stream, nonce, answer).await;
|
|
buf.drain(0..(4 + len as usize));
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
},
|
|
Err(ref e) if e.kind() == ::std::io::ErrorKind::WouldBlock => { continue; },
|
|
Err(e) => eprintln!("Error reading from stream: {:?}", e),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
// Create a struct to hold queries behind an Arc<Mutex<>> to enable async access
|
|
// TODO: It might be a good idea to just make this a generic struct and write it in actual code
|
|
// rather than in this macro
|
|
let queries_struct = quote! {
|
|
#[derive(Clone)]
|
|
struct #queries_struct_name {
|
|
queries: ::std::sync::Arc<::std::sync::Mutex<::std::collections::HashMap<u64, #query_enum_name>>>,
|
|
}
|
|
impl #queries_struct_name {
|
|
fn new() -> Self {
|
|
Self {
|
|
queries: ::std::sync::Arc::new(::std::sync::Mutex::new(::std::collections::HashMap::new())),
|
|
}
|
|
}
|
|
|
|
pub fn insert(&self, nonce: u64, query: #query_enum_name) {
|
|
self.queries.lock().unwrap().insert(nonce, query);
|
|
}
|
|
|
|
pub fn get(&self, nonce: &u64) -> Option<#query_enum_name> {
|
|
self.queries.lock().unwrap().get(nonce).cloned()
|
|
}
|
|
|
|
pub fn set_answer(&self, nonce: u64, answer: #answer_enum_name) {
|
|
if let Some(query) = self.queries.lock().unwrap().get_mut(&nonce) {
|
|
query.set_answer(answer);
|
|
}
|
|
}
|
|
|
|
pub fn len(&self) -> usize {
|
|
self.queries.lock().unwrap().len()
|
|
}
|
|
}
|
|
};
|
|
|
|
// Create a struct to handle the connection from the client to the server
|
|
let cc_struct = quote! {
|
|
struct #client_connection_struct_name {
|
|
to_send: tokio::sync::mpsc::Receiver<(u64, #question_enum_name)>,
|
|
received: tokio::sync::mpsc::Sender<(u64, #answer_enum_name)>,
|
|
ready: std::sync::Arc<tokio::sync::Notify>,
|
|
stream: #stream_type,
|
|
}
|
|
impl #client_connection_struct_name {
|
|
pub fn new(
|
|
to_send: tokio::sync::mpsc::Receiver<(u64, #question_enum_name)>,
|
|
received: tokio::sync::mpsc::Sender<(u64, #answer_enum_name)>,
|
|
ready: std::sync::Arc<tokio::sync::Notify>,
|
|
stream: #stream_type,
|
|
) -> Self {
|
|
Self {
|
|
to_send,
|
|
received,
|
|
ready,
|
|
stream,
|
|
}
|
|
}
|
|
|
|
pub async fn run(mut self) {
|
|
use tokio::io::AsyncWriteExt;
|
|
use tokio::io::AsyncReadExt;
|
|
let mut buf = Vec::with_capacity(1024);
|
|
loop {
|
|
tokio::select! {
|
|
Some(msg) = self.to_send.recv() => {
|
|
let serialized = ron::ser::to_string(&msg).expect("Failed to serialize query!");
|
|
let len = serialized.len() as u32;
|
|
#debug("Sending `{}`", serialized);
|
|
self.stream.write_all(&len.to_le_bytes()).await.expect("Failed to write length!");
|
|
self.stream.write_all(serialized.as_bytes()).await.expect("Failed to write query!");
|
|
},
|
|
Ok(_) = self.stream.readable() => {
|
|
let mut read_buf = [0; 1024];
|
|
match self.stream.try_read(&mut read_buf) {
|
|
Ok(0) => { break; },
|
|
Ok(n) => {
|
|
#debug("Received {} bytes (client)", n);
|
|
buf.extend_from_slice(&read_buf[..n]);
|
|
while buf.len() >= 4 {
|
|
let len = u32::from_le_bytes(buf[..4].try_into().expect("Failed to convert bytes to u32"));
|
|
if buf.len() >= (4 + len as usize) {
|
|
let serialized = std::str::from_utf8(&buf[4..(4 + len as usize)]).expect("Failed to convert bytes to string");
|
|
let response: (u64, #answer_enum_name) = ron::de::from_str(serialized).expect("Failed to deserialize response!");
|
|
if let #answer_enum_name::Ready = response.1 {
|
|
#debug("Received ready signal");
|
|
self.ready.notify_one();
|
|
} else {
|
|
self.received.send(response).await.expect("Failed to send response!");
|
|
}
|
|
buf.drain(0..(4 + len as usize));
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
},
|
|
Err(ref e) if e.kind() == ::std::io::ErrorKind::WouldBlock => { continue; },
|
|
Err(e) => eprintln!("Error reading from stream: {:?}", e),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
// Create a struct which the client will use to communicate
|
|
let client_recv_queue_wrapper = format_ident!("__{}RecvQueueWrapper", name);
|
|
let client_struct = quote! {
|
|
#[derive(Clone)]
|
|
struct #client_recv_queue_wrapper {
|
|
recv_queue: ::std::sync::Arc<::tokio::sync::Mutex<tokio::sync::mpsc::Receiver<(u64, #answer_enum_name)>>>,
|
|
}
|
|
impl #client_recv_queue_wrapper {
|
|
fn new(recv_queue: tokio::sync::mpsc::Receiver<(u64, #answer_enum_name)>) -> Self {
|
|
Self {
|
|
recv_queue: ::std::sync::Arc::new(::tokio::sync::Mutex::new(recv_queue)),
|
|
}
|
|
}
|
|
async fn recv(&self) -> Option<(u64, #answer_enum_name)> {
|
|
self.recv_queue.lock().await.recv().await
|
|
}
|
|
}
|
|
#[derive(Clone)]
|
|
#vis struct #client_struct_name {
|
|
queries: #queries_struct_name,
|
|
send_queue: tokio::sync::mpsc::Sender<(u64, #question_enum_name)>,
|
|
recv_queue: #client_recv_queue_wrapper,
|
|
ready: ::std::sync::Arc<tokio::sync::Mutex<bool>>,
|
|
ready_notify: ::std::sync::Arc<tokio::sync::Notify>,
|
|
connection_task: Option<::std::sync::Arc<tokio::task::JoinHandle<()>>>,
|
|
}
|
|
impl #client_struct_name {
|
|
pub fn new(send_queue: tokio::sync::mpsc::Sender<(u64, #question_enum_name)>,
|
|
recv_queue: tokio::sync::mpsc::Receiver<(u64, #answer_enum_name)>,
|
|
connection_task: Option<::std::sync::Arc<tokio::task::JoinHandle<()>>>,
|
|
ready_notify: ::std::sync::Arc<tokio::sync::Notify>) -> Self {
|
|
Self {
|
|
queries: #queries_struct_name::new(),
|
|
recv_queue: #client_recv_queue_wrapper::new(recv_queue),
|
|
ready: ::std::sync::Arc::new(false.into()),
|
|
ready_notify,
|
|
send_queue,
|
|
connection_task,
|
|
}
|
|
}
|
|
pub async fn connect<A: #stream_addr_trait + ::std::fmt::Display>(addr: A) -> Result<Self, std::io::Error> {
|
|
#info("Connecting to server at address {}", addr);
|
|
let stream = #stream_type::connect(addr).await?;
|
|
let (send_queue, to_send) = tokio::sync::mpsc::channel(16);
|
|
let (to_recv, recv_queue) = tokio::sync::mpsc::channel(16);
|
|
let ready_notify = ::std::sync::Arc::new(tokio::sync::Notify::new());
|
|
let connection = #client_connection_struct_name::new(to_send, to_recv, ready_notify.clone(), stream);
|
|
let connection_task = tokio::spawn(connection.run());
|
|
Ok(Self::new(send_queue, recv_queue, Some(::std::sync::Arc::new(connection_task)), ready_notify))
|
|
}
|
|
pub fn close(self) {
|
|
if let Some(task) = self.connection_task {
|
|
task.abort();
|
|
}
|
|
}
|
|
async fn send(&self, query: #question_enum_name) -> Result<u64, #error_enum_name> {
|
|
// Wait until the connection is ready
|
|
if !*self.ready.lock().await {
|
|
self.ready_notify.notified().await;
|
|
let mut ready = self.ready.lock().await;
|
|
*ready = true;
|
|
}
|
|
let nonce = self.queries.len() as u64;
|
|
let res = self.send_queue.send((nonce, query.clone())).await;
|
|
match res {
|
|
Ok(_) => {
|
|
self.queries.insert(nonce, query.into());
|
|
Ok(nonce)
|
|
}
|
|
Err(e) => Err(#error_enum_name::SendError(e)),
|
|
}
|
|
}
|
|
async fn recv_until(&self, id: u64) -> Result<#answer_enum_name, #error_enum_name> {
|
|
loop {
|
|
// Check if we've received the answer for the query we're looking for
|
|
if let Some(query) = self.queries.get(&id) {
|
|
if let Some(answer) = query.get_answer() {
|
|
#info("Found answer for query {}", id);
|
|
return Ok(answer);
|
|
}
|
|
}
|
|
match self.recv_queue.recv().await {
|
|
Some((nonce, answer)) => {
|
|
#info("Received answer for query {}", nonce);
|
|
self.queries.set_answer(nonce, answer.clone());
|
|
}
|
|
None => return Err(#error_enum_name::Closed),
|
|
};
|
|
}
|
|
}
|
|
#(#client_impl)*
|
|
}
|
|
};
|
|
|
|
let expanded = quote! {
|
|
#error_enum
|
|
#answer_enum
|
|
#question_enum
|
|
#query_enum
|
|
#queries_struct
|
|
#server_trait
|
|
#sc_struct
|
|
#cc_struct
|
|
#client_struct
|
|
};
|
|
expanded
|
|
}
|
|
|
|
fn ident_to_snake_case(ident: &Ident) -> Ident {
|
|
let ident = ident.to_string();
|
|
let mut out = String::new();
|
|
for (i, c) in ident.chars().enumerate() {
|
|
if c.is_uppercase() {
|
|
if i != 0 {
|
|
out.push('_');
|
|
}
|
|
out.push(c.to_lowercase().next().unwrap());
|
|
} else {
|
|
out.push(c);
|
|
}
|
|
}
|
|
Ident::new(&out, ident.span())
|
|
}
|
|
|
|
fn field_to_args(field: &Field) -> proc_macro2::TokenStream {
|
|
let type_ = &field.ty;
|
|
if let syn::Type::Tuple(tuple) = type_ {
|
|
let mut args = Vec::new();
|
|
for (i, elem) in tuple.elems.iter().enumerate() {
|
|
let arg = Ident::new(&format!("arg{}", i), elem.span());
|
|
args.push(quote! { #arg: #elem });
|
|
}
|
|
quote! { #( #args ), * }
|
|
} else {
|
|
quote! { arg: #type_ }
|
|
}
|
|
}
|
|
|
|
fn field_to_tuple_args(field: &Field) -> proc_macro2::TokenStream {
|
|
let type_ = &field.ty;
|
|
if let syn::Type::Tuple(tuple) = type_ {
|
|
let mut args = Vec::new();
|
|
for (i, elem) in tuple.elems.iter().enumerate() {
|
|
let arg = Ident::new(&format!("arg{}", i), elem.span());
|
|
args.push(quote! { #arg });
|
|
}
|
|
quote! { ( #( #args ), * ) }
|
|
} else {
|
|
quote! { (arg) }
|
|
}
|
|
}
|
|
|
|
fn field_to_handler_args(field: &Field) -> proc_macro2::TokenStream {
|
|
let type_ = &field.ty;
|
|
if let syn::Type::Tuple(tuple) = type_ {
|
|
let mut args = Vec::new();
|
|
for (i, elem) in tuple.elems.iter().enumerate() {
|
|
let arg = Ident::new(&format!("arg{}", i), elem.span());
|
|
args.push(quote! { #arg });
|
|
}
|
|
quote! { #( #args ), * }
|
|
} else {
|
|
quote! { arg }
|
|
}
|
|
}
|