Add localhost-only middleware to but-server

Copilot · Byron · Copilot · commit d9c72ce2eb56 · 2026-01-24T13:38:58.000Z
Co-authored-by: Byron &lt;63622+Byron@users.noreply.github.com&gt;
diff --git a/crates/but-server/src/lib.rs b/crates/but-server/src/lib.rs
@@ -5,9 +5,10 @@ use axum::{
     Json, Router,
     body::Body,
     extract::{
-        WebSocketUpgrade,
+        ConnectInfo, WebSocketUpgrade,
         ws::{Message, WebSocket},
     },
+    http::StatusCode,
     middleware::Next,
     response::IntoResponse,
     routing::{any, get},
@@ -19,6 +20,7 @@ use futures_util::{SinkExt, StreamExt as _};
 use gitbutler_project::ProjectId;
 use serde::{Deserialize, Serialize};
 use serde_json::json;
+use std::net::SocketAddr;
 use tokio::sync::Mutex;
 use tower_http::cors::{Any, CorsLayer};
 
@@ -52,6 +54,21 @@ struct AppState {
     app_settings: AppSettingsWithDiskSync,
 }
 
+/// Middleware to ensure all connections are from localhost only
+async fn localhost_only_middleware(
+    ConnectInfo(addr): ConnectInfo<SocketAddr>,
+    req: axum::extract::Request<Body>,
+    next: Next,
+) -> Result<impl IntoResponse, StatusCode> {
+    // Check if the connection is from localhost (127.0.0.1 or ::1)
+    if addr.ip().is_loopback() {
+        Ok(next.run(req).await)
+    } else {
+        tracing::warn!("Rejected non-localhost connection from: {}", addr);
+        Err(StatusCode::FORBIDDEN)
+    }
+}
+
 pub async fn run() {
     let cors = CorsLayer::new()
         .allow_methods(Any)
@@ -104,6 +121,8 @@ pub async fn run() {
                 tokio::task::spawn(next.run(req)).await.unwrap()
             },
         ))
+        // Middleware to ensure only localhost connections are accepted
+        .layer(axum::middleware::from_fn(localhost_only_middleware))
         .layer(cors)
         .with_state(state);
 
@@ -112,7 +131,12 @@ pub async fn run() {
     let url = format!("{host}:{port}");
     let listener = tokio::net::TcpListener::bind(&url).await.unwrap();
     println!("Running at {url}");
-    axum::serve(listener, app).await.unwrap();
+    axum::serve(
+        listener,
+        app.into_make_service_with_connect_info::<SocketAddr>(),
+    )
+    .await
+    .unwrap();
 }
 
 async fn post_handle_json_command(
