Add file storage & retrieval tests

This commit is contained in:
Niels Andriesse 2021-03-19 11:18:03 +11:00
parent 592c204747
commit 955192f5fe
6 changed files with 63 additions and 22 deletions

4
.gitignore vendored
View File

@ -1,4 +1,6 @@
/target /target
*.db *.db
.DS_Store .DS_Store
*.pem *.pem
/files
/rooms

View File

@ -33,7 +33,7 @@ pub async fn store_file(base64_encoded_bytes: &str, pool: &storage:: DatabaseCon
// Generate UUID // Generate UUID
let id = Uuid::new_v4(); let id = Uuid::new_v4();
let mut buffer = Uuid::encode_buffer(); let mut buffer = Uuid::encode_buffer();
let id = id.to_simple().encode_lower(&mut buffer); let id: String = id.to_simple().encode_lower(&mut buffer).to_string();
// Update the database // Update the database
// We do this * before * storing the actual file, so that in case something goes // We do this * before * storing the actual file, so that in case something goes
// wrong we're not left with files that'll never be pruned. // wrong we're not left with files that'll never be pruned.
@ -41,7 +41,7 @@ pub async fn store_file(base64_encoded_bytes: &str, pool: &storage:: DatabaseCon
let mut conn = pool.get().map_err(|_| Error::DatabaseFailedInternally)?; let mut conn = pool.get().map_err(|_| Error::DatabaseFailedInternally)?;
let tx = conn.transaction().map_err(|_| Error::DatabaseFailedInternally)?; let tx = conn.transaction().map_err(|_| Error::DatabaseFailedInternally)?;
let stmt = format!("INSERT INTO {} (id, timestamp) VALUES (?1, ?2)", storage::FILES_TABLE); let stmt = format!("INSERT INTO {} (id, timestamp) VALUES (?1, ?2)", storage::FILES_TABLE);
let _ = match tx.execute(&stmt, params![ &id, now ]) { let _ = match tx.execute(&stmt, params![ id, now ]) {
Ok(rows) => rows, Ok(rows) => rows,
Err(e) => { Err(e) => {
println!("Couldn't insert file record due to error: {}.", e); println!("Couldn't insert file record due to error: {}.", e);
@ -51,7 +51,7 @@ pub async fn store_file(base64_encoded_bytes: &str, pool: &storage:: DatabaseCon
tx.commit().map_err(|_| Error::DatabaseFailedInternally)?; tx.commit().map_err(|_| Error::DatabaseFailedInternally)?;
// Write to file // Write to file
let mut pos = 0; let mut pos = 0;
let mut buffer = match fs::File::create(&id) { let mut buffer = match fs::File::create(format!("files/{}", &id)) {
Ok(buffer) => buffer, Ok(buffer) => buffer,
Err(e) => { Err(e) => {
println!("Couldn't store file due to error: {}.", e); println!("Couldn't store file due to error: {}.", e);
@ -72,7 +72,7 @@ pub async fn store_file(base64_encoded_bytes: &str, pool: &storage:: DatabaseCon
return Ok(warp::reply::json(&id).into_response()); return Ok(warp::reply::json(&id).into_response());
} }
pub async fn get_file(id: &str, pool: &storage:: DatabaseConnectionPool) -> Result<Response, Rejection> { pub async fn get_file(id: &str) -> Result<String, Rejection> {
// Check that the ID is a valid UUID // Check that the ID is a valid UUID
match Uuid::parse_str(id) { match Uuid::parse_str(id) {
Ok(_) => (), Ok(_) => (),
@ -81,8 +81,6 @@ pub async fn get_file(id: &str, pool: &storage:: DatabaseConnectionPool) -> Resu
return Err(warp::reject::custom(Error::ValidationFailed)); return Err(warp::reject::custom(Error::ValidationFailed));
} }
}; };
// Get a database connection
let conn = pool.get().map_err(|_| Error::DatabaseFailedInternally)?;
// Try to read the file // Try to read the file
let bytes = match fs::read(format!("files/{}", id)) { let bytes = match fs::read(format!("files/{}", id)) {
Ok(bytes) => bytes, Ok(bytes) => bytes,
@ -94,7 +92,7 @@ pub async fn get_file(id: &str, pool: &storage:: DatabaseConnectionPool) -> Resu
// Base64 encode the result // Base64 encode the result
let base64_encoded_bytes = base64::encode(bytes); let base64_encoded_bytes = base64::encode(bytes);
// Return // Return
return Ok(warp::reply::json(&base64_encoded_bytes).into_response()); return Ok(base64_encoded_bytes);
} }
// Authentication // Authentication

View File

@ -40,5 +40,5 @@ async fn main() {
.key_path("tls_private_key.pem") .key_path("tls_private_key.pem")
.run(([0, 0, 0, 0], 443)); .run(([0, 0, 0, 0], 443));
// Keep futures alive // Keep futures alive
join!(prune_pending_tokens_future, prune_tokens_future, serve_routes_future); join!(prune_pending_tokens_future, prune_tokens_future, prune_files_future, serve_routes_future);
} }

View File

@ -1,7 +1,7 @@
use std::collections::HashMap; use std::collections::HashMap;
use serde::{Serialize, Deserialize}; use serde::{Serialize, Deserialize};
use warp::{Rejection, reply::Response}; use warp::{Rejection, reply::Reply, reply::Response};
use super::errors::Error; use super::errors::Error;
use super::handlers; use super::handlers;
@ -70,7 +70,7 @@ async fn handle_get_request(rpc_call: RpcCall, uri: http::Uri, pool: &storage::D
return Err(warp::reject::custom(Error::InvalidRpcCall)); return Err(warp::reject::custom(Error::InvalidRpcCall));
} }
let file_id = components[1]; let file_id = components[1];
return handlers::get_file(file_id, pool).await; return handlers::get_file(file_id).await.map(|json| warp::reply::json(&json).into_response());
} }
match uri.path() { match uri.path() {
"/messages" => return handlers::get_messages(query_options, pool).await, "/messages" => return handlers::get_messages(query_options, pool).await,

View File

@ -289,7 +289,7 @@ async fn prune_files() {
} }
// Remove the file records from the database (only for the files that were actually deleted) // Remove the file records from the database (only for the files that were actually deleted)
let stmt = format!("DELETE FROM {} WHERE id IN (?1)", FILES_TABLE); let stmt = format!("DELETE FROM {} WHERE id IN (?1)", FILES_TABLE);
match tx.execute(&stmt, params![ deleted_ids ]) { match tx.execute(&stmt, deleted_ids) {
Ok(_) => (), Ok(_) => (),
Err(e) => return println!("Couldn't prune files due to error: {}.", e) Err(e) => return println!("Couldn't prune files due to error: {}.", e)
}; };

File diff suppressed because one or more lines are too long