Store files by room

This commit is contained in:
Niels Andriesse 2021-04-27 15:51:53 +10:00
parent bf3771aab2
commit 1bbc521893
4 changed files with 57 additions and 31 deletions

View File

@ -125,7 +125,8 @@ pub fn get_all_rooms() -> Result<Response, Rejection> {
// Files
pub async fn store_file(
base64_encoded_bytes: &str, auth_token: Option<String>, pool: &storage::DatabaseConnectionPool,
room_id: Option<String>, base64_encoded_bytes: &str, auth_token: Option<String>,
pool: &storage::DatabaseConnectionPool,
) -> Result<Response, Rejection> {
// It'd be nice to use the UUID crate for the file ID, but clients want an integer ID
const UPPER_BOUND: u64 = 2u64.pow(53); // JS has trouble if we go higher than this
@ -166,7 +167,17 @@ pub async fn store_file(
}
};
// Write to file
let raw_path = format!("files/{}", &id);
// room_id is guaranteed to be present at this point
let room_id = room_id.unwrap();
match std::fs::create_dir(format!("files/{}", &room_id)) {
Ok(_) => (),
Err(e) => {
error!("Couldn't store file due to error: {}.", e);
return Err(warp::reject::custom(Error::DatabaseFailedInternally));
}
};
let raw_path = format!("files/{}/{}", &room_id, &id);
println!("mmmyess");
let path = Path::new(&raw_path);
let mut file = match File::create(path).await {
Ok(file) => file,
@ -193,7 +204,8 @@ pub async fn store_file(
}
pub async fn get_file(
id: u64, auth_token: Option<String>, pool: &storage::DatabaseConnectionPool,
room_id: Option<String>, id: u64, auth_token: Option<String>,
pool: &storage::DatabaseConnectionPool,
) -> Result<GenericStringResponse, Rejection> {
// Doesn't return a response directly for testing purposes
// Check authorization level if needed
@ -210,7 +222,8 @@ pub async fn get_file(
}
// Try to read the file
let mut bytes = vec![];
let raw_path = format!("files/{}", id);
// room_id is guaranteed to be present at this point
let raw_path = format!("files/{}/{}", room_id.unwrap(), id);
let path = Path::new(&raw_path);
let mut file = match File::open(path).await {
Ok(file) => file,

View File

@ -48,10 +48,14 @@ pub async fn handle_rpc_call(rpc_call: RpcCall) -> Result<Response, Rejection> {
};
// Get the auth token if possible
let auth_token = get_auth_token(&rpc_call);
// Get the room ID
let room_id = get_room_id(&rpc_call);
// Switch on the HTTP method
match rpc_call.method.as_ref() {
"GET" => return handle_get_request(rpc_call, &path, auth_token, query_params).await,
"POST" => return handle_post_request(rpc_call, &path, auth_token).await,
"GET" => {
return handle_get_request(room_id, rpc_call, &path, auth_token, query_params).await
}
"POST" => return handle_post_request(room_id, rpc_call, &path, auth_token).await,
"DELETE" => {
let pool = get_pool_for_room(&rpc_call)?;
return handle_delete_request(rpc_call, &path, auth_token, &pool).await;
@ -64,7 +68,7 @@ pub async fn handle_rpc_call(rpc_call: RpcCall) -> Result<Response, Rejection> {
}
async fn handle_get_request(
rpc_call: RpcCall, path: &str, auth_token: Option<String>,
room_id: Option<String>, rpc_call: RpcCall, path: &str, auth_token: Option<String>,
query_params: HashMap<String, String>,
) -> Result<Response, Rejection> {
// Handle routes that don't require authorization first
@ -110,7 +114,7 @@ async fn handle_get_request(
return Err(warp::reject::custom(Error::InvalidRpcCall));
}
};
return handlers::get_file(file_id, auth_token, &pool)
return handlers::get_file(room_id, file_id, auth_token, &pool)
.await
.map(|json| warp::reply::json(&json).into_response());
}
@ -168,7 +172,7 @@ async fn handle_get_request(
}
async fn handle_post_request(
rpc_call: RpcCall, path: &str, auth_token: Option<String>,
room_id: Option<String>, rpc_call: RpcCall, path: &str, auth_token: Option<String>,
) -> Result<Response, Rejection> {
// Handle routes that don't require authorization first
if path == "compact_poll" {
@ -203,7 +207,7 @@ async fn handle_post_request(
return Err(warp::reject::custom(Error::InvalidRpcCall));
}
};
return handlers::store_file(&json.file, auth_token, &pool).await;
return handlers::store_file(room_id, &json.file, auth_token, &pool).await;
}
// Check that the auth token is present
let auth_token = auth_token.ok_or(warp::reject::custom(Error::NoAuthToken))?;
@ -357,21 +361,13 @@ async fn handle_delete_request(
// Utilities
fn get_pool_for_room(rpc_call: &RpcCall) -> Result<storage::DatabaseConnectionPool, Rejection> {
let room_id: String;
match MODE {
// In file server mode we don't have a concept of rooms, but for convenience (i.e. so
// we can use the same database structure) we just always use the main room
Mode::FileServer => room_id = "main".to_string(),
Mode::OpenGroupServer => {
room_id = match get_room_id(&rpc_call) {
Some(room_id) => room_id,
None => {
warn!("Missing room ID.");
return Err(warp::reject::custom(Error::InvalidRpcCall));
}
};
let room_id = match get_room_id(&rpc_call) {
Some(room_id) => room_id,
None => {
warn!("Missing room ID.");
return Err(warp::reject::custom(Error::InvalidRpcCall));
}
}
};
return Ok(storage::pool_by_room_id(&room_id));
}
@ -383,10 +379,17 @@ fn get_auth_token(rpc_call: &RpcCall) -> Option<String> {
}
fn get_room_id(rpc_call: &RpcCall) -> Option<String> {
if rpc_call.headers.is_empty() {
return None;
match MODE {
// In file server mode we don't have a concept of rooms, but for convenience (i.e. so
// we can use the same database structure) we just always use the main room
Mode::FileServer => return Some("main".to_string()),
Mode::OpenGroupServer => {
if rpc_call.headers.is_empty() {
return None;
}
return rpc_call.headers.get("Room").map(|s| s.to_string());
}
}
return rpc_call.headers.get("Room").map(|s| s.to_string());
}
fn reject_if_file_server_mode(path: &str) -> Result<(), Rejection> {

View File

@ -267,7 +267,7 @@ pub async fn prune_files(file_expiration: i64) {
// Delete the files
let mut deleted_ids: Vec<String> = vec![];
for id in ids {
match fs::remove_file(format!("files/{}", id)) {
match fs::remove_file(format!("files/{}/{}", room, id)) {
Ok(_) => deleted_ids.push(id),
Err(e) => error!("Couldn't delete file due to error: {}.", e),
}

View File

@ -82,7 +82,14 @@ async fn test_file_handling() {
// Get an auth token
let (auth_token, _) = get_auth_token();
// Store the test file
handlers::store_file(TEST_FILE, Some(auth_token.clone()), &pool).await.unwrap();
handlers::store_file(
Some(test_room_id.to_string()),
TEST_FILE,
Some(auth_token.clone()),
&pool,
)
.await
.unwrap();
// Check that there's a file record
let conn = pool.get().unwrap();
let raw_query = format!("SELECT id FROM {}", storage::FILES_TABLE);
@ -91,13 +98,16 @@ async fn test_file_handling() {
let id = id_as_string.parse::<u64>().unwrap();
// Retrieve the file and check the content
let base64_encoded_file =
handlers::get_file(id, Some(auth_token.clone()), &pool).await.unwrap().result;
handlers::get_file(Some(test_room_id.to_string()), id, Some(auth_token.clone()), &pool)
.await
.unwrap()
.result;
assert_eq!(base64_encoded_file, TEST_FILE);
// Prune the file and check that it's gone
// Will evaluate to now + 60
storage::prune_files(-60).await;
// It should be gone now
fs::read(format!("files/{}", id)).unwrap_err();
fs::read(format!("files/{}/{}", test_room_id, id)).unwrap_err();
// Check that the file record is also gone
let conn = pool.get().unwrap();
let raw_query = format!("SELECT id FROM {}", storage::FILES_TABLE);