Add better error handling

dragynfruit · Dec 2, 2024 · c75e48f · c75e48f
1 parent dfa74ce
commit c75e48f
Show file tree

Hide file tree

Showing 9 changed files with 286 additions and 124 deletions.
diff --git a/src/client.rs b/src/client.rs
@@ -1,5 +1,6 @@
 use scraper::Html;
 use ureq::{Agent, AgentBuilder};
+use std::fmt;
 
 #[derive(Clone)]
 pub struct Client {
@@ -24,6 +25,15 @@ impl From<std::io::Error> for ClientError {
     }
 }
 
+impl fmt::Display for ClientError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            ClientError::UreqError(err) => write!(f, "HTTP request error: {}", err),
+            ClientError::IoError(err) => write!(f, "IO error: {}", err),
+        }
+    }
+}
+
 impl Client {
     pub fn new() -> Self {
         Self {

diff --git a/src/routes/error.rs b/src/routes/error.rs
@@ -9,10 +9,77 @@ use tera::Context;
 use crate::client::ClientError;
 use crate::templates::TEMPLATES;
 
-#[derive(Serialize)]
+#[derive(Serialize, Debug)]
+pub enum ErrorSource {
+    Upstream,
+    Internal,
+}
+
+#[derive(Serialize, Debug)]
 pub struct Error {
     status: u16,
-    message: Option<String>,
+    message: String,
+    details: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    stack_trace: Option<String>,
+    source: ErrorSource,
+}
+
+impl Error {
+    pub fn new(status: u16, message: String, source: ErrorSource) -> Self {
+        Self {
+            status,
+            message,
+            details: None,
+            stack_trace: None,
+            source,
+        }
+    }
+}
+
+impl From<ClientError> for Error {
+    fn from(err: ClientError) -> Self {
+        match err {
+            ClientError::UreqError(error) => {
+                let (status, message) = match error {
+                    ureq::Error::Status(code, response) => {
+                        (code, response.status_text().to_string())
+                    }
+                    ureq::Error::Transport(transport) => {
+                        (StatusCode::BAD_GATEWAY.as_u16(), transport.to_string())
+                    }
+                };
+                Error::new(status, message, ErrorSource::Upstream)
+            }
+            ClientError::IoError(error) => Error {
+                status: StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
+                message: "Internal Server Error".to_string(),
+                details: Some(error.to_string()),
+                stack_trace: if cfg!(debug_assertions) {
+                    Some(format!("{:?}", error))
+                } else {
+                    None
+                },
+                source: ErrorSource::Internal,
+            }
+        }
+    }
+}
+
+impl From<tera::Error> for Error {
+    fn from(err: tera::Error) -> Self {
+        Error {
+            status: StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
+            message: "Template rendering error".to_string(),
+            details: Some(err.to_string()),
+            stack_trace: if cfg!(debug_assertions) {
+                Some(format!("{:?}", err))
+            } else {
+                None
+            },
+            source: ErrorSource::Internal,
+        }
+    }
 }
 
 pub fn render_error(error: Error) -> Response<Body> {
@@ -28,28 +95,15 @@ pub fn render_error(error: Error) -> Response<Body> {
 }
 
 pub async fn error_404() -> impl IntoResponse {
-    return render_error(Error {
-        status: StatusCode::NOT_FOUND.as_u16(),
-        message: None,
-    });
+    render_error(Error::new(
+        StatusCode::NOT_FOUND.as_u16(),
+        "Page not found".to_string(),
+        ErrorSource::Internal,
+    ))
 }
 
 pub fn construct_error(error: ClientError) -> Response<Body> {
-    match error {
-        ClientError::UreqError(error) => {
-            let response = error.into_response().unwrap();
-            render_error(Error {
-                status: response.status(),
-                message: None,
-            })
-            .into_response()
-        }
-        ClientError::IoError(error) => render_error(Error {
-            status: 500,
-            message: Some(error.to_string()),
-        })
-        .into_response(),
-    }
+    render_error(Error::from(error))
 }
 
 #[cfg(test)]

diff --git a/src/routes/imgs.rs b/src/routes/imgs.rs
@@ -1,41 +1,54 @@
 use axum::{
-    body::Body,
-    extract::{Path, State},
-    response::{IntoResponse, Response},
-    routing, Router,
+    body::Body, extract::{Path, State}, http::StatusCode, response::{IntoResponse, Response}, routing, Router
 };
 
 use crate::{state::AppState, constants::URL};
 
+use super::error::{render_error, Error, ErrorSource};
+
 pub fn get_router(state: AppState) -> Router {
     Router::new()
         .route("/guest.png", routing::get(guest))
         .route("/:id0/:id1/:id2/:id3.jpg", routing::get(icon))
         .with_state(state)
 }
 
-async fn guest() -> impl IntoResponse {
+async fn guest() -> Result<Response<Body>, Response<Body>> {
     Response::builder()
         .status(200)
         .header("Content-Type", "image/png")
         .header("Cache-Control", "public, max-age=31536000, immutable")
         .body(Body::from(include_bytes!("assets/guest.png").to_vec()))
-        .unwrap()
+        .map_err(|e| render_error(Error::new(
+            StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
+            "Failed to serve guest image".to_string(),
+            ErrorSource::Internal
+        )))
 }
 
 async fn icon(
     State(state): State<AppState>,
     Path((id0, id1, id2, id3)): Path<(String, String, String, String)>,
-) -> impl IntoResponse {
-    let id3 = id3.split_once(".").unwrap().0;
+) -> Result<Response<Body>, Response<Body>> {
+    let id3 = id3.split_once(".")
+        .ok_or_else(|| render_error(Error::new(
+            StatusCode::BAD_REQUEST.as_u16(),
+            "Invalid image path".to_string(),
+            ErrorSource::Internal
+        )))?
+        .0;
+
     let path = format!("{id0}/{id1}/{id2}/{id3}");
     let tree = state.db.open_tree("icons").unwrap();
 
     let icon = if tree.contains_key(&path).unwrap() {
         tree.get(&path).unwrap().unwrap().to_vec()
     } else {
-        let icon = state.client.get_bytes(format!("{URL}/cache/img/{path}.jpg").as_str()).unwrap(); //fix
-        tree.insert(&path, icon.clone()).unwrap();
+        let icon = state.client.get_bytes(format!("{URL}/cache/img/{path}.jpg").as_str()).unwrap();
+        let save_icon = icon.clone();
+        tokio::spawn(async move {
+            tree.insert(&path, save_icon).ok();
+        });
         icon
     };
 
@@ -44,7 +57,11 @@ async fn icon(
         .header("Content-Type", "image/jpeg")
         .header("Cache-Control", "public, max-age=31536000, immutable")
         .body(Body::from(icon))
-        .unwrap()
+        .map_err(|e| render_error(Error::new(
+            StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
+            "Failed to serve icon image".to_string(),
+            ErrorSource::Internal
+        )))
 }
 
 #[cfg(test)]
@@ -53,7 +70,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_guest() {
-        let response = guest().await;
+        let response = guest().await.unwrap();
         assert_eq!(response.into_response().status(), 200);
     }
 
@@ -69,7 +86,8 @@ mod tests {
                 "10674139.jpg".to_string(),
             )),
         )
-        .await;
+        .await
+        .unwrap();
         assert_eq!(response.into_response().status(), 200);
     }
 }
diff --git a/src/routes/info.rs b/src/routes/info.rs
@@ -1,6 +1,7 @@
 use axum::{
     body::Body,
     extract::State,
+    http::StatusCode,
     response::{IntoResponse, Response},
     routing, Json, Router,
 };
@@ -9,6 +10,7 @@ use std::sync::OnceLock;
 use tera::Context;
 
 use crate::{state::AppState, templates::TEMPLATES};
+use super::error::{Error, ErrorSource, render_error};
 
 pub static DEPLOY_DATE: OnceLock<String> = OnceLock::new();
 
@@ -29,6 +31,11 @@ fn get_info(state: AppState) -> InstanceInfo {
     let commit = include_str!("../../.git/FETCH_HEAD")
         .lines()
         .next()
+        .ok_or_else(|| Error::new(
+            StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
+            "Failed to read commit info".to_string(),
+            ErrorSource::Internal
+        ))
         .unwrap()
         .split('\t')
         .next()
@@ -57,21 +64,21 @@ pub fn get_router(state: AppState) -> Router {
         .with_state(state)
 }
 
-async fn info(State(state): State<AppState>) -> impl IntoResponse {
-    Response::builder()
-        .status(200)
-        .header("Content-Type", "text/html")
-        .body(Body::new(
-            TEMPLATES
-                .render(
-                    "info.html",
-                    &Context::from_serialize(&get_info(state)).unwrap(),
-                )
-                .unwrap(),
-        ))
-        .unwrap()
+async fn info(State(state): State<AppState>) -> Result<Response<Body>, Response<Body>> {
+    let info = get_info(state);
+    TEMPLATES
+        .render("info.html", &Context::from_serialize(&info).unwrap())
+        .map(|html| {
+            Response::builder()
+                .status(200)
+                .header("Content-Type", "text/html")
+                .body(Body::new(html))
+                .unwrap()
+        })
+        .map_err(|e| render_error(Error::from(e)))
 }
 
-async fn info_json(State(state): State<AppState>) -> Json<InstanceInfo> {
-    Json(get_info(state))
+async fn info_json(State(state): State<AppState>) -> Result<Json<InstanceInfo>, Response<Body>> {
+    let info = get_info(state);
+    Ok(Json(info))
 }