From ccd4da433b4f267cde0a5db6e740bd50f5608b62 Mon Sep 17 00:00:00 2001
From: Ben Bridle <bridle.benjamin@gmail.com>
Date: Sat, 7 Aug 2021 14:20:59 +1200
Subject: [PATCH] Make the TcpServer work with bytes instead of strings

This makes it possible to implement byte-based protocols using TcpServer, instead of only being able to implement string-based protocols.
---
 src/connection.rs       |  4 ++--
 src/request_response.rs |  4 ++--
 src/tcp_server.rs       | 15 ++++-----------
 3 files changed, 8 insertions(+), 15 deletions(-)

diff --git a/src/connection.rs b/src/connection.rs
index e589007..83145a9 100644
--- a/src/connection.rs
+++ b/src/connection.rs
@@ -15,13 +15,13 @@ impl<Req: Request, Res: Response> Connection<Req, Res> {
             stream,
             address,
             token,
-            state: RequestState::Incoming(String::new()),
+            state: RequestState::Incoming(Vec::new()),
         }
     }
 }
 
 pub enum RequestState<Req: Request, Res: Response> {
-    Incoming(String),
+    Incoming(Vec<u8>),
     Processing(Req),
     Outgoing(Res),
 }
diff --git a/src/request_response.rs b/src/request_response.rs
index b6eaaf5..6e4c3eb 100644
--- a/src/request_response.rs
+++ b/src/request_response.rs
@@ -1,13 +1,13 @@
 pub trait Request {
     type Response;
-    fn from_string(request_string: &str) -> RequestParseResult<Self, Self::Response>
+    fn from_bytes(request_bytes: &[u8]) -> RequestParseResult<Self, Self::Response>
     where
         Self: Sized,
         Self::Response: Response;
 }
 
 pub trait Response {
-    fn to_string(self) -> String;
+    fn to_bytes(self) -> Vec<u8>;
 }
 
 pub type ProcessRequest<Req, Res> = fn(request: &Req) -> Res;
diff --git a/src/tcp_server.rs b/src/tcp_server.rs
index 7cc0128..b49f654 100644
--- a/src/tcp_server.rs
+++ b/src/tcp_server.rs
@@ -144,19 +144,13 @@ impl<
                         self.remove_connection(event.token()).unwrap();
                         return;
                     }
-                    Ok(_) => match std::str::from_utf8(&buffer) {
-                        Ok(s) => request_data.push_str(s.trim_matches(char::from(0))),
-                        Err(e) => panic!(
-                            "Incoming data from {:?} is not valid UTF-8: {}",
-                            connection.address, e
-                        ),
-                    },
+                    Ok(len) => request_data.extend_from_slice(&buffer[..len]),
                     Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => break,
                     Err(e) => panic!("Unexpected error: {}", e),
                 };
             }
 
-            match Req::from_string(&request_data) {
+            match Req::from_bytes(&request_data) {
                 RequestParseResult::Complete(request) => {
                     let mut connection =
                         std::mem::replace(&mut self.connections[token.0], None).unwrap();
@@ -180,9 +174,8 @@ impl<
         let token = event.token();
         let mut connection = std::mem::replace(&mut self.connections[token.0], None).unwrap();
         if let RequestState::Outgoing(response) = connection.state {
-            let response_string = response.to_string();
-            let bytes = response_string.as_bytes();
-            connection.stream.write_all(bytes).unwrap();
+            let bytes = response.to_bytes();
+            connection.stream.write_all(&bytes).unwrap();
         } else {
             info!("Received write event for non-outgoing connection")
         }