渣渣兔 发表于 2024-3-29 17:26:23

40. 干货系列从零用Rust编写负载均衡及代理,websocket的实现

wmproxy

wmproxy已用Rust实现http/https代理, socks5代理, 反向代理, 静态文件服务器,四层TCP/UDP转发,七层负载均衡,内网穿透,后续将实现websocket代理等,会将实现过程分享出来,感兴趣的可以一起造个轮子
项目地址

国内: https://gitee.com/tickbh/wmproxy
github: https://github.com/tickbh/wmproxy
简单介绍websocket

WebSocket 是一种在单个 TCP 连接上进行全双工通信的协议,它使得客户端和服务器之间的数据交换变得更加简单,允许服务端主动向客户端推送数据。WebSocket 通信协议于 2011 年被 IETF 定为标准 RFC 6455,并由 RFC7936 补充规范。WebSocket API 也被 W3C 定为标准。
也就是在web环境中,websocket就是socket的一种标准形式的体现。类似的还要SSE基于HTTP中的text/event-stream
源码文件含义

协议层的编码解码主要在webparse/ws

[*]frame_header协议头的解码与编码
[*]dataframe 基础单位为帧,存在多帧组成一个数据包的情况
[*]message 协议的基本信息,包含Text必须为UTF-8字符串文本,Binary二进制数据流,Close关闭信息,Ping,Pong用来做心跳包相关的信息。
[*]mask 是否为数据进行基本的加密,服务端要求客户端传来的数据必须加密
网络处理层的源码主要在wenmeng/ws

[*]codec/framed_read 每一帧的读,以帧为单位进行读取
[*]codec/framed_write 每一帧的写,以帧为单位进行写入
[*]state/state_handshake websocket连接内部的握手状态
[*]client_connection 客户端的状态连接
[*]server_connection 服务端的状态连接
[*]control 状态的控制,写入读取的pending等,核心处理源码
[*]handshake 定义on_open回调后的WsHandshake类
[*]option 定义on_open回调后返回的WsOption类,当下只包含定时器,即客户端多久时间唤醒一次interval
[*]ws_trait websocket的核心回调
#
pub trait WsTrait: Send {
    /// 通过请求连接构建出返回的握手连接信息
    #
    fn on_request(&mut self, req: &RecvRequest) -> ProtResult<RecvResponse> {
      // warn!("Handler received request:\n{}", req);
      WsHandshake::build_request(req)
    }

    /// 握手完成后之后的回调,服务端返回了Response之后就认为握手成功
    fn on_open(&mut self, shake: WsHandshake) -> ProtResult<Option<WsOption>>;

    /// 接受到远端的关闭消息
    async fn on_close(&mut self, reason: &Option<CloseData>) {}

    /// 服务内部出现了错误代码
    async fn on_error(&mut self, err: ProtError) {}

    /// 收到来在远端的ping消息, 默认返回pong消息
    async fn on_ping(&mut self, val: Vec<u8>) -> ProtResult<OwnedMessage> {
      return Ok(OwnedMessage::Pong(val));
    }

    /// 收到来在远端的pong消息, 默认不做任何处理, 可自定义处理如ttl等
    async fn on_pong(&mut self, val: Vec<u8>) {}

    /// 收到来在远端的message消息, 必须覆写该函数
    async fn on_message(&mut self, msg: OwnedMessage) -> ProtResult<()>;

    /// 定时器定时按间隔时间返回
    async fn on_interval(&mut self, option: &mut Option<WsOption>) -> ProtResult<()> {
      Ok(())
    }
   
    /// 将当前trait转化成Any,仅需当需要重新获取回调处理的时候进行处理
    fn as_any(&self) -> Option<&dyn Any> {
      None
    }

    /// 将当前trait转化成mut Any,仅需当需要重新获取回调处理的时候进行处理
    fn as_any_mut(&mut self) -> Option<&mut dyn Any> {
      None
    }
}服务端基础demo

建立一个本地监听8081的ws端口,完整源码ws_server
建立监听类:
struct Operate {
    sender: Option<Sender<OwnedMessage>>,
}

#
impl WsTrait for Operate {
    fn on_open(&mut self, shake: WsHandshake) -> ProtResult<Option<WsOption>> {
      self.sender = Some(shake.sender);
      Ok(Some(WsOption::new(Duration::from_secs(10))))
    }

    async fn on_message(&mut self, msg: OwnedMessage) -> ProtResult<()> {
      println!("callback on message = {:?}", msg);
      let _ = self
            .sender
            .as_mut()
            .unwrap()
            .send(OwnedMessage::Text("from server".to_string()))
            .await;
      let _ = self.sender.as_mut().unwrap().send(msg).await;
      Ok(())
    }

    async fn on_interval(&mut self, _option: &mut Option<WsOption>) -> ProtResult<()> {
      println!("on_interval!!!!!!!");
      Ok(())
    }
}然后启动服务器监听:
async fn run_main() -> Result<(), Box<dyn Error>> {
    let addr = "127.0.0.1:8081".to_string();
    let server = TcpListener::bind(&addr).await?;
    println!("Listening on: {}", addr);
    loop {
      let (stream, addr) = server.accept().await?;
      tokio::spawn(async move {
            let mut server = Server::new(stream, Some(addr));
            let operate = Operate { sender: None };
            // 设置服务回调
            server.set_callback_ws(Box::new(operate));
            let e = server.incoming().await;
            println!("close server ==== addr = {:?} e = {:?}", addr, e);
      });
    }
}此时即可实现websocket的监听及处理。
客户端demo

当下客户端demo需要能接受终端的输入,并向服务器发送数据,所以需要自己构建sender
建立客户端连接,在这里我们手动构建了一个sender/receiver对。
async fn run_main() -> ProtResult<()> {
    // 自己手动构建数据对,并将receiver传给服务端
    let (sender, receiver) = channel(10);
    let sender_clone = sender.clone();
    tokio::spawn(async move {
      let url = "ws://127.0.0.1:8081";
      let mut client = Client::builder()
            .url(url)
            .unwrap()
            .connect()
            .await
            .unwrap();
      client.set_callback_ws(Box::new(Operate { sender:Some(sender_clone), receiver: Some(receiver) }));
      client.wait_ws_operate().await.unwrap();
    });
    loop {
      let mut buffer = String::new();
      let stdin = io::stdin(); // We get `Stdin` here.
      stdin.read_line(&mut buffer)?;
      sender.send(OwnedMessage::Text(buffer)).await?;
    }
    Ok(())
}监听实现
struct Operate {
    sender: Option<Sender<OwnedMessage>>,
    receiver: Option<Receiver<OwnedMessage>>,
}

#
impl WsTrait for Operate {
    fn on_open(&mut self, shake: WsHandshake) -> ProtResult<Option<WsOption>> {
      // 将receiver传给控制中心, 以让其用该receiver做接收
      let mut option = WsOption::new(Duration::from_secs(1000));
      if self.receiver.is_some() {
            option.set_receiver(self.receiver.take().unwrap());
      }
      if self.sender.is_none() {
            self.sender = Some(shake.sender);
      }
      Ok(Some(option))
    }

    async fn on_message(&mut self, msg: OwnedMessage) -> ProtResult<()> {
      println!("callback on message = {:?}", msg);
      let _ = self
            .sender
            .as_mut()
            .unwrap()
            .send(OwnedMessage::Text("from client".to_string()))
            .await;
      let _ = self.sender.as_mut().unwrap().send(msg).await;
      Ok(())
    }

    async fn on_interval(&mut self, _option: &mut Option<WsOption>) -> ProtResult<()> {
      println!("on_interval!!!!!!!");
      Ok(())
    }
}接口说明

Client和Server为了同时兼容HTTP服务,即握手用的为HTTP的前半段请求,选择了回调用Box的形式来做回调函数的处理。
pub struct Server<T>
where
    T: AsyncRead + AsyncWrite + Unpin + Sized,
{
    /// http的接口回调, 处理http服务器
    callback_http: Option<Box<dyn HttpTrait>>,
    /// websocket的接口回调, 处理websocket服务器
    callback_ws: Option<Box<dyn WsTrait>>,
    // ...
}他们两个可能是同时存在,或者单个存在的,即当作服务的时候,可能仅对/ws进行websocket的升级,其它的仅仅是http服务,所以需要能单独又能聚合的处理数据。而单存的websocket仅需WsTrait回调。
即在pub async fn incoming(&mut self) -> ProtResult处理服务的时候不在传入回调地址,改成预先设置。达到灵活处理的目的。且接口比较清晰。
小结

wenmeng库当前已支持HTTP1.1/HTTP2/WEBSOCKET,在浏览器的环境中websocket是必不可缺少的存在,当然有很多原生的服务中用的都是socket,下一章中,我们将实现websocket与tcp的互转,以便一些tcp的程序可以服务web的服务。
点击 [关注],[在看],[点赞] 是对作者最大的支持

免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作!
页: [1]
查看完整版本: 40. 干货系列从零用Rust编写负载均衡及代理,websocket的实现