Experimental: websockets

src/main.rs

@@ -1,10 +1,10 @@
-use std::{env, fs::{self, File}, thread};
+use std::{collections::HashMap, env, fs::{self, File}, sync::{Arc, Mutex}, thread};
 
 use chrono::Utc;
 use kv::{Config, Store, Json, Bucket, Value};
 use paste::Paste;
 use rand::{Rng, distributions::Alphanumeric};
-use rouille::{Response, router, try_or_400, post_input, Request, Server};
+use rouille::{post_input, router, try_or_400, websocket, Request, Response, Server};
 use signal_hook::{iterator::Signals, consts::{SIGINT, SIGTERM}};
 
 mod paste;
@@ -30,6 +30,7 @@ fn generate_key<T : Value>(key: Option<String>, store: &Bucket<String, T>) -> St
     }
 }
 
+// Get query parameters
 fn get_query(req: &Request, name: &str) -> Option<String> {
     let params: Vec<&str> = req.raw_query_string().split('&').collect();
     let param = params.iter().find(|p|p.starts_with(&format!("{}=", name)));
@@ -58,6 +59,8 @@ fn main() {
     let links = store.bucket::<String, String>(Some("links"))
         .expect("Failed to open links bucket");
 
+    // Temp pastes
+    let temp_pastes = Arc::new(Mutex::new(HashMap::new()));
     
     // Create server
     let server = Server::new(format!("0.0.0.0:{}", port), move |req| {
@@ -68,6 +71,10 @@ fn main() {
             }
         }
 
+        if req.url().starts_with("/ws") {
+            handle_websocket(temp_pastes.clone(), &req);
+        }
+
         router!(req,
             (GET) (/) => {
                 match File::open(format!("{}/index.html", &source_dir)) {
@@ -115,7 +122,7 @@ fn main() {
             },
             (POST) (/{id: String}) => { 
                 let id = if id.is_empty() { None } else { Some(id) };
-                register_paste(&pastes, req, id, prefix.clone()) 
+                register_paste(&pastes, req, id, prefix.clone(), None) 
             },
             _ => Response::empty_404()
         )
@@ -143,8 +150,7 @@ fn main() {
 }
 
 // Register paste handler
-fn register_paste(pastes: &Bucket<String, Json<Paste>>, req: &Request, id: Option<String>, prefix: Option<String>) -> Response {
+fn register_paste(pastes: &Bucket<String, Json<Paste>>, req: &Request, id: Option<String>, prefix: Option<String>, allow_edit: Option<bool>) -> Response {
-
     // Try read body from form data, and if not present from request body
     let body = match post_input!(req, {
         content: String
@@ -165,6 +171,7 @@ fn register_paste(pastes: &Bucket<String, Json<Paste>>, req: &Request, id: Optio
         content: body,
         expires: None, 
         language, 
+        allow_edit,
         created: Utc::now()
     });
 
@@ -210,3 +217,31 @@ fn register_link(links: &Bucket<String,String>, req: &Request, id: Option<String
 
     Response::text(url)
 }
+
+fn handle_websocket(pastes: Arc<Mutex<HashMap<String, String>>>, req: &Request) -> Response {
+    let key = get_query(&req, "key").unwrap();
+
+    let mut current_text = String::new();
+    match pastes.lock().expect("Failed to access pastes DB").get(&key) {
+        Some(paste) => {
+            current_text = paste.clone()
+        },
+        None => {
+            // OK, ignore
+        }
+    }
+
+    let (response, websocket) = try_or_400!(websocket::start(req, Some("pastabble")));
+    if let Ok(mut ws) = websocket.recv() {
+        thread::spawn(move || {
+            if !current_text.is_empty() {
+                _ = ws.send_text(&current_text);
+            }
+
+            
+        });
+        return response.with_additional_header("Access-Control-Allow-Origin", "*");
+    } else {
+        return Response::empty_400()
+    }
+}

src/paste.rs

@@ -9,5 +9,7 @@ pub struct Paste {
 
     #[serde(with = "ts_seconds_option")]
     pub expires: Option<DateTime<Utc>>,
-    pub created: DateTime<Utc>
+    pub created: DateTime<Utc>,
+
+    pub allow_edit: Option<bool>
 }