More error handling improvements

compose.yaml

@@ -9,6 +9,7 @@ services:
       - 5432:5432
     volumes:
       - postgres_data:/var/lib/postgresql/data
+    restart: unless-stopped
 
 volumes:
   postgres_data:

src/auth.rs

@@ -2,8 +2,7 @@ use std::{array::TryFromSliceError, sync::LazyLock};
 
 use axum::{
     extract::{FromRef, FromRequestParts},
-    http::{request::Parts, StatusCode},
-    response::IntoResponse,
+    http::request::Parts,
     RequestPartsExt,
 };
 use axum_extra::{
@@ -16,7 +15,7 @@ use rand::{rngs::OsRng, RngCore};
 use serde::{Deserialize, Serialize};
 use subtle::ConstantTimeEq;
 
-use crate::{db, Pool};
+use crate::prelude::*;
 
 pub const HASH_LENGTH: usize = 64;
 pub const SALT_LENGTH: usize = 64;
@@ -122,6 +121,8 @@ pub struct Claims {
     pub exp: i64,
 }
 
+const JWT_ALGORITHM: jsonwebtoken::Algorithm = jsonwebtoken::Algorithm::HS256;
+
 impl Claims {
     pub fn new(user_id: i32) -> Self {
         Self {
@@ -130,13 +131,9 @@ impl Claims {
         }
     }
 
-    pub fn encode(self) -> Result<Token, Error> {
-        let access_token = encode(
-            &Header::new(jsonwebtoken::Algorithm::HS256),
-            &self,
-            &KEYS.encoding_key,
-        )
-        .map_err(|_| Error::TokenCreation)?;
+    pub fn encode(self) -> Result<Token, GeneralError> {
+        let access_token = encode(&Header::new(JWT_ALGORITHM), &self, &KEYS.encoding_key)
+            .handle_internal("Token creation error")?;
         let token = Token {
             access_token,
             token_type: "Bearer",
@@ -145,51 +142,40 @@ impl Claims {
     }
 }
 
-#[derive(Debug)]
-pub enum Error {
-    WrongCredentials,
-    TokenCreation,
-    Validation,
-    InvalidToken,
-}
-
-impl IntoResponse for Error {
-    fn into_response(self) -> axum::response::Response {
-        let (status, error_message) = match self {
-            Error::WrongCredentials => (StatusCode::UNAUTHORIZED, "Wrong credentials"),
-            Error::TokenCreation => (StatusCode::INTERNAL_SERVER_ERROR, "Token creation error"),
-            Error::Validation => (StatusCode::INTERNAL_SERVER_ERROR, "Token validation error"),
-            Error::InvalidToken => (StatusCode::BAD_REQUEST, "Invalid token"),
-        };
-        (status, error_message).into_response()
-    }
-}
-
 #[axum::async_trait]
 impl<T> FromRequestParts<T> for Claims
 where
     Pool: FromRef<T>,
     T: Sync,
 {
-    type Rejection = Error;
+    type Rejection = GeneralError;
 
     async fn from_request_parts(parts: &mut Parts, state: &T) -> Result<Self, Self::Rejection> {
+        const INVALID_TOKEN: GeneralError =
+            GeneralError::const_message(StatusCode::UNAUTHORIZED, "Invalid token");
+
         let pool = Pool::from_ref(state);
         let TypedHeader(Authorization(bearer)) = parts
             .extract::<TypedHeader<Authorization<Bearer>>>()
             .await
-            .map_err(|_| Error::InvalidToken)?;
-        let claims: Claims = decode(bearer.token(), &KEYS.decoding_key, &Validation::default())
-            .map_err(|_| Error::InvalidToken)?
+            .map_err(|_| INVALID_TOKEN)?;
+
+        let claims: Claims = decode(
+            bearer.token(),
+            &KEYS.decoding_key,
+            &Validation::new(JWT_ALGORITHM),
+        )
+        .map_err(|_| INVALID_TOKEN)?
         .claims;
-        match db::users::exists(claims.user_id, &pool).await {
-            Ok(true) => Ok(claims),
-            Ok(false) => Err(Error::WrongCredentials),
-            Err(err) => {
-                tracing::error!(%err);
-                Err(Error::Validation)
-            }
-        }
+
+        db::users::exists(claims.user_id, &pool)
+            .await
+            .handle_internal("Token validation error")?
+            .then_some(claims)
+            .ok_or(GeneralError::const_message(
+                StatusCode::UNAUTHORIZED,
+                "Wrong credentials",
+            ))
     }
 }
 

src/db/permissions.rs

@@ -38,12 +38,7 @@ impl PermissionType {
     }
 
     fn can_read_guard(self) -> GeneralResult<()> {
-        if !self.can_read() {
-            return Err(GeneralError::message(
-                StatusCode::NOT_FOUND,
-                "Item not found",
-            ));
-        }
+        self.can_read().then_some(()).item_not_found()?;
         Ok(())
     }
 

src/endpoints/file/download.rs

@@ -20,7 +20,7 @@ pub async fn download(
     let mut name = db::file::get_name(params.file_id, &state.pool)
         .await
         .handle_internal("Error getting file info")?
-        .ok_or_else(GeneralError::item_not_found)?;
+        .item_not_found()?;
     name = name
         .chars()
         .fold(String::with_capacity(name.len()), |mut result, char| {

src/endpoints/file/modify.rs

@@ -36,7 +36,7 @@ pub async fn modify(
         .write(params.file_id)
         .await
         .handle_internal("Error writing to the file")?
-        .ok_or_else(GeneralError::item_not_found)?;
+        .item_not_found()?;
 
     let (hash, size) = crate::FileStorage::write_to_file(&mut file, &mut field)
         .await

src/endpoints/folder/get_structure.rs

@@ -29,12 +29,12 @@ pub async fn structure(
     let folder_id = db::folder::process_id(params.folder_id, claims.user_id, &pool)
         .await
         .handle_internal("Error processing id")?
-        .ok_or_else(GeneralError::item_not_found)?;
+        .item_not_found()?;
 
     let folder = db::folder::get_by_id(folder_id, &pool)
         .await
         .handle_internal("Error getting folder info")?
-        .ok_or_else(GeneralError::item_not_found)?;
+        .item_not_found()?;
 
     let mut response: FolderStructure = folder.into();
     let mut stack = vec![&mut response];

src/endpoints/permissions/set.rs

@@ -29,7 +29,7 @@ pub async fn set(
     let folder_info = db::folder::get_by_id(params.folder_id, &pool)
         .await
         .handle_internal("Error getting folder info")?
-        .ok_or_else(GeneralError::item_not_found)?;
+        .item_not_found()?;
     if folder_info.owner_id == params.user_id {
         return Err(GeneralError::message(
             StatusCode::BAD_REQUEST,

src/endpoints/users/login.rs

@@ -1,7 +1,7 @@
 use axum::Form;
 
 use crate::{
-    auth::{authenticate_user, Error, Token},
+    auth::{authenticate_user, Token},
     prelude::*,
 };
 
@@ -14,10 +14,13 @@ pub struct Params {
 pub async fn login(
     State(pool): State<Pool>,
     Form(payload): Form<Params>,
-) -> Result<Json<Token>, Error> {
+) -> GeneralResult<Json<Token>> {
     let user_id = authenticate_user(&payload.username, &payload.password, &pool)
         .await
-        .map_err(|_| Error::WrongCredentials)?
-        .ok_or(Error::WrongCredentials)?;
+        .handle_internal("Error getting user from database")?
+        .handle(
+            StatusCode::NOT_FOUND,
+            "User with this name and password doesn't exist",
+        )?;
     Claims::new(user_id).encode().map(Json)
 }

src/endpoints/users/put.rs

@@ -15,7 +15,7 @@ pub async fn put(
     claims: Claims,
     Json(params): Json<Params>,
 ) -> GeneralResult<Json<db::users::UserInfo>> {
-    params.validate().map_err(GeneralError::validation)?;
+    params.validate().handle_validation()?;
     db::users::update(claims.user_id, &params.username, &params.email, &pool)
         .await
         .handle_internal("Error updating the user")

src/endpoints/users/register.rs

@@ -1,10 +1,9 @@
 use axum::Form;
-use axum_extra::either::Either;
 use itertools::Itertools;
 use validator::{Validate, ValidationError};
 
 use crate::{
-    auth::{Error, HashedBytes, Token},
+    auth::{HashedBytes, Token},
     prelude::*,
 };
 
@@ -48,23 +47,17 @@ fn validate_password(password: &str) -> Result<(), ValidationError> {
 pub async fn register(
     State(pool): State<Pool>,
     Form(params): Form<Params>,
-) -> Result<Json<Token>, Either<GeneralError, Error>> {
-    params
-        .validate()
-        .map_err(GeneralError::validation)
-        .map_err(Either::E1)?;
+) -> GeneralResult<Json<Token>> {
+    params.validate().handle_validation()?;
 
     let password = HashedBytes::hash_bytes(params.password.as_bytes()).as_bytes();
     let id = db::users::create_user(&params.username, &params.email, &password, &pool)
         .await
-        .handle_internal("Error creating the user")
-        .map_err(Either::E1)?
+        .handle_internal("Error creating the user")?
         .handle(
             StatusCode::BAD_REQUEST,
             "The username or the email are taken",
-        )
-        .map_err(Either::E1)?;
+        )?;
 
-    let token = Claims::new(id).encode().map_err(Either::E2)?;
-    Ok(Json(token))
+    Claims::new(id).encode().map(Json)
 }

src/errors.rs

@@ -4,14 +4,16 @@ use axum::{http::StatusCode, response::IntoResponse};
 
 type BoxError = Box<dyn std::error::Error>;
 
+/// Common error type for the project
 pub struct GeneralError {
+    /// Response status code
     pub status_code: StatusCode,
+    /// Message to send to the user
     pub message: Cow<'static, str>,
+    /// Error to log
     pub error: Option<BoxError>,
 }
 
-pub type GeneralResult<T> = Result<T, GeneralError>;
-
 impl GeneralError {
     pub fn message(status_code: StatusCode, message: impl Into<Cow<'static, str>>) -> Self {
         Self {
@@ -21,15 +23,10 @@ impl GeneralError {
         }
     }
 
-    #[allow(clippy::needless_pass_by_value)]
-    pub fn validation(error: validator::ValidationErrors) -> Self {
-        Self::message(StatusCode::BAD_REQUEST, error.to_string())
-    }
-
-    pub const fn item_not_found() -> Self {
-        GeneralError {
-            status_code: StatusCode::NOT_FOUND,
-            message: Cow::Borrowed("Item not found"),
+    pub const fn const_message(status_code: StatusCode, message: &'static str) -> Self {
+        Self {
+            status_code,
+            message: Cow::Borrowed(message),
             error: None,
         }
     }
@@ -44,6 +41,8 @@ impl IntoResponse for GeneralError {
     }
 }
 
+pub type GeneralResult<T> = Result<T, GeneralError>;
+
 pub trait ErrorHandlingExt<T, E>
 where
     Self: Sized,
@@ -86,3 +85,25 @@ impl<T> ErrorHandlingExt<T, Infallible> for Option<T> {
         })
     }
 }
+
+pub trait ItemNotFoundExt<T> {
+    fn item_not_found(self) -> Result<T, GeneralError>;
+}
+
+impl<T> ItemNotFoundExt<T> for Option<T> {
+    fn item_not_found(self) -> GeneralResult<T> {
+        const ITEM_NOT_FOUND_ERROR: GeneralError =
+            GeneralError::const_message(StatusCode::NOT_FOUND, "Item not found");
+        self.ok_or(ITEM_NOT_FOUND_ERROR)
+    }
+}
+
+pub trait ValidationExt<T> {
+    fn handle_validation(self) -> GeneralResult<T>;
+}
+
+impl<T> ValidationExt<T> for Result<T, validator::ValidationErrors> {
+    fn handle_validation(self) -> GeneralResult<T> {
+        self.map_err(|err| GeneralError::message(StatusCode::BAD_REQUEST, err.to_string()))
+    }
+}

src/prelude.rs

@@ -1,7 +1,10 @@
 pub(crate) use crate::{
     auth::Claims,
     db::{self, permissions::PermissionExt as _},
-    errors::{ErrorHandlingExt as _, GeneralError, GeneralResult},
+    errors::{
+        ErrorHandlingExt as _, GeneralError, GeneralResult, ItemNotFoundExt as _,
+        ValidationExt as _,
+    },
     AppState, Pool,
 };
 pub use axum::{