Rust Reverse Proxy WebSocket

chatgpt.md

Sure! Below is a Rust program that acts as a reverse proxy in front of a website, supports WebSockets, and serves static files. We'll use the warp framework, which is a powerful and easy-to-use web framework for Rust.

---

Step 1: Create a New Rust Project

First, create a new Rust project:

cargo new reverse_proxy
cd reverse_proxy

Step 2: Update Cargo.toml

Add the necessary dependencies in your Cargo.toml file:

[dependencies]
warp = "0.3"
tokio = { version = "1", features = ["full"] }
hyper = { version = "0.14", features = ["client", "websocket"] }
bytes = "1.0"
futures = "0.3"

Step 3: Write the Code

Replace the contents of src/main.rs with the following code:

use warp::Filter;
use warp::hyper::{Body, Client, Request, Uri};
use warp::ws::{Message, WebSocket};
use std::convert::Infallible;
use futures::{StreamExt, FutureExt};
use bytes::Bytes;

#[tokio::main]
async fn main() {
    // Define the target website
    let target_website = "http://example.com".to_string();

    // Static file route (e.g., serving files from ./static directory)
    let static_files = warp::path("static")
        .and(warp::fs::dir("./static"));

    // WebSocket route
    let websocket_route = warp::path("ws")
        .and(warp::ws())
        .map(move |ws: warp::ws::Ws| {
            ws.on_upgrade(handle_websocket)
        });

    // Reverse proxy route
    let proxy_route = warp::any()
        .and(warp::method())
        .and(warp::header::headers_cloned())
        .and(warp::path::full())
        .and(warp::body::bytes())
        .and_then(move |method, headers, full_path, body| {
            let target_website = target_website.clone();
            async move {
                proxy_handler(method, headers, full_path, body, target_website).await
            }
        });

    // Combine all routes
    let routes = static_files
        .or(websocket_route)
        .or(proxy_route);

    // Start the server
    warp::serve(routes)
        .run(([0, 0, 0, 0], 3030))
        .await;
}

async fn proxy_handler(
    method: warp::http::Method,
    headers: warp::http::HeaderMap,
    full_path: warp::path::FullPath,
    body: Bytes,
    target_website: String,
) -> Result<impl warp::Reply, Infallible> {
    // Build the URI
    let uri_str = format!("{}{}", target_website, full_path.as_str());
    let uri = uri_str.parse::<Uri>().unwrap();

    // Create the client
    let client = Client::new();

    // Build the request
    let mut req = Request::builder()
        .method(method)
        .uri(uri);

    // Copy headers
    *req.headers_mut().unwrap() = headers;

    // Set the body
    let req = req
        .body(Body::from(body))
        .unwrap();

    // Send the request
    match client.request(req).await {
        Ok(res) => {
            let mut response = warp::http::Response::builder()
                .status(res.status());

            // Copy response headers
            *response.headers_mut().unwrap() = res.headers().clone();

            // Aggregate the body
            let body_bytes = hyper::body::to_bytes(res.into_body()).await.unwrap();

            Ok(response.body(body_bytes).unwrap())
        },
        Err(err) => {
            eprintln!("Error: {}", err);
            Ok(warp::http::Response::builder()
                .status(500)
                .body("Internal Server Error".into())
                .unwrap())
        }
    }
}

async fn handle_websocket(ws: WebSocket) {
    let (mut tx, mut rx) = ws.split();

    // Echo incoming messages back to the client
    while let Some(result) = rx.next().await {
        match result {
            Ok(msg) => {
                if tx.send(msg).await.is_err() {
                    break;
                }
            },
            Err(e) => {
                eprintln!("WebSocket error: {}", e);
                break;
            }
        }
    }
}

Step 4: Set Up Static Files

Create a static directory in your project root and add any static files you want to serve.

Explanation

• Reverse Proxy Handler (proxy_handler): Captures all incoming requests not matched by previous routes and forwards them to the target website. It rebuilds the HTTP request with the same method, headers, and body, then sends it using Hyper's client. The response from the target site is then returned to the client.

• WebSocket Handler (handle_websocket): Manages WebSocket connections on the /ws path. In this example, it simply echoes back any messages received. You can modify this function to proxy WebSocket messages to another server if needed.

• Static File Serving: Any requests to /static/* will serve files from the local ./static directory.

Notes

• Target Website: Replace "http://example.com" with the URL of the website you want to proxy.

• Error Handling: This is a basic implementation. For production use, you should add more robust error handling, logging, and possibly HTTPS support.

• WebSocket Proxying: If you need to proxy WebSocket connections to another server, you'll need to modify the handle_websocket function to connect to the target WebSocket server and relay messages between the client and the server.

Running the Server

Run the server using:

cargo run

The server will start listening on http://0.0.0.0:3030.

Testing the Reverse Proxy

• HTTP Requests: Navigate to http://localhost:3030/ in your browser. This should proxy the request to your target website.

• Static Files: Access static files at http://localhost:3030/static/yourfile.ext.

• WebSocket: Connect to the WebSocket server at ws://localhost:3030/ws.

Dependencies Overview

• warp: Web framework for building web servers and APIs.

• tokio: Asynchronous runtime for Rust.

• hyper: Low-level HTTP library used here for making outbound HTTP requests.

• bytes: Utilities for working with bytes across asynchronous boundaries.

• futures: Abstractions for asynchronous programming.

Customizing the Proxy

• Changing the Target Site: Update the target_website variable in main to proxy to a different site.

• Adding Middleware: You can add logging, authentication, or other middleware by composing additional filters in Warp.

• SSL/TLS Support: To serve over HTTPS, you'll need to configure TLS settings in Warp using the warp::tls module.

---

Feel free to ask if you have any questions or need further assistance customizing the proxy!