Merge branch 'main' of github.com:Tensamin/Omega

src/server/server.rs

@@ -1,3 +1,4 @@
+<<<<<<< HEAD
 use base64::Engine;
 use base64::engine::general_purpose::STANDARD;
 use bytes::Bytes;
@@ -17,12 +18,27 @@ use std::result::Result::Ok;
 use std::{future::Future, pin::Pin, time::Duration};
 use tokio::net::TcpListener;
 use tower::Service;
+=======
+use axum::{
+    Router,
+    body::Body,
+    extract::{ConnectInfo, OriginalUri, Path, ws::WebSocketUpgrade},
+    response::{IntoResponse, Redirect},
+    routing::get,
+};
+
+use pnet::datalink::NetworkInterface;
+use std::net::SocketAddr;
+use std::time::Duration;
+use tokio::net::TcpListener;
+>>>>>>> 7f78c8669b36cbe39755d69cccd6971e56e10290
 
 use crate::log;
 use crate::server::api;
 use crate::server::short_link::get_short_link;
 use crate::server::socket;
 
+<<<<<<< HEAD
 // --- ApiService for HTTP/2 ---
 
 #[derive(Clone)]
@@ -34,14 +50,28 @@ impl Service<HttpRequest<Incoming>> for ApiService {
     type Response = HttpResponse<Full<Bytes>>;
     type Error = io::Error;
     type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
+=======
+pub async fn start(port: u16) -> bool {
+    let app = Router::new()
+        .route("/ws/omikron", get(ws_handler))
+        .route("/direct/{short}", get(direct_handler))
+        .fallback(fallback_handler);
+    run_http_server(port, app).await
+}
 
-    fn poll_ready(
-        &mut self,
-        _cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<Result<(), Self::Error>> {
-        std::task::Poll::Ready(std::io::Result::Ok(()))
-    }
+async fn ws_handler(
+    ws: WebSocketUpgrade,
+    OriginalUri(uri): OriginalUri,
+    ConnectInfo(_): ConnectInfo<SocketAddr>,
+) -> impl IntoResponse {
+    log!("Attempting WebSocket upgrade on {}", uri.path());
+    let path = uri.path().to_string();
+>>>>>>> 7f78c8669b36cbe39755d69cccd6971e56e10290
 
+    ws.on_upgrade(async move |socket| socket::handle(path, socket))
+}
+
+<<<<<<< HEAD
     fn call(&mut self, req: HttpRequest<Incoming>) -> Self::Future {
         let (parts, body) = req.into_parts();
         let path = parts.uri.path().to_string();
@@ -286,4 +316,74 @@ fn calculate_accept_key(key: &str) -> String {
     sha1.update(websocket_guid.as_bytes());
     let result = sha1.finalize();
     STANDARD.encode(result)
+=======
+async fn direct_handler(Path(short): Path<String>) -> impl IntoResponse {
+    match get_short_link(&short).await {
+        Ok(long) => Redirect::temporary(&long),
+        Err(_) => Redirect::temporary("https://tensamin.net"),
+    }
+}
+
+async fn fallback_handler(
+    OriginalUri(uri): OriginalUri,
+    headers: axum::http::HeaderMap,
+    body: Body,
+) -> impl IntoResponse {
+    let path = uri.path().to_string();
+
+    let whole_body = tokio::time::timeout(
+        Duration::from_secs(10),
+        axum::body::to_bytes(body, 1024 * 1024 * 10),
+    )
+    .await;
+
+    let body_string = match whole_body {
+        Ok(Ok(bytes)) => String::from_utf8(bytes.to_vec()).ok(),
+        _ => None,
+    };
+
+    api::handle(&path, headers, body_string).await
+}
+
+async fn run_http_server(port: u16, app: Router) -> bool {
+    let ip = find_local_ip();
+    let listener = match TcpListener::bind(format!("0.0.0.0:{}", port)).await {
+        Ok(l) => l,
+        Err(e) => {
+            log!("Failed to bind to port {}: {:?}", port, e);
+            return false;
+        }
+    };
+
+    log!(
+        "Standard Server listening for HTTP and WS on {}:{}",
+        ip,
+        port
+    );
+
+    axum::serve(
+        listener,
+        app.into_make_service_with_connect_info::<SocketAddr>(),
+    )
+    .await
+    .map(|_| true)
+    .unwrap_or_else(|e| {
+        log!("Server error: {}", e);
+        false
+    })
+}
+
+fn find_local_ip() -> String {
+    for iface in pnet::datalink::interfaces() {
+        let iface: NetworkInterface = iface;
+        if !iface.ips.is_empty() {
+            let ipsv = format!("{}", iface.ips[0]);
+            let ips: &str = ipsv.split('/').next().unwrap();
+            if ips.starts_with("10.") || ips.starts_with("192.") {
+                return ips.to_string();
+            }
+        }
+    }
+    "0.0.0.0".to_string()
+>>>>>>> 7f78c8669b36cbe39755d69cccd6971e56e10290
 }