refactor: Use PublicKey and Signature as parameters

### Changes
- Implement `salvo_oapi::Extractible` and `salvo_oapi::EndpointArgRegister`
for `PublicKey`
- Implement `salvo_oapi::Extractible` and `salvo_oapi::ToParameters` for `Signature`

### Notes
I did not implement `salvo_oapi::ToParameters` for `PublicKey` because it will
not be used as an OpenAPI parameter. Instead, Salvo will register it using the
`EndpointArgRegister` trait.

Similarly, I did not implement `salvo_oapi::EndpointArgRegister` for `Signature`
because it will not be used as an argument in the endpoint. Instead, the
signature will be verified by the signature middleware, and we will only use it
as a parameter.

Reviewed-by: Amjad Alsharafi <me@amjad.alsharafi.dev>
Reviewed-on: #33
Fixes: #21
Helped-by: Amjad Alsharafi <me@amjad.alsharafi.dev>
Signed-off-by: Awiteb <a@4rs.nl>

Cargo.lock

@@ -1954,8 +1954,10 @@ dependencies = [
  "k256",
  "rand",
  "salvo-oapi",
+ "salvo_core",
  "sea-orm",
  "serde",
+ "serde_json",
  "sha2",
  "thiserror",
 ]

Cargo.toml

@@ -22,8 +22,9 @@ thiserror  = "1.0.61"
 log        = "0.4.21"
 logcall    = "0.1.9"
 chrono     = "0.4.38"
+serde_json = "1.0.117"
 sea-orm    = { version = "0.12.15", features = ["with-chrono", "macros"] }
-salvo_core = { version = "0.68.3", default-features = false, features = ["rustls"] }
+salvo_core = { version = "0.68.3", default-features = false }
 salvo-oapi = { version = "0.68.3", default-features = false, features = ["rapidoc","redoc","scalar","swagger-ui"] }
 
 [profile.release]

crates/oxidetalis/Cargo.toml

@@ -21,12 +21,12 @@ sea-orm               = { workspace = true }
 serde                 = { workspace = true }
 thiserror             = { workspace = true }
 chrono                = { workspace = true }
+serde_json            = { workspace = true }
 salvo                 = { version = "0.68.2", features = ["rustls", "affix", "logging", "oapi", "rate-limiter", "websocket"] }
 tokio                 = { version = "1.38.0", features = ["macros", "rt-multi-thread"] }
 uuid                  = { version = "1.9.1", default-features = false, features = ["v4"] }
 derive-new            = "0.6.0"
 pretty_env_logger     = "0.5.0"
-serde_json            = "1.0.117"
 once_cell             = "1.19.0"
 futures               = "0.3.30"
 rayon                 = "1.10.0"

crates/oxidetalis/src/middlewares/mod.rs

@@ -24,10 +24,8 @@ use salvo::{
     Response,
 };
 
-mod public_key;
 mod signature;
 
-pub use public_key::*;
 pub use signature::*;
 
 use crate::{routes::write_json_body, schemas::MessageSchema};

crates/oxidetalis/src/middlewares/public_key.rs

@@ -1,32 +0,0 @@
-// OxideTalis Messaging Protocol homeserver implementation
-// Copyright (C) 2024 Awiteb <a@4rs.nl>, OxideTalis Contributors
-//
-// This program is free software: you can redistribute it and/or modify
-// it under the terms of the GNU Affero General Public License as
-// published by the Free Software Foundation, either version 3 of the
-// License, or (at your option) any later version.
-//
-// This program is distributed in the hope that it will be useful,
-// but WITHOUT ANY WARRANTY; without even the implied warranty of
-// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-// GNU Affero General Public License for more details.
-//
-// You should have received a copy of the GNU Affero General Public License
-// along with this program. If not, see <https://gnu.org/licenses/agpl-3.0>.
-
-//! Request sender public key middleware.
-
-use salvo::{handler, http::StatusCode, FlowCtrl, Request, Response};
-
-use crate::utils;
-
-/// Middleware to check the public key of the request sender.
-///
-/// If the public key is valid, the request will be passed to the next handler.
-/// Otherwise, a 401 Unauthorized response will be returned.
-#[handler]
-pub async fn public_key_check(req: &mut Request, res: &mut Response, ctrl: &mut FlowCtrl) {
-    if let Err(err) = utils::extract_public_key(req) {
-        super::write_error(res, ctrl, err.to_string(), StatusCode::UNAUTHORIZED)
-    }
-}

crates/oxidetalis/src/middlewares/signature.rs

@@ -16,6 +16,7 @@
 
 //! Request signature middleware.
 
+use oxidetalis_core::types::{PublicKey, Signature};
 use salvo::{
     handler,
     http::{Body, StatusCode},
@@ -23,6 +24,7 @@ use salvo::{
     FlowCtrl,
     Request,
     Response,
+    Writer,
 };
 
 use crate::{extensions::DepotExt, utils};
@@ -37,6 +39,8 @@ pub async fn signature_check(
     res: &mut Response,
     depot: &mut Depot,
     ctrl: &mut FlowCtrl,
+    sender_public_key: PublicKey,
+    signature: Signature,
 ) {
     const UNAUTHORIZED: StatusCode = StatusCode::UNAUTHORIZED;
     let mut write_err =
@@ -54,22 +58,6 @@ pub async fn signature_check(
         }
     };
 
-    let signature = match utils::extract_signature(req) {
-        Ok(s) => s,
-        Err(err) => {
-            write_err(&err.to_string(), UNAUTHORIZED);
-            return;
-        }
-    };
-
-    let sender_public_key = match utils::extract_public_key(req) {
-        Ok(k) => k,
-        Err(err) => {
-            write_err(&err.to_string(), UNAUTHORIZED);
-            return;
-        }
-    };
-
     if !utils::is_valid_nonce(&signature, &depot.nonce_cache()).await
         || !depot.config().server.private_key.verify(
             data.as_bytes(),

crates/oxidetalis/src/routes/user.rs

@@ -17,7 +17,7 @@
 //! REST API endpoints for user management
 
 use oxidetalis_core::types::{PublicKey, Signature};
-use salvo::{http::StatusCode, oapi::endpoint, writing::Json, Depot, Request, Router, Writer};
+use salvo::{http::StatusCode, oapi::endpoint, writing::Json, Depot, Router, Writer};
 
 use super::{ApiError, ApiResult};
 use crate::{
@@ -26,7 +26,6 @@ use crate::{
     middlewares,
     parameters::Pagination,
     schemas::{BlackListedUser, EmptySchema, MessageSchema, WhiteListedUser},
-    utils,
 };
 
 /// (🔓) Register a user
@@ -38,22 +37,18 @@ use crate::{
     tags("User"),
     responses(
         (status_code = 201, description = "User registered"),
-        (status_code = 400, description = "The entered public key is already registered", content_type = "application/json", body = MessageSchema),
-        (status_code = 401, description = "The entered signature or public key is invalid", content_type = "application/json", body = MessageSchema),
+        (status_code = 400, description = "Invalid public key", content_type = "application/json", body = MessageSchema),
+        (status_code = 401, description = "Invalid signature", content_type = "application/json", body = MessageSchema),
         (status_code = 403, description = "Server registration is closed", content_type = "application/json", body = MessageSchema),
+        (status_code = 409, description = "The entered public key is already registered", content_type = "application/json", body = MessageSchema),
         (status_code = 429, description = "Too many requests", content_type = "application/json", body = MessageSchema),
         (status_code = 500, description = "Internal server error", content_type = "application/json", body = MessageSchema),
     ),
-    parameters(
-        ("X-OTMP-SIGNATURE" = Signature, Header, description = "Signature of the request"),
-        ("X-OTMP-PUBLIC"    = PublicKey, Header, description = "Public key of the sender"),
-    ),
+    parameters(Signature),
 )]
-pub async fn register(req: &Request, depot: &mut Depot) -> ApiResult<EmptySchema> {
+pub async fn register(public_key: PublicKey, depot: &mut Depot) -> ApiResult<EmptySchema> {
     let db = depot.db_conn();
     let config = depot.config();
-    let public_key =
-        utils::extract_public_key(req).expect("Public key should be checked in the middleware");
 
     if !db.users_exists_in_database().await? {
         db.register_user(&public_key, true).await?;
@@ -72,28 +67,22 @@ pub async fn register(req: &Request, depot: &mut Depot) -> ApiResult<EmptySchema
     tags("User"),
     responses(
         (status_code = 200, description = "Returns whitelisted users", content_type = "application/json", body = Vec<WhiteListedUser>),
-        (status_code = 400, description = "Wrong query parameter", content_type = "application/json", body = MessageSchema),
-        (status_code = 401, description = "The entered signature or public key is invalid", content_type = "application/json", body = MessageSchema),
+        (status_code = 400, description = "Invalid parameters or public key", content_type = "application/json", body = MessageSchema),
+        (status_code = 401, description = "Invalid signature", content_type = "application/json", body = MessageSchema),
         (status_code = 403, description = "Not registered user, must register first", content_type = "application/json", body = MessageSchema),
         (status_code = 429, description = "Too many requests", content_type = "application/json", body = MessageSchema),
         (status_code = 500, description = "Internal server error", content_type = "application/json", body = MessageSchema),
     ),
-    parameters(
-        ("X-OTMP-PUBLIC"    = PublicKey, Header, description = "Public key of the sender"),
-        ("X-OTMP-SIGNATURE" = Signature, Header, description = "Signature of the request"),
-    ),
+    parameters(Signature),
 )]
 async fn user_whitelist(
-    req: &mut Request,
     depot: &mut Depot,
     pagination: Pagination,
+    public_key: PublicKey,
 ) -> ApiResult<Json<Vec<WhiteListedUser>>> {
     let conn = depot.db_conn();
     let user = conn
-        .get_user_by_pubk(
-            &utils::extract_public_key(req)
-                .expect("Public key should be checked in the middleware"),
-        )
+        .get_user_by_pubk(&public_key)
         .await?
         .ok_or(ApiError::NotRegisteredUser)?;
     Ok(Json(
@@ -111,28 +100,22 @@ async fn user_whitelist(
     tags("User"),
     responses(
         (status_code = 200, description = "Returns blacklisted users", content_type = "application/json", body = Vec<BlackListedUser>),
-        (status_code = 400, description = "Wrong query parameter", content_type = "application/json", body = MessageSchema),
-        (status_code = 401, description = "The entered signature or public key is invalid", content_type = "application/json", body = MessageSchema),
+        (status_code = 400, description = "Invalid parameters or public key", content_type = "application/json", body = MessageSchema),
+        (status_code = 401, description = "Invalid signature", content_type = "application/json", body = MessageSchema),
         (status_code = 403, description = "Not registered user, must register first", content_type = "application/json", body = MessageSchema),
         (status_code = 429, description = "Too many requests", content_type = "application/json", body = MessageSchema),
         (status_code = 500, description = "Internal server error", content_type = "application/json", body = MessageSchema),
     ),
-    parameters(
-        ("X-OTMP-PUBLIC"    = PublicKey, Header, description = "Public key of the sender"),
-        ("X-OTMP-SIGNATURE" = Signature, Header, description = "Signature of the request"),
-    ),
+    parameters(Signature),
 )]
 async fn user_blacklist(
-    req: &mut Request,
     depot: &mut Depot,
     pagination: Pagination,
+    public_key: PublicKey,
 ) -> ApiResult<Json<Vec<BlackListedUser>>> {
     let conn = depot.db_conn();
     let user = conn
-        .get_user_by_pubk(
-            &utils::extract_public_key(req)
-                .expect("Public key should be checked in the middleware"),
-        )
+        .get_user_by_pubk(&public_key)
         .await?
         .ok_or(ApiError::NotRegisteredUser)?;
     Ok(Json(
@@ -150,6 +133,5 @@ pub fn route() -> Router {
         .push(Router::with_path("register").post(register))
         .push(Router::with_path("whitelist").get(user_whitelist))
         .push(Router::with_path("blacklist").get(user_blacklist))
-        .hoop(middlewares::public_key_check)
         .hoop(middlewares::signature_check)
 }

crates/oxidetalis/src/utils.rs

@@ -16,17 +16,10 @@
 
 //! Oxidetalis server utilities, utilities shared across the crate.
 
-use std::str::FromStr;
-
 use chrono::Utc;
 use logcall::logcall;
 use oxidetalis_config::Postgres;
-use oxidetalis_core::{
-    types::{PublicKey, Signature},
-    PUBLIC_KEY_HEADER,
-    SIGNATURE_HEADER,
-};
-use salvo::Request;
+use oxidetalis_core::types::Signature;
 
 use crate::nonce::NonceCache;
 
@@ -48,31 +41,3 @@ pub(crate) async fn is_valid_nonce(signature: &Signature, nonce_cache: &NonceCac
     let unused_nonce = new_timestamp && nonce_cache.add_nonce(signature.nonce()).await;
     new_timestamp && unused_nonce
 }
-
-/// Extract the sender public key from the request
-///
-/// Returns the public key of the sender extracted from the request, or the
-/// reason why it failed.
-pub(crate) fn extract_public_key(req: &Request) -> Result<PublicKey, String> {
-    req.headers()
-        .get(PUBLIC_KEY_HEADER)
-        .map(|v| {
-            PublicKey::from_str(v.to_str().map_err(|err| err.to_string())?)
-                .map_err(|err| err.to_string())
-        })
-        .ok_or_else(|| "The public key is missing".to_owned())?
-}
-
-/// Extract the signature from the request
-///
-/// Returns the signature extracted from the request, or the reason why it
-/// failed.
-pub(crate) fn extract_signature(req: &Request) -> Result<Signature, String> {
-    req.headers()
-        .get(SIGNATURE_HEADER)
-        .map(|v| {
-            Signature::from_str(v.to_str().map_err(|err| err.to_string())?)
-                .map_err(|err| err.to_string())
-        })
-        .ok_or_else(|| "The signature is missing".to_owned())?
-}

crates/oxidetalis/src/websocket/mod.rs

@@ -33,6 +33,7 @@ use salvo::{
     Request,
     Response,
     Router,
+    Writer,
 };
 use sea_orm::DatabaseConnection;
 use tokio::{sync::RwLock, task::spawn as tokio_spawn, time::sleep as tokio_sleep};
@@ -49,7 +50,6 @@ use crate::{
     extensions::{DepotExt, OnlineUsersExt},
     middlewares,
     nonce::NonceCache,
-    utils,
 };
 
 /// Online users type
@@ -96,12 +96,11 @@ impl SocketUserData {
 pub async fn user_connected(
     req: &mut Request,
     res: &mut Response,
+    public_key: PublicKey,
     depot: &Depot,
 ) -> Result<(), StatusError> {
     let nonce_cache = depot.nonce_cache();
     let db_conn = depot.db_conn();
-    let public_key =
-        utils::extract_public_key(req).expect("The public key was checked in the middleware");
     let shared_secret = depot.config().server.private_key.shared_secret(&public_key);
 
     WebSocketUpgrade::new()
@@ -259,5 +258,4 @@ pub fn route() -> Router {
     Router::new()
         .push(Router::with_path("chat").get(user_connected))
         .hoop(middlewares::signature_check)
-        .hoop(middlewares::public_key_check)
 }

crates/oxidetalis_core/Cargo.toml

@@ -16,6 +16,8 @@ thiserror  = { workspace = true }
 salvo-oapi = { workspace = true, optional = true }
 serde      = { workspace = true, optional = true }
 sea-orm    = { workspace = true, optional = true }
+salvo_core = { workspace = true, optional = true }
+serde_json = { workspace = true, optional = true }
 cbc        = { version = "0.1.2", features = ["alloc", "std"] }
 k256       = { version = "0.13.3", default-features = false, features = ["ecdh"] }
 rand       = { version = "0.8.5", default-features = false, features = ["std_rng", "std"] }
@@ -25,7 +27,7 @@ hmac       = "0.12.1"
 sha2       = "0.10.8"
 
 [features]
-openapi = ["dep:salvo-oapi"]
+openapi = ["dep:salvo-oapi", "dep:salvo_core", "dep:serde_json"]
 serde   = ["dep:serde"]
 sea-orm   = ["dep:sea-orm"]
 

crates/oxidetalis_core/src/types/cipher.rs

@@ -24,15 +24,6 @@
 use std::{fmt, str::FromStr};
 
 use base58::{FromBase58, ToBase58};
-#[cfg(feature = "openapi")]
-use salvo_oapi::{
-    schema::{
-        Schema as OapiSchema,
-        SchemaFormat as OapiSchemaFormat,
-        SchemaType as OapiSchemaType,
-    },
-    ToSchema,
-};
 
 use crate::cipher::{hmac_sha256, CipherError};
 
@@ -206,23 +197,3 @@ impl From<[u8; 56]> for Signature {
         }
     }
 }
-
-#[cfg(feature = "openapi")]
-impl ToSchema for PublicKey {
-    fn to_schema(_components: &mut salvo_oapi::Components) -> salvo_oapi::RefOr<OapiSchema> {
-        salvo_oapi::Object::new()
-            .schema_type(OapiSchemaType::String)
-            .format(OapiSchemaFormat::Custom("base58".to_owned()))
-            .into()
-    }
-}
-
-#[cfg(feature = "openapi")]
-impl ToSchema for Signature {
-    fn to_schema(_components: &mut salvo_oapi::Components) -> salvo_oapi::RefOr<OapiSchema> {
-        salvo_oapi::Object::new()
-            .schema_type(OapiSchemaType::String)
-            .format(OapiSchemaFormat::Custom("hex".to_owned()))
-            .into()
-    }
-}

crates/oxidetalis_core/src/types/impl_openapi.rs

@@ -0,0 +1,168 @@
+// OxideTalis Messaging Protocol homeserver core implementation
+// Copyright (C) 2024 Awiteb <a@4rs.nl>, OxideTalis Contributors
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+//! OpenAPI schema for some core types.
+
+use std::str::FromStr;
+
+use salvo_core::{extract::Metadata as ExtractMetadata, http::StatusError, Extractible, Request};
+use salvo_oapi::{
+    schema::{
+        Schema as OapiSchema,
+        SchemaFormat as OapiSchemaFormat,
+        SchemaType as OapiSchemaType,
+    },
+    Components as OapiComponents,
+    EndpointArgRegister,
+    Parameter,
+    ParameterIn,
+    Parameters,
+    ToParameters,
+    ToSchema,
+};
+
+use super::{PublicKey as CorePublicKey, Signature};
+
+impl ToSchema for CorePublicKey {
+    fn to_schema(_components: &mut salvo_oapi::Components) -> salvo_oapi::RefOr<OapiSchema> {
+        salvo_oapi::Object::new()
+            .name(crate::PUBLIC_KEY_HEADER)
+            .description("User's public key")
+            .schema_type(OapiSchemaType::String)
+            .format(OapiSchemaFormat::Custom("base58".to_owned()))
+            .required(crate::PUBLIC_KEY_HEADER)
+            // A 33-byte base58 string can be either 44 or 45 characters long
+            .example("rW8FMG5D75NVNJV3Wd498dEh65BgUuhwY1Yk5zYJPpRe".into())
+            .max_length(45)
+            .min_length(44)
+            .into()
+    }
+}
+
+impl ToSchema for Signature {
+    fn to_schema(_components: &mut salvo_oapi::Components) -> salvo_oapi::RefOr<OapiSchema> {
+        salvo_oapi::Object::new()
+            .name(crate::SIGNATURE_HEADER)
+            .description("Signature of the request")
+            .schema_type(OapiSchemaType::String)
+            .format(OapiSchemaFormat::Custom("hex".to_owned()))
+            .required(crate::SIGNATURE_HEADER)
+            // 56 bytes in hex (valid signature)
+            .example("0".repeat(112).into())
+            .max_length(112)
+            .min_length(112)
+            .into()
+    }
+}
+
+impl<'ex> Extractible<'ex> for CorePublicKey {
+    fn metadata() -> &'ex ExtractMetadata {
+        static METADATA: ExtractMetadata = ExtractMetadata::new("");
+        &METADATA
+    }
+
+    #[allow(refining_impl_trait)]
+    async fn extract(req: &'ex mut Request) -> Result<Self, StatusError> {
+        extract_header(req, crate::PUBLIC_KEY_HEADER).and_then(|public_key| {
+            CorePublicKey::from_str(public_key).map_err(|err| {
+                StatusError::bad_request()
+                    .brief("Invalid public key")
+                    .cause(err.to_string())
+            })
+        })
+    }
+
+    #[allow(refining_impl_trait)]
+    async fn extract_with_arg(req: &'ex mut Request, _: &str) -> Result<Self, StatusError> {
+        Self::extract(req).await
+    }
+}
+
+impl EndpointArgRegister for CorePublicKey {
+    fn register(components: &mut OapiComponents, operation: &mut salvo_oapi::Operation, _: &str) {
+        operation.parameters.insert(
+            Parameter::new(crate::PUBLIC_KEY_HEADER)
+                .parameter_in(ParameterIn::Header)
+                .required(true)
+                .description("User's public key")
+                .example("2BiUSWkJUy5bcdJB8qszq9K6a5EXVHvK41vQWZVkUBUM8".into())
+                .schema(CorePublicKey::to_schema(components)),
+        )
+    }
+}
+
+impl<'ex> Extractible<'ex> for Signature {
+    fn metadata() -> &'ex ExtractMetadata {
+        static METADATA: ExtractMetadata = ExtractMetadata::new("");
+        &METADATA
+    }
+
+    #[allow(refining_impl_trait)]
+    async fn extract(req: &'ex mut Request) -> Result<Self, StatusError> {
+        extract_header(req, crate::SIGNATURE_HEADER)
+            .and_then(|sig| {
+                Signature::from_str(sig).map_err(|err| {
+                    StatusError::unauthorized()
+                        .brief("Invalid signature")
+                        .cause(err.to_string())
+                })
+            })
+            .map_err(|err| {
+                StatusError::unauthorized().brief(err.brief).cause(
+                    err.cause
+                        .expect("The cause was set when we extract the header"),
+                )
+            })
+    }
+}
+
+impl ToParameters<'_> for Signature {
+    fn to_parameters(components: &mut OapiComponents) -> Parameters {
+        Parameters::new().parameter(
+            Parameter::new(crate::SIGNATURE_HEADER)
+                .parameter_in(ParameterIn::Header)
+                .required(true)
+                .description("Signature of the request")
+                .example("0".repeat(112).into())
+                .schema(Self::to_schema(components)),
+        )
+    }
+}
+
+fn extract_header<'req>(req: &'req Request, name: &str) -> Result<&'req str, StatusError> {
+    req.headers()
+        .get(name)
+        .map(|v| {
+            v.to_str().map_err(|_| {
+                StatusError::bad_request()
+                    .brief("Invalid header value")
+                    .cause("Header value must be a valid ascii string")
+            })
+        })
+        .transpose()?
+        .ok_or_else(|| {
+            StatusError::bad_request()
+                .brief(format!("Could not find {name} in headers"))
+                .cause(format!(
+                    "{name} is required to authenication and authorization"
+                ))
+        })
+}

crates/oxidetalis_core/src/types/mod.rs

@@ -22,6 +22,8 @@
 //! Oxidetalis server types
 
 mod cipher;
+#[cfg(feature = "openapi")]
+mod impl_openapi;
 #[cfg(feature = "sea-orm")]
 mod impl_sea_orm;
 #[cfg(feature = "serde")]