wss

2026-05-29 11:00:38 +00:00 · 2022-09-29 18:12:11 +03:00
parent dabb6d2675
commit fcc6d52221
3 changed files with 170 additions and 5 deletions
--- a/proxy/Cargo.toml
+++ b/proxy/Cargo.toml
@@ -16,12 +16,14 @@ hashbrown = "0.12"
 hex = "0.4.3"
 hmac = "0.12.1"
 hyper = "0.14"
+hyper-tungstenite = "0.8.1"
 itertools = "0.10.3"
 md5 = "0.7.0"
 once_cell = "1.13.0"
 parking_lot = "0.12"
 pin-project-lite = "0.2.7"
 rand = "0.8.3"
+regex = "1.4.5"
 reqwest = { version = "0.11", default-features = false, features = ["blocking", "json", "rustls-tls"] }
 routerify = "3"
 rustls = "0.20.0"
@@ -37,6 +39,7 @@ tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", re
 tokio-rustls = "0.23.0"
 url = "2.2.2"
 uuid = { version = "0.8.2", features = ["v4", "serde"]}
+webpki-roots = "0.22.5"
 x509-parser = "0.13.2"

 utils = { path = "../libs/utils" }
--- a/proxy/src/http/server.rs
+++ b/proxy/src/http/server.rs
@@ -1,15 +1,113 @@
+use regex::Regex;
+use rustls::{OwnedTrustAnchor, RootCertStore};
+use std::{net::TcpListener, sync::Arc};
+
 use anyhow::anyhow;
+use futures::{sink::SinkExt, stream::StreamExt};
 use hyper::{Body, Request, Response, StatusCode};
-use std::net::TcpListener;
-use utils::http::{endpoint, error::ApiError, json::json_response, RouterBuilder, RouterService};
+use hyper_tungstenite::tungstenite::Message;
+use hyper_tungstenite::HyperWebsocket;
+use tokio::io::AsyncReadExt;
+use tokio::io::AsyncWriteExt;
+use tokio::net::TcpStream;
+use utils::http::{
+    endpoint, error::ApiError, json::json_response, request::parse_request_param, RouterBuilder,
+    RouterService,
+};

 async fn status_handler(_: Request<Body>) -> Result<Response<Body>, ApiError> {
    json_response(StatusCode::OK, "")
 }

+async fn ws_handler(mut request: Request<Body>) -> Result<Response<Body>, ApiError> {
+    // Check if the request is a websocket upgrade request.
+    if hyper_tungstenite::is_upgrade_request(&request) {
+        let (response, websocket) = hyper_tungstenite::upgrade(&mut request, None)
+            .map_err(|e| ApiError::BadRequest(e.into()))?;
+
+        let connstr: String = parse_request_param(&request, "connstr")?;
+        let r = Regex::new(r"@(.*:(\d+)").unwrap();
+        let caps = r.captures(&connstr).unwrap();
+        let endpoint: String = caps.get(1).unwrap().as_str().to_owned();
+
+        // connect to itself, but through remote endpoint, lol
+        let neon_sock = TcpStream::connect(endpoint)
+            .await
+            .map_err(|e| ApiError::InternalServerError(e.into()))?;
+
+        // Spawn a task to handle the websocket connection.
+        tokio::spawn(async move {
+            if let Err(e) = serve_websocket(websocket, neon_sock).await {
+                eprintln!("Error in websocket connection: {}", e);
+            }
+        });
+
+        // Return the response so the spawned future can continue.
+        Ok(response)
+    } else {
+        json_response(StatusCode::OK, "hi")
+    }
+}
+
+async fn serve_websocket(
+    websocket: HyperWebsocket,
+    mut neon_sock: TcpStream,
+) -> anyhow::Result<()> {
+    let mut websocket = websocket.await?;
+    let mut buf = [0u8; 8192];
+
+    tokio::select! {
+        Some(message) = websocket.next() => {
+            match message? {
+                Message::Text(msg) => {
+                    println!("Received text message: {}", msg);
+                }
+                Message::Binary(msg) => {
+                    println!("Received binary message: {:02X?}", msg);
+                    neon_sock.write_all(&msg).await?;
+                    websocket
+                        .send(Message::binary(b"Thank you, come again.".to_vec()))
+                        .await?;
+                }
+                Message::Ping(msg) => {
+                    // No need to send a reply: tungstenite takes care of this for you.
+                    println!("Received ping message: {:02X?}", msg);
+                }
+                Message::Pong(msg) => {
+                    println!("Received pong message: {:02X?}", msg);
+                }
+                Message::Close(msg) => {
+                    // No need to send a reply: tungstenite takes care of this for you.
+                    if let Some(msg) = &msg {
+                        println!(
+                            "Received close message with code {} and message: {}",
+                            msg.code, msg.reason
+                        );
+                    } else {
+                        println!("Received close message");
+                    }
+                }
+                Message::Frame(_msg) => {
+                    unreachable!();
+                }
+            }
+        },
+        res = neon_sock.read(&mut buf) => {
+            let res = res?;
+            websocket
+            .send(Message::binary(&buf[0..res]))
+            .await?;
+        }
+    }
+
+    Ok(())
+}
+
 fn make_router() -> RouterBuilder<hyper::Body, ApiError> {
    let router = endpoint::make_router();
-    router.get("/v1/status", status_handler)
+    router
+        .get("/v1/status", status_handler)
+        .get("/v1/ws", ws_handler)
 }

 pub async fn thread_main(http_listener: TcpListener) -> anyhow::Result<()> {