|
|
@@ -9,9 +9,7 @@ these work. */
|
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
|
use tokio::net::TcpListener;
|
|
|
|
|
|
-use std::str;
|
|
|
-
|
|
|
-pub async fn listen(addr: String, fun: fn(String) -> String) {
|
|
|
+pub async fn listen(addr: String, fun: fn(Vec<u8>) -> Vec<u8>) {
|
|
|
let listener = TcpListener::bind(&addr)
|
|
|
.await
|
|
|
.expect("Failed to create TcpListener");
|
|
|
@@ -22,11 +20,23 @@ pub async fn listen(addr: String, fun: fn(String) -> String) {
|
|
|
let (mut socket, _) = listener.accept().await.expect("Failed to create socket");
|
|
|
|
|
|
tokio::spawn(async move {
|
|
|
- let mut buf = vec![0; 1024];
|
|
|
-
|
|
|
- // read data, perform function on it, write result
|
|
|
loop {
|
|
|
- let n = socket
|
|
|
+ // get number of bytes to receive
|
|
|
+ let mut nbuf: [u8; 8] = [0; 8];
|
|
|
+ socket
|
|
|
+ .read(&mut nbuf)
|
|
|
+ .await
|
|
|
+ .expect("Failed to get number of bytes to read");
|
|
|
+ let n = usize::from_be_bytes(nbuf);
|
|
|
+
|
|
|
+ if n == 0 {
|
|
|
+ return;
|
|
|
+ }
|
|
|
+
|
|
|
+ let mut buf = vec![0; n];
|
|
|
+
|
|
|
+ // receive data
|
|
|
+ socket
|
|
|
.read(&mut buf)
|
|
|
.await
|
|
|
.expect("Failed to read data from socket");
|
|
|
@@ -35,15 +45,18 @@ pub async fn listen(addr: String, fun: fn(String) -> String) {
|
|
|
return;
|
|
|
}
|
|
|
|
|
|
- // I think this is a problem if there's more data than fits in the buffer...
|
|
|
- // But that's a problem for future me.
|
|
|
- let s = str::from_utf8(&buf[0..n])
|
|
|
- .expect("Invalid UTF-8 sequence")
|
|
|
- .to_string();
|
|
|
- let response = fun(s);
|
|
|
+ let response = fun(buf);
|
|
|
+
|
|
|
+ // send number of bytes in response
|
|
|
+ let response_size = usize::to_be_bytes(response.len());
|
|
|
+ socket
|
|
|
+ .write_all(&response_size)
|
|
|
+ .await
|
|
|
+ .expect("Failed to write number of bytes to listen for");
|
|
|
|
|
|
+ // send response
|
|
|
socket
|
|
|
- .write_all(response.as_bytes())
|
|
|
+ .write_all(&response)
|
|
|
.await
|
|
|
.expect("Failed to write data to socket");
|
|
|
}
|