35. 乾貨系列從零用Rust編寫負載均衡及代理,代理伺服器的原始碼升級改造

2023-12-19 09:00:23

wmproxy

wmproxy已用Rust實現http/https代理, socks5代理, 反向代理, 靜態檔案伺服器,四層TCP/UDP轉發,七層負載均衡,內網穿透,後續將實現websocket代理等,會將實現過程分享出來,感興趣的可以一起造個輪子

專案地址

國內: https://gitee.com/tickbh/wmproxy

github: https://github.com/tickbh/wmproxy

專案設計目標

在同一個埠上同時支援HTTP/HTTPS/SOCKS5代理,即假設監聽8090埠,那麼可以設定如下:

curl --proxy socks5://127.0.0.1:8090 http://www.baidu.com

curl --proxy http://127.0.0.1:8090 http://www.baidu.com

curl --proxy http://127.0.0.1:8090 https://www.baidu.com

以上方案需要都可以相容打通,才算成功。

初始方案

不做HTTP伺服器,僅簡單的解析資料流,然後進行資料轉發

pub async fn process<T>(
    username: &Option<String>,
    password: &Option<String>,
    mut inbound: T,
) -> Result<(), ProxyError<T>>
where
    T: AsyncRead + AsyncWrite + Unpin,
{
    let mut outbound;
    let mut request;
    let mut buffer = BinaryMut::new();
    loop {
        let size = {
            let mut buf = ReadBuf::uninit(buffer.chunk_mut());
            inbound.read_buf(&mut buf).await?;
            buf.filled().len()
        };

        if size == 0 {
            return Err(ProxyError::Extension("empty"));
        }
        unsafe {
            buffer.advance_mut(size);
        }
        request = webparse::Request::new();
        // 通過該方法解析檔頭是否合法, 若是partial(部分)則繼續讀資料
        // 若解析失敗, 則表示非http協定能處理, 則丟擲錯誤
        // 此處clone為淺拷貝,不確定是否一定能解析成功,不能影響偏移
        match request.parse_buffer(&mut buffer.clone()) {
            Ok(_) => match request.get_connect_url() {
                Some(host) => {
                    match HealthCheck::connect(&host).await {
                        Ok(v) => outbound = v,
                        Err(e) => {
                            Self::err_server_status(inbound, 503).await?;
                            return Err(ProxyError::from(e));
                        }
                    }
                    break;
                }
                None => {
                    if !request.is_partial() {
                        Self::err_server_status(inbound, 503).await?;
                        return Err(ProxyError::UnknownHost);
                    }
                }
            },
            Err(WebError::Http(HttpError::Partial)) => {
                continue;
            }
            Err(_) => {
                return Err(ProxyError::Continue((Some(buffer), inbound)));
            }
        }
    }

    match request.method() {
        &Method::Connect => {
            log::trace!(
                "https connect {:?}",
                String::from_utf8_lossy(buffer.chunk())
            );
            inbound.write_all(b"HTTP/1.1 200 OK\r\n\r\n").await?;
        }
        _ => {
            outbound.write_all(buffer.chunk()).await?;
        }
    }
    let _ = copy_bidirectional(&mut inbound, &mut outbound).await?;
    Ok(())
}

此方案僅做淺解析,處理相當高效,但遇到如下問題:

  • HTTP/HTTPS代理伺服器需要驗證密碼
  • HTTP服務存在不同的協定,此方法只相容HTTP/1.1,無法相容明確的HTTP/2協定
  • 請求的協定頭有些得做修改,此方法無法修改

改造方案

  • 引入HTTP伺服器介入
  • 但是因為需要相容不同協定,只有等確定協定後才能引入協定,需要預讀資料,進行協定判定。
  • HTTPS代理協定只處理一組Connect協定,之後需要解除http協定進行雙向繫結。

完整原始碼

  1. 預讀資料
  • Socks5:第一個位元組為0X05,非ascii字元,其它協定不會影響
  • Https: https代理必須傳送Connect方法,所以必須以CONNECT或者connect開頭,且查詢其它HTTP方法沒有以C開頭的,這裡僅判斷第一個字元為C或者c,該協定僅處理一條http請求不參與後續TLS握手協定等保證資料安全
  • 其它開頭的均被認為http代理
let mut buffer = BinaryMut::with_capacity(24);
let size = {
    let mut buf = ReadBuf::uninit(buffer.chunk_mut());
    inbound.read_buf(&mut buf).await?;
    buf.filled().len()
};

if size == 0 {
    return Err(ProxyError::Extension("empty"));
}
unsafe {
    buffer.advance_mut(size);
}
// socks5 協定, 直接返回, 交給socks5層處理
if buffer.as_slice()[0] == 5 {
    return Err(ProxyError::Continue((Some(buffer), inbound)));
}

let mut max_req_num = usize::MAX;
// https 協定, 以connect開頭, 僅處理一條HTTP請求
if buffer.as_slice()[0] == b'C' || buffer.as_slice()[0] == b'c' {
    max_req_num = 1;
}
  1. 構建HTTP伺服器,構建服務類:
/// http代理類處理類
struct Operate {
    /// 使用者名稱
    username: Option<String>,
    /// 密碼
    password: Option<String>,
    /// Stream類, https連線後給後續https使用
    stream: Option<TcpStream>,
    /// http代理keep-alive的複用
    sender: Option<Sender<RecvRequest>>,
    /// http代理keep-alive的複用
    receiver: Option<Receiver<ProtResult<RecvResponse>>>,
}

構建HTTP服務

// 需要將已讀的資料buffer重新加到server的已讀cache中, 否則解析會出錯
let mut server = Server::new_by_cache(inbound, None, buffer);
// 構建HTTP服務回撥
let mut operate = Operate {
    username: username.clone(),
    password: password.clone(),
    stream: None,
    sender: None,
    receiver: None,
};
server.set_max_req(max_req_num);
let _e = server.incoming(&mut operate).await?;
if let Some(outbound) = &mut operate.stream {
    let mut inbound = server.into_io();
    let _ = copy_bidirectional(&mut inbound, outbound).await?;
}

此時我們已將資料用HTTP服務進行處理,收到相應的請求再進行給遠端做轉發:

HTTP核心處理回撥,此處我們用的是async_trait非同步回撥


#[async_trait]
impl OperateTrait for &mut Operate {
    async fn operate(&mut self, request: &mut RecvRequest) -> ProtResult<RecvResponse> {
        // 已連線直接進行後續處理
        if let Some(sender) = &self.sender {
            sender.send(request.replace_clone(Body::empty())).await?;
            if let Some(res) = self.receiver.as_mut().unwrap().recv().await {
                return Ok(res?)
            }
            return Err(ProtError::Extension("already close by other"))
        }
        // 獲取要連線的物件
        let stream = if let Some(host) = request.get_connect_url() {
            match HealthCheck::connect(&host).await {
                Ok(v) => v,
                Err(e) => {
                    return Err(ProtError::from(e));
                }
            }
        } else {
            return Err(ProtError::Extension("unknow tcp stream"));
        };

        // 賬號密碼存在,將獲取`Proxy-Authorization`進行校驗,如果檢驗錯誤返回407協定
        if self.username.is_some() && self.password.is_some() {
            let mut is_auth = false;
            if let Some(auth) = request.headers_mut().remove(&"Proxy-Authorization") {
                if let Some(val) = auth.as_string() {
                    is_auth = self.check_basic_auth(&val);
                }
            }
            if !is_auth {
                return Ok(Response::builder().status(407).body("")?.into_type());
            }
        }

        // 判斷使用者協定
        match request.method() {
            &Method::Connect => {
                // https返回200內容直接進行遠端和使用者端的雙向繫結
                self.stream = Some(stream);
                return Ok(Response::builder().status(200).body("")?.into_type());
            }
            _ => {
                // http協定,需要將使用者端的內容轉發到伺服器端,並將伺服器端資料轉回使用者端
                let client = Client::new(ClientOption::default(), MaybeHttpsStream::Http(stream));
                let (mut recv, sender) = client.send2(request.replace_clone(Body::empty())).await?;
                match recv.recv().await {
                    Some(res) => {
                        self.sender = Some(sender);
                        self.receiver = Some(recv);
                        return Ok(res?)
                    },
                    None => return Err(ProtError::Extension("already close by other")),
                }
            }
        }

    }
}

密碼校驗,由Basic的密碼加密方法,先用base64解密,再用:做拆分,再與使用者密碼比較

pub fn check_basic_auth(&self, value: &str) -> bool
{
    use base64::engine::general_purpose;
    use std::io::Read;

    let vals: Vec<&str> = value.split_whitespace().collect();
    if vals.len() == 1 {
        return false;
    }

    let mut wrapped_reader = Cursor::new(vals[1].as_bytes());
    let mut decoder = base64::read::DecoderReader::new(
        &mut wrapped_reader,
        &general_purpose::STANDARD);
    // handle errors as you normally would
    let mut result: Vec<u8> = Vec::new();
    decoder.read_to_end(&mut result).unwrap();

    if let Ok(value) = String::from_utf8(result) {
        let up: Vec<&str> = value.split(":").collect();
        if up.len() != 2 {
            return false;
        }
        if up[0] == self.username.as_ref().unwrap() ||
            up[1] == self.password.as_ref().unwrap() {
            return true;
        }
    }

    return false;
}

小結

代理在計算機網路很常見,比如伺服器群組內部通常只會開一個口進行對外存取,就可以通過內網代理來進行處理,從而更好的保護內網伺服器。代理讓我們網路更安全,但是警惕非正規的代理可能會竊取您的資料。請用HTTPS內容存取更安全。

點選 [關注][在看][點贊] 是對作者最大的支援