Rate limiting, caching, patching vulnerabilities

Author: breynard0

Cargo.lock

@@ -8,4 +8,5 @@ actix-web = "4.14.1"
 anyhow = "1.0.104"
 reqwest = { version = "0.13.4", features = ["json"] }
 serde = { version = "1.0.229", features = ["derive"] }
+serde_json = "1.0.151"
 tokio = { version = "1.53.1", features = ["macros", "rt-multi-thread"] }

Cargo.toml

@@ -1,8 +1,15 @@
+use std::{
+    collections::HashMap,
+    sync::OnceLock,
+    time::{Duration, Instant},
+};
+
 use anyhow::{Ok, Result};
-use reqwest::Response;
+use reqwest::Client;
 use serde::{Deserialize, Serialize};
+use tokio::sync::Mutex;
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub struct GithubPathEntry {
     path: String,
     mode: String,
@@ -11,7 +18,7 @@ pub struct GithubPathEntry {
     sha: String,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub struct GithubRepoResponse {
     truncated: bool,
     sha: String,
@@ -20,24 +27,56 @@ pub struct GithubRepoResponse {
 }
 
 fn github_token() -> Result<String> {
-    std::env::var("GITHUB_TOKEN").map_err(|_| anyhow::anyhow!("GITHUB_TOKEN not set"))
+    std::env::var("GITHUB_TOKEN").map_err(|_| anyhow::anyhow!("github credentials unavailable"))
 }
 
-pub async fn send_github_request(url: &str) -> Result<Response> {
-    let client = reqwest::Client::builder()
+static CLIENT: OnceLock<Client> = OnceLock::new();
+static CACHE: OnceLock<Mutex<HashMap<String, (String, Instant)>>> = OnceLock::new();
+
+const REQUEST_TIMEOUT_SECONDS: u64 = 30;
+const CACHE_SECONDS: u64 = 300;
+
+fn github_client() -> Result<&'static Client> {
+    if let Some(client) = CLIENT.get() {
+        return Ok(client);
+    }
+
+    let client = Client::builder()
         .user_agent("HackClubPyterm (brendan@hackclub.com)")
+        .timeout(Duration::from_secs(REQUEST_TIMEOUT_SECONDS))
         .build()?;
-    client
+
+    Ok(CLIENT.get_or_init(|| client))
+}
+
+pub async fn send_github_request(url: &str) -> Result<String> {
+    let cache = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
+
+    if let Some((data, stored_at)) = cache.lock().await.get(url) {
+        if Instant::now().duration_since(*stored_at).as_secs() < CACHE_SECONDS {
+            return Ok(data.clone());
+        }
+    }
+
+    let output = github_client()?
         .get(url)
         .bearer_auth(github_token()?)
         .send()
-        .await
-        .map_err(|e| anyhow::anyhow!(e))
+        .await?
+        .error_for_status()?
+        .text()
+        .await?;
+
+    let mut handle = cache.lock().await;
+    handle.retain(|_, x| Instant::now().duration_since(x.1).as_secs() < CACHE_SECONDS);
+    handle.insert(url.to_string(), (output.clone(), Instant::now()));
+
+    Ok(output)
 }
 
 pub async fn get_repo_files(user: &String, repo: &String) -> Result<Vec<String>> {
     let url = format!("https://api.github.com/repos/{user}/{repo}/git/trees/HEAD?recursive=1");
-    let response: GithubRepoResponse = send_github_request(&url).await?.json().await?;
+    let response: GithubRepoResponse = serde_json::from_str(&send_github_request(&url).await?)?;
 
     Ok(response.tree.iter().map(|x| x.path.clone()).collect())
 }
@@ -47,9 +86,7 @@ pub async fn get_file_in_repo(user: &String, repo: &String, file_path: &String)
         "https://raw.githubusercontent.com/{user}/{repo}/HEAD/{}",
         file_path.trim_start_matches("/")
     );
-    let response = send_github_request(&url).await?;
-
-    Ok(response.text().await?)
+    Ok(send_github_request(&url).await?)
 }
 
 pub async fn get_python_code(user: &String, repo: &String) -> Result<String> {
@@ -66,11 +103,18 @@ pub async fn get_python_code(user: &String, repo: &String) -> Result<String> {
     for query in files_as_priority {
         let query_string = query.to_string();
         if files.contains(&query_string) {
-            path = Some(query_string)
+            path = Some(query_string);
+            break;
         }
     }
+
     if path.is_none() {
-        path = files.iter().find(|x| x.contains(".py")).cloned();
+        path = files.iter().find(|x| x.ends_with(".py")).cloned();
     }
+
+    if path.is_none() {
+        return Ok("No Python files in directory".to_string());
+    }
+
     get_file_in_repo(user, repo, &path.unwrap()).await
 }

src/github.rs

@@ -1,5 +1,10 @@
 pub mod github;
 
+use std::{
+    sync::{Mutex, MutexGuard, OnceLock},
+    time::Instant,
+};
+
 use actix_web::{
     App, HttpRequest, HttpResponse, HttpServer, Responder, http::header, middleware::DefaultHeaders,
 };
@@ -8,7 +13,135 @@ use crate::github::get_python_code;
 
 const PAGE_HTML: &str = include_str!("ui/page.html");
 
+static CONNECTIONS: OnceLock<Mutex<Vec<(String, Instant)>>> = OnceLock::new();
+static RATE_LIMITED: OnceLock<Mutex<Vec<(String, Instant)>>> = OnceLock::new();
+static RATE_LIMITED_THIS_INSTANCE: OnceLock<Mutex<Vec<String>>> = OnceLock::new();
+static PEER_CONNECTIONS: OnceLock<Mutex<Vec<(String, Instant)>>> = OnceLock::new();
+
+const MAX_PER_TEN_SECONDS: u32 = 70; // ten full page loads
+const MAX_PER_TEN_SECONDS_PEER: u32 = 700; // one hundred full page loads
+const RATE_LIMIT_MINUTES_FIRST: u64 = 1;
+const RATE_LIMIT_MINUTES_SECOND: u64 = 30;
+
+fn escape_html(input: &str) -> String {
+    input
+        .replace("&", "&amp;")
+        .replace("<", "&lt;")
+        .replace(">", "&gt;")
+        .replace("\"", "&quot;")
+        .replace("'", "&#39;")
+}
+
+fn extract_global<'a, T>(input: &'a OnceLock<Mutex<Vec<T>>>) -> MutexGuard<'a, Vec<T>> {
+    input
+        .get_or_init(|| Mutex::new(vec![]))
+        .lock()
+        .unwrap_or_else(|x| x.into_inner())
+}
+
 async fn dispatch(req: HttpRequest) -> impl Responder {
+    {
+        let ip = req
+            .connection_info()
+            .realip_remote_addr()
+            .map(|x| x.to_string())
+            .unwrap_or(String::new());
+
+        let peer_ip = req
+            .peer_addr()
+            .map(|x| x.ip().to_string())
+            .unwrap_or(String::new());
+        if ip != peer_ip {
+            eprintln!("forwarded ip {ip} claimed by peer {peer_ip}");
+        }
+
+        let mut peer_connections = extract_global(&PEER_CONNECTIONS);
+        *peer_connections = peer_connections
+            .iter()
+            .filter(|x| std::time::Instant::now().duration_since(x.1).as_secs() < 10)
+            .map(|x| (x.0.clone(), x.1))
+            .collect::<Vec<_>>();
+
+        let peer_count = peer_connections.iter().filter(|x| x.0 == peer_ip).count() as u32;
+        if peer_count >= MAX_PER_TEN_SECONDS_PEER {
+            return HttpResponse::TooManyRequests()
+                .content_type("text/plain; charset=utf-8")
+                .body("rate limited, please wait");
+        }
+
+        peer_connections.push((peer_ip, std::time::Instant::now()));
+
+        let mut rate_limited = extract_global(&RATE_LIMITED);
+
+        let mut connections = extract_global(&CONNECTIONS);
+
+        let mut rate_limited_this_instance = extract_global(&RATE_LIMITED_THIS_INSTANCE);
+        rate_limited_this_instance.sort();
+        rate_limited_this_instance.dedup();
+
+        fn check_rate_limit_not_ready_to_clear(
+            x: &(String, Instant),
+            rate_limited_this_instance: &[String],
+        ) -> bool {
+            let rate_limit_time = if !rate_limited_this_instance.contains(&x.0) {
+                RATE_LIMIT_MINUTES_FIRST * 60
+            } else {
+                RATE_LIMIT_MINUTES_SECOND * 60
+            };
+
+            std::time::Instant::now().duration_since(x.1).as_secs() < rate_limit_time
+        }
+
+        let rate_limited_this_instance_immutable = rate_limited_this_instance.clone();
+
+        let expired = rate_limited
+            .iter()
+            .filter(|x| {
+                !check_rate_limit_not_ready_to_clear(x, &rate_limited_this_instance_immutable)
+            })
+            .map(|x| x.0.clone())
+            .collect::<Vec<_>>();
+
+        *rate_limited = rate_limited
+            .iter()
+            .filter(|x| {
+                check_rate_limit_not_ready_to_clear(x, &rate_limited_this_instance_immutable)
+            })
+            .map(|x| (x.0.clone(), x.1))
+            .collect::<Vec<_>>();
+
+        rate_limited_this_instance.extend(expired);
+
+        if rate_limited.iter().find(|x| x.0 == ip).is_some() {
+            return HttpResponse::TooManyRequests()
+                .content_type("text/plain; charset=utf-8")
+                .body("rate limited, please wait");
+        }
+
+        *connections = connections
+            .iter()
+            .filter(|x| std::time::Instant::now().duration_since(x.1).as_secs() < 10)
+            .map(|x| (x.0.clone(), x.1))
+            .collect::<Vec<_>>();
+        connections.sort();
+        let mut count = 1;
+        let mut last = String::new();
+        for (ip, _) in connections.iter() {
+            if *ip == last {
+                count += 1;
+            } else {
+                count = 1;
+            }
+            if count >= MAX_PER_TEN_SECONDS {
+                rate_limited.push((ip.clone(), std::time::Instant::now()));
+            }
+
+            last = ip.clone();
+        }
+
+        connections.push((ip, std::time::Instant::now()));
+    }
+
     let path = req.path();
     let terms = path
         .split("/")
@@ -17,36 +150,56 @@ async fn dispatch(req: HttpRequest) -> impl Responder {
 
     if terms.len() < 2 {
         eprintln!("Invalid path: {}", path);
-        return HttpResponse::Ok().body("Invalid path");
+        return HttpResponse::Ok()
+            .content_type("text/plain; charset=utf-8")
+            .body("Invalid path");
     }
 
     let user = terms[0].to_string();
     let repo = terms[1].to_string();
     if terms.len() != 0 {
         match *terms.last().unwrap() {
-            "term_style.css" => return HttpResponse::Ok().body(include_str!("ui/term_style.css")),
-            "term_config.js" => return HttpResponse::Ok().body(include_str!("ui/term_config.js")),
-            "conf.json" => return HttpResponse::Ok().body("{}"),
+            "term_style.css" => {
+                return HttpResponse::Ok()
+                    .content_type("text/css; charset=utf-8")
+                    .body(include_str!("ui/term_style.css"));
+            }
+            "term_config.js" => {
+                return HttpResponse::Ok()
+                    .content_type("text/javascript; charset=utf-8")
+                    .body(include_str!("ui/term_config.js"));
+            }
+            "conf.json" => {
+                return HttpResponse::Ok()
+                    .content_type("application/json")
+                    .body("{}");
+            }
             "script.py" => {
                 let python_code = match get_python_code(&user, &repo).await {
                     Ok(code) => code,
                     Err(e) => {
                         eprintln!("Error: {}", e);
-                        return HttpResponse::Ok().body(format!("Error: {}", e));
+                        return HttpResponse::InternalServerError()
+                            .content_type("text/plain; charset=utf-8")
+                            .body("Internal server error");
                     }
                 };
-                return HttpResponse::Ok().body(python_code);
+                return HttpResponse::Ok()
+                    .content_type("text/plain; charset=utf-8")
+                    .body(python_code);
             }
             _ => {}
         }
     }
 
-    HttpResponse::Ok().body(
-        PAGE_HTML
-            .replace("{PAGE_TITLE}", &repo.to_string())
-            .replace("{USER}", &user)
-            .replace("{REPO}", &repo),
-    )
+    HttpResponse::Ok()
+        .content_type("text/html; charset=utf-8")
+        .body(
+            PAGE_HTML
+                .replace("{PAGE_TITLE}", &escape_html(&repo))
+                .replace("{USER}", &escape_html(&user))
+                .replace("{REPO}", &escape_html(&repo)),
+        )
 }
 
 #[actix_web::main]
@@ -57,7 +210,8 @@ async fn main() -> std::io::Result<()> {
                 DefaultHeaders::new()
                     .add((header::CROSS_ORIGIN_OPENER_POLICY, "same-origin"))
                     .add((header::CROSS_ORIGIN_EMBEDDER_POLICY, "require-corp"))
-                    .add((header::CROSS_ORIGIN_RESOURCE_POLICY, "cross-origin")),
+                    .add((header::CROSS_ORIGIN_RESOURCE_POLICY, "cross-origin"))
+                    .add((header::X_CONTENT_TYPE_OPTIONS, "nosniff")),
             )
             .default_service(actix_web::web::to(dispatch))
     })