Fix screenshot functionality and add API key based ratelimiting

backend/Cargo.toml

@@ -6,7 +6,10 @@ edition = "2024"
 [dependencies]
 async-openai = "0.29.2"
 axum = "0.8.4"
+dashmap = "6.1.0"
 dotenv = "0.15.0"
+governor = "0.10.1"
+once_cell = "1.21.3"
 serde = "1.0.219"
 sqlx = { version = "0.7", features = ["runtime-tokio-rustls", "postgres", "chrono"] }
 tokio = { version = "1", features = ["full"] }

backend/src/lib.rs

@@ -1,2 +1,3 @@
 pub mod openai;
 pub mod database;
+pub mod ratelimit;

backend/src/main.rs

@@ -3,6 +3,8 @@ use serde::{Deserialize, Serialize};
 use std::net::SocketAddr;
 use backend::openai::{get_step_by_step_guidance, get_final_answer};
 use backend::database::{is_valid_api_key, is_root_api_key, add_api_key as db_add_api_key, init_db, test_connection};
+use backend::ratelimit::check_rate_limit;
+
 // An enum to determine the type of request
 #[derive(Debug, Deserialize)]
 pub enum OpenAiRequestType {
@@ -71,6 +73,12 @@ async fn openai(Json(payload): Json<OpenAiRequest>) -> (StatusCode, Json<OpenAiR
                     response: "Invalid API key".to_string(),
                 }));
             }
+            if !check_rate_limit(&payload.api_key) {
+                return (StatusCode::TOO_MANY_REQUESTS, Json(OpenAiResponse {
+                    success: false,
+                    response: "Rate limit exceeded".to_string(),
+                }));
+            }
         }
         Err(e) => {
             return (StatusCode::INTERNAL_SERVER_ERROR, Json(OpenAiResponse {

backend/src/openai.rs

@@ -25,9 +25,16 @@ pub async fn get_step_by_step_guidance(image_b64: &str) -> Result<String, Box<dy
         name: None,
     };
 
+    // Handle both raw base64 and data URLs
+    let image_url = if image_b64.starts_with("data:") {
+        image_b64.to_string()
+    } else {
+        format!("data:image/jpeg;base64,{}", image_b64)
+    };
+
     let image_content = ChatCompletionRequestMessageContentPartImage {
         image_url: ImageUrl {
-            url: format!("data:image/jpeg;base64,{}", image_b64),
+            url: image_url,
             detail: None,
         },
     };
@@ -68,9 +75,16 @@ pub async fn get_final_answer(image_b64: &str) -> Result<String, Box<dyn std::er
         name: None,
     };
 
+    // Handle both raw base64 and data URLs
+    let image_url = if image_b64.starts_with("data:") {
+        image_b64.to_string()
+    } else {
+        format!("data:image/jpeg;base64,{}", image_b64)
+    };
+
     let image_content = ChatCompletionRequestMessageContentPartImage {
         image_url: ImageUrl {
-            url: format!("data:image/jpeg;base64,{}", image_b64),
+            url: image_url,
             detail: None,
         },
     };

backend/src/ratelimit.rs

@@ -0,0 +1,27 @@
+use governor::{Quota, clock::DefaultClock, state::InMemoryState, RateLimiter};
+use std::{num::NonZeroU32, sync::Arc};
+use dashmap::DashMap;
+use once_cell::sync::Lazy;
+use governor::state::NotKeyed;
+
+type ApiKey = String;
+
+const RATE_LIMIT: u32 = 3;
+
+// Global, thread-safe map: API key → rate limiter
+static API_KEY_LIMITERS: Lazy<DashMap<ApiKey, Arc<RateLimiter<NotKeyed, InMemoryState, DefaultClock>>>> =
+    Lazy::new(DashMap::new);
+
+/// Checks if the given API key is within its rate limit.
+/// Returns `true` if allowed, `false` if rate limit is exceeded.
+pub fn check_rate_limit(api_key: &str) -> bool {
+    let limiter = API_KEY_LIMITERS
+        .entry(api_key.to_string())
+        .or_insert_with(|| {
+            let quota = Quota::per_minute(NonZeroU32::new(RATE_LIMIT).unwrap());
+            Arc::new(RateLimiter::direct_with_clock(quota, DefaultClock::default()))
+        })
+        .clone();
+
+    limiter.check().is_ok()
+}

frontend/maths-online-app/src-tauri/tauri.conf.json

@@ -1,6 +1,6 @@
 {
   "$schema": "../node_modules/@tauri-apps/cli/config.schema.json",
-  "productName": "maths-online-app",
+  "productName": "Maths Online App",
   "version": "0.1.0",
   "identifier": "com.connorakey.dev",
   "build": {