Refactor captcha solutions and AppState

src/captcha_solutions.rs

@@ -1,67 +1,56 @@
-use std::{ops::Deref, sync::Mutex};
-
+/// Stores captcha solutions in a wrapping way.
+/// If the capacity is exeeded, old values are overwritten.
+///
+/// The wrapping prevents attacks that keep filling the memory with solutions.
+/// The maximum number of solutions before wrapping is `u16::MAX as usize + 1` which is `65536`.
+///
+/// If you have more than 65536 persons trying to use the contact form at the same time, contact me!
+///
+/// Anyway, if the solution is overwritten, you get a new id in the contact form and can just submit
+/// a new captcha answer without loosing the content of the other fields.
 pub struct CaptchaSolutions {
+    /// The id of the last value stored.
     last_id: u16,
+    /// The vector of solutions.
     solutions: Vec<Option<String>>,
 }
 
-impl CaptchaSolutions {
-    fn get(&self, id: u16) -> &Option<String> {
-        &self.solutions[id as usize]
-    }
-
-    fn set(&mut self, id: u16, sol: Option<String>) {
-        self.solutions[id as usize] = sol;
-    }
-
-    #[must_use = "The new id has to be stored for future verification of an answer!"]
-    fn push(&mut self, sol: String) -> u16 {
-        self.last_id = self.last_id.wrapping_add(1);
-        self.set(self.last_id, Some(sol));
-
-        self.last_id
-    }
-}
-
-pub struct SharedCaptchaSolutions {
-    inner: Mutex<CaptchaSolutions>,
-}
-
-impl Default for SharedCaptchaSolutions {
+impl Default for CaptchaSolutions {
+    /// Default with `None` values.
     fn default() -> Self {
+        // Wrap the index around u16.
         let max_size = u16::MAX as usize + 1;
 
         Self {
-            inner: Mutex::new(CaptchaSolutions {
-                last_id: 0,
-                solutions: vec![None; max_size],
-            }),
+            last_id: 0,
+            solutions: vec![None; max_size],
         }
     }
 }
 
-impl Deref for SharedCaptchaSolutions {
-    type Target = Mutex<CaptchaSolutions>;
-
-    fn deref(&self) -> &Self::Target {
-        &self.inner
-    }
-}
-
-impl SharedCaptchaSolutions {
+impl CaptchaSolutions {
+    /// Pushes a solution and returns its id.
     #[must_use = "The new id has to be stored for future verification of an answer!"]
-    pub fn store_solution(&self, sol: String) -> u16 {
-        let mut sols = self.lock().unwrap();
+    pub fn push(&mut self, sol: String) -> u16 {
+        self.last_id = self.last_id.wrapping_add(1);
+        // Safety: Inbounds because of u16 and wrapping_add.
+        unsafe {
+            *self.solutions.get_unchecked_mut(self.last_id as usize) = Some(sol);
+        }
 
-        sols.push(sol)
+        self.last_id
     }
 
-    pub fn check_answer(&self, id: u16, answer: &str) -> bool {
-        let mut sols = self.lock().unwrap();
+    /// Checks an answer and returns if it matches the solution with the given id.
+    ///
+    /// The solution is removed in the case of a match to prevent using a true answer more than once.
+    pub fn check_answer(&mut self, id: u16, answer: &str) -> bool {
+        // Safety: Inbounds because of u16.
+        let sol_ref = unsafe { self.solutions.get_unchecked_mut(id as usize) };
 
-        if let Some(sol) = sols.get(id) {
-            if sol == answer.trim() {
-                sols.set(id, None);
+        if let Some(sol) = sol_ref {
+            if sol == answer {
+                *sol_ref = None;
                 return true;
             }
         }

src/mailer.rs

@@ -5,7 +5,7 @@ use lettre::{
     Message, SmtpTransport, Transport,
 };
 
-use crate::config;
+use crate::config::Config;
 
 pub struct Mailer {
     mailer: SmtpTransport,
@@ -13,7 +13,7 @@ pub struct Mailer {
 }
 
 impl Mailer {
-    pub fn new(config: &mut config::Config) -> Result<Self> {
+    pub fn new(config: &Config) -> Result<Self> {
         let creds = Credentials::new(
             config.email_server.email.clone(),
             config.email_server.password.clone(),

src/main.rs

@@ -14,35 +14,29 @@ use axum::{
     routing::{get, get_service, Router},
     Server,
 };
-use std::{future::ready, net::SocketAddr, process, sync::Arc};
+use std::{future::ready, net::SocketAddr, process};
 use tower_http::services::ServeDir;
 use tracing::info;
 
-async fn init() -> Result<()> {
-    let mut config = config::Config::new()?;
-    let path_prefix = config.path_prefix.clone();
-    let mailer = Arc::new(mailer::Mailer::new(&mut config)?);
+use crate::states::AppState;
 
-    let socket_address = config
+async fn init() -> Result<()> {
+    let app_state = AppState::build()?;
+
+    let path_prefix = app_state.config.path_prefix.clone();
+
+    let socket_address = app_state
+        .config
         .socket_address
         .parse::<SocketAddr>()
         .context("Failed to parse the socket address: {e:?}")?;
 
     logging::init_logger(
-        &config.log_file,
-        config.utc_offset.hours,
-        config.utc_offset.minutes,
+        &app_state.config.log_file,
+        app_state.config.utc_offset.hours,
+        app_state.config.utc_offset.minutes,
     )?;
 
-    let config = Arc::new(config);
-    let captcha_solutions = Arc::new(captcha_solutions::SharedCaptchaSolutions::default());
-
-    let app_state = states::AppState {
-        config,
-        mailer,
-        captcha_solutions,
-    };
-
     let static_service =
         get_service(ServeDir::new("static")).handle_error(|_| ready(StatusCode::NOT_FOUND));
 
@@ -61,7 +55,7 @@ async fn init() -> Result<()> {
     Ok(())
 }
 
-/// Single thread
+/// Single thread.
 #[tokio::main(flavor = "current_thread")]
 async fn main() {
     if let Err(e) = init().await {

src/routes.rs

@@ -1,15 +1,20 @@
 use anyhow::{Context, Result};
 use askama_axum::IntoResponse;
-use axum::extract::{Form, State};
-use axum::response::Response;
-use std::sync::Arc;
+use axum::{
+    extract::{Form, State},
+    response::Response,
+};
+use std::sync::{Arc, Mutex};
 use tracing::{error, info};
 
-use crate::{captcha_solutions, config, errors, forms, mailer, templates};
+use crate::{
+    captcha_solutions::CaptchaSolutions, config::Config, errors::AppError, forms::ContactForm,
+    mailer::Mailer, templates,
+};
 
 pub struct IndexParams<'a> {
-    config: Arc<config::Config>,
-    captcha_solutions: Arc<captcha_solutions::SharedCaptchaSolutions>,
+    config: Arc<Config>,
+    captcha_solutions: Arc<Mutex<CaptchaSolutions>>,
     was_validated: bool,
     name: Option<String>,
     email: Option<String>,
@@ -19,9 +24,9 @@ pub struct IndexParams<'a> {
 }
 
 pub async fn index(
-    State(config): State<Arc<config::Config>>,
-    State(captcha_solutions): State<Arc<captcha_solutions::SharedCaptchaSolutions>>,
-) -> Result<Response, errors::AppError> {
+    State(config): State<Arc<Config>>,
+    State(captcha_solutions): State<Arc<Mutex<CaptchaSolutions>>>,
+) -> Result<Response, AppError> {
     info!("Visited get(index)");
 
     render_contact_form(IndexParams {
@@ -37,13 +42,13 @@ pub async fn index(
     .await
 }
 
-pub async fn render_contact_form(params: IndexParams<'_>) -> Result<Response, errors::AppError> {
+pub async fn render_contact_form(params: IndexParams<'_>) -> Result<Response, AppError> {
     let captcha = captcha::by_name(captcha::Difficulty::Easy, captcha::CaptchaName::Lucy);
     let captcha_base64 = captcha.as_base64().context("Failed to create a captcha!")?;
 
     let solution = captcha.chars_as_string();
 
-    let id = params.captcha_solutions.store_solution(solution);
+    let id = params.captcha_solutions.lock().unwrap().push(solution);
 
     let template = templates::ContactForm {
         base: templates::Base {
@@ -65,11 +70,11 @@ pub async fn render_contact_form(params: IndexParams<'_>) -> Result<Response, er
 }
 
 async fn failed_submission(
-    config: Arc<config::Config>,
-    captcha_solutions: Arc<captcha_solutions::SharedCaptchaSolutions>,
+    config: Arc<Config>,
+    captcha_solutions: Arc<Mutex<CaptchaSolutions>>,
     error_message: &str,
-    form: forms::ContactForm,
-) -> Result<Response, errors::AppError> {
+    form: ContactForm,
+) -> Result<Response, AppError> {
     let params = IndexParams {
         config,
         captcha_solutions,
@@ -85,12 +90,17 @@ async fn failed_submission(
 }
 
 pub async fn submit(
-    State(config): State<Arc<config::Config>>,
-    State(captcha_solutions): State<Arc<captcha_solutions::SharedCaptchaSolutions>>,
-    State(mailer): State<Arc<mailer::Mailer>>,
-    Form(form): Form<forms::ContactForm>,
-) -> Result<Response, errors::AppError> {
-    if !captcha_solutions.check_answer(form.id, &form.captcha_answer) {
+    State(config): State<Arc<Config>>,
+    State(captcha_solutions): State<Arc<Mutex<CaptchaSolutions>>>,
+    State(mailer): State<Arc<Mailer>>,
+    Form(form): Form<ContactForm>,
+) -> Result<Response, AppError> {
+    let right_captcha_answer = captcha_solutions
+        .lock()
+        .unwrap()
+        .check_answer(form.id, form.captcha_answer.trim());
+
+    if !right_captcha_answer {
         info!("Wrong CAPTCHA");
 
         return failed_submission(
@@ -120,7 +130,7 @@ pub async fn submit(
     success(config).await
 }
 
-pub async fn success(config: Arc<config::Config>) -> Result<Response, errors::AppError> {
+pub async fn success(config: Arc<Config>) -> Result<Response, AppError> {
     info!("Successful contact form submission");
 
     let template = templates::Success {

src/states.rs

@@ -1,11 +1,28 @@
+use anyhow::Result;
 use axum::extract::FromRef;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
-use crate::{captcha_solutions, config, mailer};
+use crate::{captcha_solutions::CaptchaSolutions, config::Config, mailer::Mailer};
 
 #[derive(Clone, FromRef)]
 pub struct AppState {
-    pub config: Arc<config::Config>,
-    pub mailer: Arc<mailer::Mailer>,
-    pub captcha_solutions: Arc<captcha_solutions::SharedCaptchaSolutions>,
+    pub config: Arc<Config>,
+    pub mailer: Arc<Mailer>,
+    pub captcha_solutions: Arc<Mutex<CaptchaSolutions>>,
+}
+
+impl AppState {
+    pub fn build() -> Result<Self> {
+        let config = Config::new()?;
+        let mailer = Arc::new(Mailer::new(&config)?);
+
+        let config = Arc::new(config);
+        let captcha_solutions = Arc::new(Mutex::new(CaptchaSolutions::default()));
+
+        Ok(Self {
+            config,
+            mailer,
+            captcha_solutions,
+        })
+    }
 }