feat: impl websocket upload file for doc_info (#20)

Cargo.toml

@@ -13,9 +13,15 @@ actix-multipart = "0.4"
 actix-session = { version = "0.5" }
 actix-identity = { version = "0.4" }
 actix-web-httpauth = { version = "0.6" }
+actix-ws = "0.2.5"
+uuid = { version = "1.6.1", features = [
+    "v4",
+    "fast-rng",
+    "macro-diagnostics",
+] }
 thiserror = "1.0"
 postgres = "0.19.7"
-sea-orm = {version = "0.12.9", features = ["sqlx-postgres", "runtime-tokio-native-tls", "macros"]}
+sea-orm = { version = "0.12.9", features = ["sqlx-postgres", "runtime-tokio-native-tls", "macros"] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1.0"
 tracing-subscriber = "0.3.18"
@@ -27,6 +33,7 @@ minio = "0.1.0"
 futures-util = "0.3.29"
 actix-multipart-extract = "0.1.5"
 regex = "1.10.2"
+tokio = { version = "1.35.1", features = ["rt", "time", "macros"] }
 
 [[bin]]
 name = "doc_gpt"

migration/src/m20220101_000001_create_table.rs

@@ -1,6 +1,7 @@
-use sea_orm_migration::{ prelude::*, sea_orm::Statement };
+use sea_orm_migration::prelude::*;
 use chrono::{ FixedOffset, Utc };
 
+#[allow(dead_code)]
 fn now() -> chrono::DateTime<FixedOffset> {
     Utc::now().with_timezone(&FixedOffset::east_opt(3600 * 8).unwrap())
 }

src/api/doc_info.rs

@@ -107,8 +107,12 @@ async fn upload(
     payload: Multipart<UploadForm>,
     data: web::Data<AppState>
 ) -> Result<HttpResponse, AppError> {
-    let uid = payload.uid;
+
-    let file_name = payload.file_field.name.as_str();
+
+    Ok(HttpResponse::Ok().body("File uploaded successfully"))
+}
+
+pub(crate) async fn _upload_file(uid: i64, did: i64, file_name: &str, bytes: &[u8], data: &web::Data<AppState>) -> Result<(), AppError> {
     async fn add_number_to_filename(
         file_name: &str,
         conn: &DbConn,
@@ -138,9 +142,9 @@ async fn upload(
         }
         new_file_name
     }
-    let fnm = add_number_to_filename(file_name, &data.conn, uid, payload.did).await;
+    let fnm = add_number_to_filename(file_name, &data.conn, uid, did).await;
 
-    let bucket_name = format!("{}-upload", payload.uid);
+    let bucket_name = format!("{}-upload", uid);
     let s3_client: &minio::s3::client::Client = &data.s3_client;
     let buckets_exists = s3_client
         .bucket_exists(&BucketExistsArgs::new(&bucket_name).unwrap()).await
@@ -152,7 +156,7 @@ async fn upload(
         print!("Existing bucket: {}", bucket_name.clone());
     }
 
-    let location = format!("/{}/{}", payload.did, fnm)
+    let location = format!("/{}/{}", did, fnm)
         .as_bytes()
         .to_vec()
         .iter()
@@ -164,8 +168,8 @@ async fn upload(
         &mut PutObjectArgs::new(
             &bucket_name,
             &location,
-            &mut BufReader::new(payload.file_field.bytes.as_slice()),
+            &mut BufReader::new(bytes),
-            Some(payload.file_field.bytes.len()),
+            Some(bytes.len()),
             None
         )?
     ).await?;
@@ -174,7 +178,7 @@ async fn upload(
         did: Default::default(),
         uid: uid,
         doc_name: fnm.clone(),
-        size: payload.file_field.bytes.len() as i64,
+        size: bytes.len() as i64,
         location,
         r#type: file_type(&fnm),
         thumbnail_base64: Default::default(),
@@ -183,9 +187,9 @@ async fn upload(
         is_deleted: Default::default(),
     }).await?;
 
-    let _ = Mutation::place_doc(&data.conn, payload.did, doc.did.unwrap()).await?;
+    let _ = Mutation::place_doc(&data.conn, did, doc.did.unwrap()).await?;
 
-    Ok(HttpResponse::Ok().body("File uploaded successfully"))
+    Ok(())
 }
 
 #[derive(Deserialize, Debug)]

src/main.rs

@@ -2,6 +2,7 @@ mod api;
 mod entity;
 mod service;
 mod errors;
+mod web_socket;
 
 use std::env;
 use actix_files::Files;
@@ -19,6 +20,7 @@ use minio::s3::http::BaseUrl;
 use sea_orm::{ Database, DatabaseConnection };
 use migration::{ Migrator, MigratorTrait };
 use crate::errors::{ AppError, UserError };
+use crate::web_socket::doc_info::upload_file_ws;
 
 #[derive(Debug, Clone)]
 struct AppState {
@@ -138,4 +140,6 @@ fn init(cfg: &mut web::ServiceConfig) {
     cfg.service(api::user_info::login);
     cfg.service(api::user_info::register);
     cfg.service(api::user_info::setting);
+
+    cfg.service(web::resource("/ws-upload-doc").route(web::get().to(upload_file_ws)));
 }

src/web_socket/doc_info.rs

@@ -0,0 +1,97 @@
+use std::io::{Cursor, Write};
+use std::time::{Duration, Instant};
+use actix_rt::time::interval;
+use actix_web::{HttpRequest, HttpResponse, rt, web};
+use actix_web::web::Buf;
+use actix_ws::Message;
+use futures_util::{future, StreamExt};
+use futures_util::future::Either;
+use uuid::Uuid;
+use crate::api::doc_info::_upload_file;
+use crate::AppState;
+use crate::errors::AppError;
+
+const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
+
+/// How long before lack of client response causes a timeout.
+const CLIENT_TIMEOUT: Duration = Duration::from_secs(10);
+
+pub async fn upload_file_ws(req: HttpRequest, stream: web::Payload, data: web::Data<AppState>) -> Result<HttpResponse, AppError> {
+    let (res, session, msg_stream) = actix_ws::handle(&req, stream)?;
+
+    // spawn websocket handler (and don't await it) so that the response is returned immediately
+    rt::spawn(upload_file_handler(data, session, msg_stream));
+
+    Ok(res)
+}
+
+async fn upload_file_handler(
+    data: web::Data<AppState>,
+    mut session: actix_ws::Session,
+    mut msg_stream: actix_ws::MessageStream,
+) {
+    let mut bytes = Cursor::new(vec![]);
+    let mut last_heartbeat = Instant::now();
+    let mut interval = interval(HEARTBEAT_INTERVAL);
+
+    let reason = loop {
+        let tick = interval.tick();
+        tokio::pin!(tick);
+
+        match future::select(msg_stream.next(), tick).await {
+            // received message from WebSocket client
+            Either::Left((Some(Ok(msg)), _)) => {
+                match msg {
+                    Message::Text(text) => {
+                        session.text(text).await.unwrap();
+                    }
+
+                    Message::Binary(bin) => {
+                        let mut pos = 0;  // notice the name of the file that will be written
+                        while pos < bin.len() {
+                            let bytes_written = bytes.write(&bin[pos..]).unwrap();
+                            pos += bytes_written
+                        };
+                        session.binary(bin).await.unwrap();
+                    }
+
+                    Message::Close(reason) => {
+                        break reason;
+                    }
+
+                    Message::Ping(bytes) => {
+                        last_heartbeat = Instant::now();
+                        let _ = session.pong(&bytes).await;
+                    }
+
+                    Message::Pong(_) => {
+                        last_heartbeat = Instant::now();
+                    }
+
+                    Message::Continuation(_) | Message::Nop => {}
+                };
+            }
+            Either::Left((Some(Err(_)), _)) => {
+                break None;
+            }
+            Either::Left((None, _)) => break None,
+            Either::Right((_inst, _)) => {
+                if Instant::now().duration_since(last_heartbeat) > CLIENT_TIMEOUT {
+                    break None;
+                }
+
+                let _ = session.ping(b"").await;
+            }
+        }
+    };
+    let _ = session.close(reason).await;
+
+    if !bytes.has_remaining() {
+        return;
+    }
+
+    let uid = bytes.get_i64();
+    let did = bytes.get_i64();
+
+    _upload_file(uid, did, &Uuid::new_v4().to_string(), &bytes.into_inner(), &data).await.unwrap();
+}

src/web_socket/mod.rs

@@ -0,0 +1 @@
+pub mod doc_info;