Code cleanup

Author: sandhose

.rustfmt.toml

@@ -0,0 +1,6 @@
+max_width = 100
+comment_width = 80
+wrap_comments = true
+imports_granularity = "Crate"
+use_small_heuristics = "Default"
+group_imports = "StdExternalCrate"

server/src/main.rs

@@ -16,13 +16,15 @@
 #![deny(clippy::all)]
 #![warn(clippy::pedantic)]
 
-use bytesize::ByteSize;
-use clap::Parser;
 use std::{
     net::{IpAddr, Ipv4Addr, SocketAddr},
     time::Duration,
 };
 
+use bytesize::ByteSize;
+use clap::Parser;
+use matrix_http_rendezvous::{DEFAULT_MAX_BYTES, DEFAULT_MAX_ENTRIES, DEFAULT_TTL};
+
 #[derive(Parser)]
 struct Options {
     /// Address on which to listen
@@ -38,18 +40,19 @@ struct Options {
     prefix: Option<String>,
 
     /// Time to live of entries, in seconds
-    #[arg(short, long, default_value_t = Duration::from_secs(60).into())]
+    #[arg(short, long, default_value_t = DEFAULT_TTL.into())]
     ttl: humantime::Duration,
 
     /// Maximum number of entries to store
-    #[arg(short, long, default_value_t = 10000)]
+    #[arg(short, long, default_value_t = DEFAULT_MAX_ENTRIES)]
     capacity: usize,
 
     /// Maximum payload size, in bytes
-    #[arg(short, long, default_value = "4KiB")]
+    #[arg(short, long, default_value_t = ByteSize(DEFAULT_MAX_BYTES as u64))]
     max_bytes: ByteSize,
 
-    /// Set this flag to test how much memory the server might use with a sessions map fully loaded
+    /// Set this flag to test how much memory the server might use with a
+    /// sessions map fully loaded
     #[arg(long)]
     mem_check: bool,
 }

src/handlers.rs

@@ -0,0 +1,140 @@
+// Copyright 2022 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use axum::{
+    body::HttpBody,
+    extract::{DefaultBodyLimit, Path, State},
+    http::{
+        header::{CONTENT_TYPE, ETAG, IF_MATCH, IF_NONE_MATCH, LOCATION},
+        StatusCode,
+    },
+    response::{IntoResponse, Response},
+    routing::{get, post},
+    Router, TypedHeader,
+};
+use bytes::Bytes;
+use headers::{ContentType, HeaderName, HeaderValue, IfMatch, IfNoneMatch};
+use tower_http::{
+    cors::{Any, CorsLayer},
+    limit::RequestBodyLimitLayer,
+    set_header::SetResponseHeaderLayer,
+};
+use ulid::Ulid;
+
+use crate::Sessions;
+
+async fn new_session(
+    State(sessions): State<Sessions>,
+    content_type: Option<TypedHeader<ContentType>>,
+    payload: Bytes,
+) -> impl IntoResponse {
+    let content_type =
+        content_type.map_or(mime::APPLICATION_OCTET_STREAM, |TypedHeader(c)| c.into());
+    let (id, session) = sessions.new_session(payload, content_type).await;
+    let headers = session.typed_headers();
+
+    let location = id.to_string();
+    let additional_headers = [(LOCATION, location)];
+    (StatusCode::CREATED, headers, additional_headers)
+}
+
+async fn delete_session(State(sessions): State<Sessions>, Path(id): Path<Ulid>) -> StatusCode {
+    if sessions.delete_session(id).await {
+        StatusCode::NO_CONTENT
+    } else {
+        StatusCode::NOT_FOUND
+    }
+}
+
+async fn update_session(
+    State(sessions): State<Sessions>,
+    Path(id): Path<Ulid>,
+    content_type: Option<TypedHeader<ContentType>>,
+    if_match: Option<TypedHeader<IfMatch>>,
+    payload: Bytes,
+) -> Response {
+    if let Some(mut session) = sessions.get_session_mut(id).await {
+        if let Some(TypedHeader(if_match)) = if_match {
+            if !if_match.precondition_passes(&session.etag()) {
+                return (StatusCode::PRECONDITION_FAILED, session.typed_headers()).into_response();
+            }
+        }
+
+        let content_type =
+            content_type.map_or(mime::APPLICATION_OCTET_STREAM, |TypedHeader(c)| c.into());
+
+        session.update(payload, content_type);
+        (StatusCode::ACCEPTED, session.typed_headers()).into_response()
+    } else {
+        StatusCode::NOT_FOUND.into_response()
+    }
+}
+
+async fn get_session(
+    State(sessions): State<Sessions>,
+    Path(id): Path<Ulid>,
+    if_none_match: Option<TypedHeader<IfNoneMatch>>,
+) -> Response {
+    let session = if let Some(session) = sessions.get_session(id).await {
+        session
+    } else {
+        return StatusCode::NOT_FOUND.into_response();
+    };
+
+    if let Some(TypedHeader(if_none_match)) = if_none_match {
+        if !if_none_match.precondition_passes(&session.etag()) {
+            return (StatusCode::NOT_MODIFIED, session.typed_headers()).into_response();
+        }
+    }
+
+    (
+        StatusCode::OK,
+        session.typed_headers(),
+        TypedHeader(session.content_type()),
+        session.data(),
+    )
+        .into_response()
+}
+
+#[must_use]
+#[allow(clippy::trait_duplication_in_bounds)]
+pub fn router<B>(prefix: &str, sessions: Sessions, max_bytes: usize) -> Router<(), B>
+where
+    B: HttpBody + Send + 'static,
+    <B as HttpBody>::Data: Send,
+    <B as HttpBody>::Error: std::error::Error + Send + Sync,
+{
+    let router = Router::with_state(sessions)
+        .route("/", post(new_session))
+        .route(
+            "/:id",
+            get(get_session).put(update_session).delete(delete_session),
+        )
+        .layer(DefaultBodyLimit::disable())
+        .layer(RequestBodyLimitLayer::new(max_bytes))
+        .layer(SetResponseHeaderLayer::if_not_present(
+            HeaderName::from_static("x-max-bytes"),
+            HeaderValue::from_str(&max_bytes.to_string())
+                .expect("Could not construct x-max-bytes header value"),
+        ))
+        .layer(
+            CorsLayer::new()
+                .allow_origin(Any)
+                .allow_methods(Any)
+                .allow_headers([CONTENT_TYPE, IF_MATCH, IF_NONE_MATCH])
+                .expose_headers([ETAG, LOCATION, HeaderName::from_static("x-max-bytes")]),
+        );
+
+    Router::new().nest(prefix, router)
+}