用 Rust 來做簡易 Ngrok 服務（下）

繼上一篇 Ngrok Server 完成後，本篇文章會示範 Ngrok Client，主要用於將外部的 HTTP 請求轉發到本地服務，並將本地服務的 HTTP Response 發送回 Ngrok Server。

開啟新專案

由於分成兩個專案，在之後的執行與測試比較好實現，所以 Ngrok Client 會另外開一個專案來運作。

$ cargo new ngrok-client

安裝專案必要套件

[dependencies]
async-tungstenite = "0.23.0"
futures-util = "0.3.28"
reqwest = { version = "0.11.20", features = ["json", "blocking"] }
serde = "1.0.188"
serde_json = "1.0.107"
tokio = { version = "1.32.0", features = ["full"] }
tokio-tungstenite = "0.20.0"

引入套件

安裝完套件後，先一次引入一些必要的 Rust 套件。主要有用於非同步處理的 futures_util 、用於 HTTP 客戶端的 reqwest、用於 JSON 處理的 serde_json ，以及用於 WebSocket 的 tokio_tungstenite。

use futures_util::sink::SinkExt;
use futures_util::stream::StreamExt;
use reqwest::blocking::Client;
use reqwest::Method;
use serde_json::json;
use serde_json::Value;
use std::str::FromStr;
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::Message;

轉發到本地服務

定義一個函式 forward_to_local_service，用於將 HTTP 請求轉發到本地服務。

fn forward_to_local_service(
    http_request: &serde_json::Value,
) -> Result<serde_json::Value, reqwest::Error> {
    let method = Method::from_str(http_request["method"].as_str().unwrap_or("GET")).unwrap();
    let path = http_request["path"].as_str().unwrap_or("/");

    let default_headers = serde_json::Map::new();
    let headers_map = http_request["headers"]
        .as_object()
        .unwrap_or(&default_headers);

    let client = Client::new();
    let local_service_url = format!("http://localhost:3000{}", path);

    let mut headers = reqwest::header::HeaderMap::new();
    for (key, value) in headers_map.iter() {
        if let Some(val_str) = value.as_str() {
            headers.insert(
                reqwest::header::HeaderName::from_bytes(key.as_bytes()).unwrap(),
                reqwest::header::HeaderValue::from_str(val_str).unwrap(),
            );
        }
    }

    let res = client
        .request(method, &local_service_url)
        .headers(headers)
        .send()?;

    let status_code = res.status().as_u16() as u64;
    let headers_map: std::collections::HashMap<String, String> = res
        .headers()
        .iter()
        .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or_default().to_string()))
        .collect();
    let body = res.text()?;
    let response = json!({
        "status_code": status_code,
        "headers": headers_map,
        "body": body
    });

    Ok(response)
}

forward_to_local_service 函式會接收一個 http_request 參數，並回傳一個 Result，其中 Ok 代表成功，Err 代表失敗。http_request 參數是一個 serde_json::Value，代表一個 HTTP 請求，例如：

{
  "method": "GET",
  "path": "/",
  "headers": {
    "Host": "localhost:3000",
    "User-Agent": "curl/7.64.1",
    "Accept": "*/*"
  }
}

主函式 main 的實作

先使用 tokio::main 標記主函式為非同步。

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // ...（省略）
}

建立 WebSocket 連接

在 main 中，使用 TcpStream::connect 和 tokio_tungstenite::client_async 建立一個 WebSocket 連接。

let socket = TcpStream::connect("127.0.0.1:8081").await?;
let (mut ws_stream, _) = tokio_tungstenite::client_async("ws://127.0.0.1:8081/ws/", socket).await?;

這裡的 127.0.0.1:8081 是 Ngrok Server 的位置，可以根據實際情況修改。

發送建立隧道的訊息

使用 ws_stream.send 發送一個訊息來建立一個新的隧道。

ws_stream
    .send(Message::Text(
        r#"{"type": "create_tunnel", "data": {"tunnel_id": "tunnel1"}}"#.to_string(),
    ))
    .await?;

處理接收到的訊息

使用 while let 迴圈來不斷地接收和處理 WebSocket 訊息。

while let Some(message) = ws_stream.next().await {
    match message {
        Ok(Message::Text(text)) => {
            let mut parsed_message: Option<Value> = None;

            if let Ok(parsed) = serde_json::from_str::<Value>(&text) {
                parsed_message = Some(parsed);
            } else {
                eprintln!("Failed to parse the message: {}", text);
            }

            if let Some(parsed_message) = &parsed_message {
            } else {
                eprintln!("Failed to parse the message: {}", text);
            }
        }
        Err(e) => {
            eprintln!("Error during the websocket communication: {}", e);
        }
        _ => {}
    }
}

轉發 HTTP 請求和響應

在迴圈中，根據接收到的訊息類型（例如 http_request 或 create_tunnel），進行相應的處理。

if let Some(msg_type) = parsed_message.get("type").and_then(|v| v.as_str()) {
    match msg_type {
        "http_request" => {
            let http_data =
                parsed_message.get("data").expect("Invalid message format");

            // 轉發 HTTP 請求到本地服務
            match forward_to_local_service(&http_data) {
                Ok(response) => {
                    // 將本地服務的 HTTP 響應發送回 ngrok 伺服器
                    let response_message = json!({
                        "type": "http_response",
                        "data": response
                    });
                    ws_stream
                        .send(Message::Text(response_message.to_string()))
                        .await?;
                }
                Err(e) => {
                    eprintln!("Failed to forward the request: {}", e);
                }
            }
        }
        "create_tunnel" => {
            // 發送訊息給 server 確認 tunnel 建立成功
            let confirm_message = r#"{
                "type": "tunnel_created_successfully",
                "data": {
                    "tunnel_id": "tunnel_id_1234567890",
                    "status": "ok"
                }
            }"#;
            ws_stream
                .send(Message::Text(confirm_message.to_string()))
                .await?;
        }
        _ => {
            println!("Unknown message type: {}", msg_type);
        }
    }
}

這樣基本的 Ngrok Client 就完成了，接下來可以進行實際測試。

如何測試

首先啟動 Ngrok Server，並監聽 8081 port。

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    HttpServer::new(|| {
        App::new()
            .route("/ws/", web::get().to(ws_index))
            .service(web::resource("/public/").route(web::get().to(public_endpoint)))
            .default_service(web::route().to(dynamic_routing))
    })
    .bind(("127.0.0.1", 8081))?
    .run()
    .await
}

接著啟動 Ngrok Client，還有要使用 Ngrok Client 轉發的本地服務，例如任一 HTTP 伺服器，位置是 http://localhost:3000。

這時候就可以使用 curl 來測試 Ngrok 是否正常運作了。

$ curl -v http://127.0.0.1:8081

如果正常運作，應該會看到類似下面的訊息：

❯ curl -v http://127.0.0.1:8081
*   Trying 127.0.0.1:8081...
* Connected to 127.0.0.1 (127.0.0.1) port 8081 (#0)
> GET / HTTP/1.1
> Host: 127.0.0.1:8081
> User-Agent: curl/8.1.2
> Accept: */*
>
< HTTP/1.1 200 OK
< content-length: 0
< date: Wed, 27 Sep 2023 07:18:19 GMT
<
* Connection #0 to host 127.0.0.1 left intact

以上就是簡易 Ngrok 的實作，如果想要更了解 Rust 的函式庫使用方式、處理記憶體機制，歡迎參考五倍學院的線上直播課程 為你自己學 Rust。