Implement basic networking code

src/lib.rs

@@ -42,7 +42,8 @@ fn derive_protocol(input: proc_macro2::TokenStream) -> proc_macro2::TokenStream
     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 server_trait_name = format_ident!("{}Server", name);
+    let client_connection_struct_name = format_ident!("__{}Connection", name);
+    let server_trait_name = format_ident!("{}ServerTrait", name);
     let client_struct_name = format_ident!("{}Client", name);
 
     let vis = &input.vis;
@@ -210,6 +211,59 @@ fn derive_protocol(input: proc_macro2::TokenStream) -> proc_macro2::TokenStream
             }
         }
     };
+    // Create a struct to handle the connection from the client to the server
+    let stream_type = quote! { tokio::net::TcpStream }; // TODO: In the future we could support other stream types
+    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)>,
+            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)>,
+                stream: #stream_type,
+            ) -> Self {
+                Self {
+                    to_send,
+                    received,
+                    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((nonce, query)) = self.to_send.recv() => {
+                            let serialized = ron::ser::to_string(&query).expect("Failed to serialize query!");
+                            let len = serialized.len() as u32;
+                            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() => {
+                            match self.stream.try_read(&mut buf) {
+                                Ok(0) => break, // Stream closed
+                                Ok(n) => {
+                                    // TODO: This doesn't cope with partial reads, we will handle that later
+                                    let len = u32::from_le_bytes(buf[..4].try_into().expect("Failed to convert bytes to u32"));
+                                    let serialized = std::str::from_utf8(&buf[4..(4 + len as usize)]).expect("Failed to convert bytes to string");
+                                    let query: #answer_enum_name = ron::de::from_str(serialized).expect("Failed to deserialize query!");
+                                    self.received.send((0, query)).await.expect("Failed to send query!");
+                                    buf.clear();
+                                },
+                                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! {
@@ -232,7 +286,7 @@ fn derive_protocol(input: proc_macro2::TokenStream) -> proc_macro2::TokenStream
             queries: #queries_struct_name,
             send_queue: tokio::sync::mpsc::Sender<(u64, #question_enum_name)>,
             recv_queue: #client_recv_queue_wrapper,
-        } // TODO: This struct will have some fields to handle the actual connection
+        }
         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)>) -> Self {
                 Self {
@@ -279,6 +333,7 @@ fn derive_protocol(input: proc_macro2::TokenStream) -> proc_macro2::TokenStream
         #query_enum
         #queries_struct
         #server_trait
+        #cc_struct
         #client_struct
     };
     expanded