limit max paste size via terminal to 2MiB

Before this commit, When reading long paste, It would read until read
bytes > 0. This means, It'll only read 1 packet, Generate URL, return it
and exit

Now it has been updated to keep reading bytes until there is nothing for
500ms

src/main.rs

@@ -26,6 +26,7 @@ mod schema;
 use crate::routes::{config, display_paste, index, insert_paste, metrics, raw_paste};
 use chrono::{Duration, Utc};
 use diesel::prelude::*;
+use models::paste;
 use rocket::fairing::AdHoc;
 use rocket::http::Header;
 use rocket::shield::{Policy, Referrer, Shield};
@@ -151,6 +152,7 @@ async fn cli() {
 }
 
 const BUFFER_SIZE: usize = 128;
+const PASTE_SIZE: usize = 2 * 1024 * 1024; // 2MiB
 
 async fn save_paste(
     db: &mut PgConnection,
@@ -160,27 +162,49 @@ async fn save_paste(
     let mut paste = vec![];
     let mut buffer = [0; BUFFER_SIZE];
 
-    loop {
-        let read = socket.read(&mut buffer).await?;
-        if read < BUFFER_SIZE {
-            paste.extend(&buffer[..read]);
+    'outer: loop {
+        if paste.len() >= PASTE_SIZE {
             break;
         }
+        let five_hundred_millis_timeout = tokio::time::Duration::from_millis(500);
 
-        paste.extend(buffer);
-        buffer.fill(0);
+        let response =
+            tokio::time::timeout(five_hundred_millis_timeout, socket.read(&mut buffer)).await;
+
+        match response {
+            Ok(Ok(read)) => {
+                if read == 0 {
+                    break 'outer;
+                }
+
+                paste.extend(buffer);
+                buffer.fill(0);
+            }
+            Ok(Err(e)) => {
+                eprintln!("error in reading from socket: {e}");
+                break;
+            }
+            Err(_) => break,
+        }
     }
 
-    use models::paste;
+    while let Some(&c) = paste.last() {
+        if c == 0x00 {
+            paste.pop();
+        } else {
+            break;
+        }
+    }
+    let string = String::from_utf8_lossy(&paste);
     let v = paste::insert(
         db,
         // Terminal paste should stay for 3 days
         Some(Utc::now() + Duration::days(90)),
-        &String::from_utf8_lossy(&paste),
+        &string,
     )
     .unwrap();
 
-    let output = format!("https://{}/{}\n", host, v);
+    let output = format!("https://{}/raw/{}\n", host, v);
     let output: Vec<u8> = output.into_bytes();
 
     socket.write(&output).await?;