40. 干货系列从零用Rust编写负载均衡及代理，websocket的实现
wenmeng库当前已支持HTTP1.1/HTTP2/WEBSOCKET，在浏览器的环境中websocket是必不可缺少的存在，当然有很多原生的服务中用的都是socket，下一章中，我们将实现websocket与tcp的互转，以便一些tcp的程序可以服务web的服务。
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的核心回调
#[async_trait]
pub trait WsTrait: Send {
/// 通过请求连接构建出返回的握手连接信息
#[inline]
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>>,
}
#[async_trait]
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>>,
}
#[async_trait]
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<dyn Trait>的形式来做回调函数的处理。
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的服务。
点击 [关注]，[在看]，[点赞] 是对作者最大的支持
