Update rocket to 0.5.0-rc1

Author: johan-bjareholt

Cargo.lock

@@ -15,3 +15,5 @@ aw-models = { path = "../aw-models" }
 [dev-dependencies]
 aw-datastore = { path = "../aw-datastore" }
 aw-server = { path = "../aw-server", default-features = false, features=[] }
+rocket = "0.5.0-rc.1"
+tokio-test = "*"

aw-client-rust/Cargo.toml

@@ -2,7 +2,9 @@ extern crate aw_client_rust;
 extern crate aw_datastore;
 extern crate aw_server;
 extern crate chrono;
+extern crate rocket;
 extern crate serde_json;
+extern crate tokio_test;
 
 #[cfg(test)]
 mod test {
@@ -13,6 +15,7 @@ mod test {
     use std::path::PathBuf;
     use std::sync::Mutex;
     use std::thread;
+    use tokio_test::block_on;
 
     // A random port, but still not guaranteed to not be bound
     // FIXME: Bind to a port that is free for certain and use that for the client instead
@@ -37,9 +40,7 @@ mod test {
         }
     }
 
-    fn setup_testserver() -> () {
-        // Start testserver and wait 10s for it to start up
-        // TODO: Properly shutdown
+    fn setup_testserver() -> rocket::Shutdown {
         use aw_server::endpoints::ServerState;
         let state = ServerState {
             datastore: Mutex::new(aw_datastore::Datastore::new_in_memory(false)),
@@ -49,10 +50,14 @@ mod test {
         let mut aw_config = aw_server::config::AWConfig::default();
         aw_config.port = PORT;
         let server = aw_server::endpoints::build_rocket(state, aw_config);
+        let server = block_on(server.ignite()).unwrap();
+        let shutdown_handler = server.shutdown();
 
         thread::spawn(move || {
-            server.launch();
+            block_on(server.launch()).unwrap();
         });
+
+        shutdown_handler
     }
 
     #[test]
@@ -62,7 +67,7 @@ mod test {
         let clientname = "aw-client-rust-test";
         let client: AwClient = AwClient::new(ip, &port, clientname);
 
-        setup_testserver();
+        let shutdown_handler = setup_testserver();
 
         wait_for_server(20, &client);
 
@@ -113,5 +118,7 @@ mod test {
         assert_eq!(count, 0);
 
         client.delete_bucket(&bucketname).unwrap();
+
+        shutdown_handler.notify();
     }
 }

aw-client-rust/tests/test.rs

@@ -14,9 +14,9 @@ name = "aw-server"
 path = "src/main.rs"
 
 [dependencies]
-rocket = "0.4"
-rocket_contrib = { version = "*", default-features = false, features = ["json"] }
-rocket_cors = "0.5"
+rocket = { version = "0.5.0-rc.1", features = ["json"] }
+# TODO: Once rocket_cors has a version for rocket 0.5, use that instead
+rocket_cors = { git = "https://github.com/lawliet89/rocket_cors", rev = "a062933" }
 multipart = { version = "0.18", default-features = false, features = ["server"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"

aw-server/Cargo.toml

@@ -12,6 +12,7 @@ use crate::dirs;
 
 use android_logger::Config;
 use log::Level;
+use rocket::serde::json::json;
 
 #[no_mangle]
 pub extern "C" fn rust_greeting(to: *const c_char) -> *mut c_char {

aw-server/src/android/mod.rs

@@ -1,7 +1,8 @@
 use std::fs::File;
 use std::io::{Read, Write};
 
-use rocket::config::{Config, Environment, Limits};
+use rocket::config::Config;
+use rocket::data::{Limits, ToByteUnit};
 use serde::{Deserialize, Serialize};
 
 use crate::dirs;
@@ -42,21 +43,23 @@ impl Default for AWConfig {
 
 impl AWConfig {
     pub fn to_rocket_config(&self) -> rocket::Config {
-        let env = if self.testing {
-            Environment::Production
+        let mut config = if self.testing {
+            Config::release_default()
         } else {
-            Environment::Development
+            Config::debug_default()
         };
+
         // Needed for bucket imports
-        let limits = Limits::new().limit("json", 1_000_000_000);
-
-        Config::build(env)
-            .address(self.address.clone())
-            .port(self.port)
-            .keep_alive(0)
-            .limits(limits)
-            .finalize()
-            .unwrap()
+        let limits = Limits::default()
+            .limit("json", 1000u64.megabytes())
+            .limit("data-form", 1000u64.megabytes());
+
+        config.address = self.address.parse().unwrap();
+        config.port = self.port;
+        config.keep_alive = 0;
+        config.limits = limits;
+
+        config
     }
 }
 

aw-server/src/config.rs

@@ -1,7 +1,6 @@
 use std::collections::HashMap;
-use std::io::Cursor;
 
-use rocket_contrib::json::Json;
+use rocket::serde::json::Json;
 
 use chrono::DateTime;
 use chrono::Utc;
@@ -11,16 +10,15 @@ use aw_models::BucketsExport;
 use aw_models::Event;
 use aw_models::TryVec;
 
-use rocket::http::Header;
 use rocket::http::Status;
-use rocket::response::Response;
 use rocket::State;
 
+use crate::endpoints::util::BucketsExportRocket;
 use crate::endpoints::{HttpErrorJson, ServerState};
 
 #[get("/")]
 pub fn buckets_get(
-    state: State<ServerState>,
+    state: &State<ServerState>,
 ) -> Result<Json<HashMap<String, Bucket>>, HttpErrorJson> {
     let datastore = endpoints_get_lock!(state.datastore);
     match datastore.get_buckets() {
@@ -32,7 +30,7 @@ pub fn buckets_get(
 #[get("/<bucket_id>")]
 pub fn bucket_get(
     bucket_id: String,
-    state: State<ServerState>,
+    state: &State<ServerState>,
 ) -> Result<Json<Bucket>, HttpErrorJson> {
     let datastore = endpoints_get_lock!(state.datastore);
     match datastore.get_bucket(&bucket_id) {
@@ -45,7 +43,7 @@ pub fn bucket_get(
 pub fn bucket_new(
     bucket_id: String,
     message: Json<Bucket>,
-    state: State<ServerState>,
+    state: &State<ServerState>,
 ) -> Result<(), HttpErrorJson> {
     let mut bucket = message.into_inner();
     if bucket.id != bucket_id {
@@ -65,7 +63,7 @@ pub fn bucket_events_get(
     start: Option<String>,
     end: Option<String>,
     limit: Option<u64>,
-    state: State<ServerState>,
+    state: &State<ServerState>,
 ) -> Result<Json<Vec<Event>>, HttpErrorJson> {
     let starttime: Option<DateTime<Utc>> = match start {
         Some(dt_str) => match DateTime::parse_from_rfc3339(&dt_str) {
@@ -107,7 +105,7 @@ pub fn bucket_events_get(
 pub fn bucket_events_create(
     bucket_id: String,
     events: Json<Vec<Event>>,
-    state: State<ServerState>,
+    state: &State<ServerState>,
 ) -> Result<Json<Vec<Event>>, HttpErrorJson> {
     let datastore = endpoints_get_lock!(state.datastore);
     let res = datastore.insert_events(&bucket_id, &events);
@@ -126,7 +124,7 @@ pub fn bucket_events_heartbeat(
     bucket_id: String,
     heartbeat_json: Json<Event>,
     pulsetime: f64,
-    state: State<ServerState>,
+    state: &State<ServerState>,
 ) -> Result<Json<Event>, HttpErrorJson> {
     let heartbeat = heartbeat_json.into_inner();
     let datastore = endpoints_get_lock!(state.datastore);
@@ -139,7 +137,7 @@ pub fn bucket_events_heartbeat(
 #[get("/<bucket_id>/events/count")]
 pub fn bucket_event_count(
     bucket_id: String,
-    state: State<ServerState>,
+    state: &State<ServerState>,
 ) -> Result<Json<u64>, HttpErrorJson> {
     let datastore = endpoints_get_lock!(state.datastore);
     let res = datastore.get_event_count(&bucket_id, None, None);
@@ -153,7 +151,7 @@ pub fn bucket_event_count(
 pub fn bucket_events_delete_by_id(
     bucket_id: String,
     event_id: i64,
-    state: State<ServerState>,
+    state: &State<ServerState>,
 ) -> Result<(), HttpErrorJson> {
     let datastore = endpoints_get_lock!(state.datastore);
     match datastore.delete_events_by_id(&bucket_id, vec![event_id]) {
@@ -165,8 +163,8 @@ pub fn bucket_events_delete_by_id(
 #[get("/<bucket_id>/export")]
 pub fn bucket_export(
     bucket_id: String,
-    state: State<ServerState>,
-) -> Result<Response, HttpErrorJson> {
+    state: &State<ServerState>,
+) -> Result<BucketsExportRocket, HttpErrorJson> {
     let datastore = endpoints_get_lock!(state.datastore);
     let mut export = BucketsExport {
         buckets: HashMap::new(),
@@ -181,20 +179,12 @@ pub fn bucket_export(
         .expect("Failed to get events for bucket");
     bucket.events = Some(TryVec::new(events));
     export.buckets.insert(bucket_id.clone(), bucket);
-    let filename = format!("aw-bucket-export_{}.json", bucket_id);
-
-    let header_content = format!("attachment; filename={}", filename);
-    Ok(Response::build()
-        .status(Status::Ok)
-        .header(Header::new("Content-Disposition", header_content))
-        .sized_body(Cursor::new(
-            serde_json::to_string(&export).expect("Failed to serialize"),
-        ))
-        .finalize())
+
+    Ok(export.into())
 }
 
 #[delete("/<bucket_id>")]
-pub fn bucket_delete(bucket_id: String, state: State<ServerState>) -> Result<(), HttpErrorJson> {
+pub fn bucket_delete(bucket_id: String, state: &State<ServerState>) -> Result<(), HttpErrorJson> {
     let datastore = endpoints_get_lock!(state.datastore);
     match datastore.delete_bucket(&bucket_id) {
         Ok(_) => Ok(()),

aw-server/src/endpoints/bucket.rs

@@ -1,18 +1,15 @@
 use std::collections::HashMap;
-use std::io::Cursor;
 
-use rocket::http::Header;
-use rocket::http::Status;
-use rocket::response::Response;
 use rocket::State;
 
 use aw_models::BucketsExport;
 use aw_models::TryVec;
 
+use crate::endpoints::util::BucketsExportRocket;
 use crate::endpoints::{HttpErrorJson, ServerState};
 
 #[get("/")]
-pub fn buckets_export(state: State<ServerState>) -> Result<Response, HttpErrorJson> {
+pub fn buckets_export(state: &State<ServerState>) -> Result<BucketsExportRocket, HttpErrorJson> {
     let datastore = endpoints_get_lock!(state.datastore);
     let mut export = BucketsExport {
         buckets: HashMap::new(),
@@ -30,14 +27,5 @@ pub fn buckets_export(state: State<ServerState>) -> Result<Response, HttpErrorJs
         export.buckets.insert(bid, bucket);
     }
 
-    Ok(Response::build()
-        .status(Status::Ok)
-        .header(Header::new(
-            "Content-Disposition",
-            "attachment; filename=aw-buckets-export.json",
-        ))
-        .sized_body(Cursor::new(
-            serde_json::to_string(&export).expect("Failed to serialize"),
-        ))
-        .finalize())
+    Ok(export.into())
 }

aw-server/src/endpoints/export.rs

@@ -7,9 +7,9 @@
 //!
 //! [1]: https://github.com/ActivityWatch/activitywatch/security/advisories/GHSA-v9fg-6g9j-h4x4
 use rocket::fairing::Fairing;
-use rocket::handler::Outcome;
 use rocket::http::uri::Origin;
 use rocket::http::{Method, Status};
+use rocket::route::Outcome;
 use rocket::{Data, Request, Rocket, Route};
 
 use crate::config::AWConfig;
@@ -29,15 +29,25 @@ impl HostCheck {
     }
 }
 
-/// Route for HostCheck Fairing error
-fn fairing_error_route<'r>(req: &'r Request<'_>, _: Data) -> Outcome<'r> {
-    let err = HttpErrorJson::new(Status::BadRequest, "Host header is invalid".to_string());
-    Outcome::from(req, err)
+/// Create a `Handler` for Fairing error handling
+#[derive(Clone)]
+struct FairingErrorRoute {}
+
+#[rocket::async_trait]
+impl rocket::route::Handler for FairingErrorRoute {
+    async fn handle<'r>(
+        &self,
+        request: &'r Request<'_>,
+        _: rocket::Data<'r>,
+    ) -> rocket::route::Outcome<'r> {
+        let err = HttpErrorJson::new(Status::BadRequest, "Host header is invalid".to_string());
+        Outcome::from(request, err)
+    }
 }
 
 /// Create a new `Route` for Fairing handling
 fn fairing_route() -> Route {
-    Route::ranked(1, Method::Get, "/", fairing_error_route)
+    Route::ranked(1, Method::Get, "/", FairingErrorRoute {})
 }
 
 fn redirect_bad_request(request: &mut Request) {
@@ -47,15 +57,16 @@ fn redirect_bad_request(request: &mut Request) {
     request.set_uri(origin);
 }
 
+#[rocket::async_trait]
 impl Fairing for HostCheck {
     fn info(&self) -> rocket::fairing::Info {
         rocket::fairing::Info {
             name: "HostCheck",
-            kind: rocket::fairing::Kind::Attach | rocket::fairing::Kind::Request,
+            kind: rocket::fairing::Kind::Ignite | rocket::fairing::Kind::Request,
         }
     }
 
-    fn on_attach(&self, rocket: Rocket) -> Result<Rocket, Rocket> {
+    async fn on_ignite(&self, rocket: Rocket<rocket::Build>) -> rocket::fairing::Result {
         match self.validate {
             true => Ok(rocket.mount(FAIRING_ROUTE_BASE, vec![fairing_route()])),
             false => {
@@ -65,7 +76,7 @@ impl Fairing for HostCheck {
         }
     }
 
-    fn on_request(&self, request: &mut Request, _: &Data) {
+    async fn on_request(&self, request: &mut Request<'_>, _: &mut Data<'_>) {
         if !self.validate {
             // host header check is disabled
             return;
@@ -112,7 +123,7 @@ mod tests {
     use crate::config::AWConfig;
     use crate::endpoints;
 
-    fn setup_testserver(address: String) -> Rocket {
+    fn setup_testserver(address: String) -> Rocket<rocket::Build> {
         let state = endpoints::ServerState {
             datastore: Mutex::new(aw_datastore::Datastore::new_in_memory(false)),
             asset_path: PathBuf::from("aw-webui/dist"),
@@ -126,7 +137,7 @@ mod tests {
     #[test]
     fn test_public_address() {
         let server = setup_testserver("0.0.0.0".to_string());
-        let client = rocket::local::Client::new(server).expect("valid instance");
+        let client = rocket::local::blocking::Client::tracked(server).expect("valid instance");
 
         // When a public address is used, request should always pass, regardless
         // if the Host header is missing
@@ -140,7 +151,7 @@ mod tests {
     #[test]
     fn test_localhost_address() {
         let server = setup_testserver("127.0.0.1".to_string());
-        let client = rocket::local::Client::new(server).expect("valid instance");
+        let client = rocket::local::blocking::Client::tracked(server).expect("valid instance");
 
         // If Host header is missing we should get a BadRequest
         let res = client