Add idempotency lock for axum server

Author: AnthonyRonning

src/lib.rs

@@ -4,12 +4,26 @@ use worker::{event, Context, Date, Env, Request};
 
 // All imports for server specific things
 #[cfg(feature = "server")]
+use axum::{
+    async_trait,
+    extract::FromRequest,
+    extract::{Path, WebSocketUpgrade},
+    http::{Request, StatusCode},
+    TypedHeader,
+};
+#[cfg(feature = "server")]
 use axum::{routing::get, Router};
 #[cfg(feature = "server")]
+use std::collections::HashMap;
+#[cfg(feature = "server")]
 use std::env;
 #[cfg(feature = "server")]
 use std::net::SocketAddr;
 #[cfg(feature = "server")]
+use std::sync::Arc;
+#[cfg(feature = "server")]
+use tokio::sync::Mutex;
+#[cfg(feature = "server")]
 use tower_http::trace::{DefaultMakeSpan, TraceLayer};
 
 mod logger;
@@ -21,15 +35,48 @@ mod cloudflare;
 #[cfg(feature = "server")]
 mod server;
 
+#[cfg(feature = "server")]
+struct IdempotencyKey(Option<String>);
+
+#[cfg(feature = "server")]
+#[async_trait]
+impl<'a, B> FromRequest<(), B> for IdempotencyKey
+where
+    B: Send + 'static,
+{
+    type Rejection = StatusCode;
+
+    async fn from_request(req: Request<B>, _: &()) -> Result<Self, Self::Rejection> {
+        let headers = req.headers().clone();
+        let key = headers
+            .get("Idempotency-Key")
+            .and_then(|v| v.to_str().ok())
+            .map(|s| s.to_string());
+        Ok(IdempotencyKey(key))
+    }
+}
+
 /// Main function for running the program as a server
 #[tokio::main]
 #[cfg(feature = "server")]
 async fn main() {
     println!("Running ln-websocket-proxy");
     tracing_subscriber::fmt::init();
 
+    let locks = Arc::new(Mutex::new(HashMap::new()));
+
     let app = Router::new()
-        .route("/v1/:ip/:port", get(crate::server::ws_handler))
+        .route(
+            "/v1/:ip/:port",
+            get(
+                |path: Path<(String, String)>,
+                 ws: WebSocketUpgrade,
+                 user_agent: Option<TypedHeader<headers::UserAgent>>,
+                 idempotency_key: IdempotencyKey| async move {
+                    server::ws_handler(path, ws, user_agent, idempotency_key.0, locks.clone())
+                },
+            ),
+        )
         .layer(
             TraceLayer::new_for_http()
                 .make_span_with(DefaultMakeSpan::default().include_headers(true)),

src/server.rs

@@ -6,27 +6,79 @@ use axum::{
     },
     response::IntoResponse,
 };
-use std::time::Duration;
+use std::{collections::HashMap, sync::Arc, time::Duration};
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     net::TcpStream,
+    sync::Mutex,
 };
 
-pub(crate) async fn ws_handler(
+struct IdempotencyLockGuard {
+    key: Option<String>,
+    locks: Arc<Mutex<HashMap<String, bool>>>,
+}
+
+impl IdempotencyLockGuard {
+    async fn new(
+        key: Option<String>,
+        locks: Arc<Mutex<HashMap<String, bool>>>,
+    ) -> Result<Self, String> {
+        if let Some(ref key) = key {
+            let mut lock_map = locks.lock().await;
+            if lock_map.get(key).copied().unwrap_or(false) {
+                return Err("Idempotency key already in use".to_string());
+            }
+            lock_map.insert(key.clone(), true);
+        }
+        Ok(Self { key, locks })
+    }
+}
+
+impl Drop for IdempotencyLockGuard {
+    fn drop(&mut self) {
+        if let Some(ref key) = self.key {
+            if let Ok(mut lock_map) = self.locks.try_lock() {
+                lock_map.insert(key.clone(), false);
+                // Log that the lock was successfully released if needed
+            } else {
+                // Log an error or handle the lock acquisition failure
+                logger::error("Failed to acquire lock to release idempotency key");
+            }
+        }
+    }
+}
+
+pub(crate) fn ws_handler(
     Path((ip, port)): Path<(String, String)>,
     ws: WebSocketUpgrade,
     user_agent: Option<TypedHeader<headers::UserAgent>>,
+    idempotency_key: Option<String>,
+    locks: Arc<Mutex<HashMap<String, bool>>>,
 ) -> impl IntoResponse {
     logger::info(&format!("ip: {ip}, port: {port}"));
     if let Some(TypedHeader(user_agent)) = user_agent {
         logger::info(&format!("`{user_agent}` connected"));
     }
 
     ws.protocols(["binary"])
-        .on_upgrade(move |socket| handle_socket(socket, ip, port))
+        .on_upgrade(move |socket| handle_socket(socket, ip, port, idempotency_key, locks))
 }
 
-pub(crate) async fn handle_socket(mut socket: WebSocket, host: String, port: String) {
+pub(crate) async fn handle_socket(
+    socket: WebSocket,
+    host: String,
+    port: String,
+    idempotency_key: Option<String>,
+    locks: Arc<Mutex<HashMap<String, bool>>>,
+) {
+    let lock_guard = match IdempotencyLockGuard::new(idempotency_key, locks).await {
+        Ok(guard) => guard,
+        Err(e) => {
+            logger::error(&e);
+            return;
+        }
+    };
+
     let server_stream = match connect_to_addr(host, port, server_tcp_connect) {
         Ok(s) => s,
         Err(e) => {