rustdesk/hbb_common
1
0
Fork 0
mirror of https://github.com/rustdesk/hbb_common.git synced 2025-07-02 07:56:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

[test] tcp success. Websocket test. (temp ver.)

Browse Source
This commit is contained in:
YinMo19 2025-04-23 08:33:01 +08:00
parent d299e4909f
commit 6798bf3780
2 changed files with 21 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
src/socket_client.rs
View File

@ -116,18 +116,15 @@ pub async fn connect_tcp_local<
) -> ResultType<Stream> { ) -> ResultType<Stream> {
let target_str = target.to_string(); let target_str = target.to_string();
// 根据目标地址协议决定连接方式 // if target_str.starts_with("ws://") || target_str.starts_with("wss://") {
if target_str.starts_with("ws://") || target_str.starts_with("wss://") { // Ok(Stream::WebSocket(
// WebSocket 连接逻辑 // websocket::WsFramedStream::new(target_str, local, None, ms_timeout).await?,
Ok(Stream::WebSocket(websocket::WsFramedStream::new( // ))
target_str, if true {
local, Ok(Stream::WebSocket(
None, websocket::WsFramedStream::new(target_str, local, None, ms_timeout).await?,
ms_timeout, ))
)
.await?))
} else { } else {
// TCP 连接逻辑
if let Some(conf) = Config::get_socks() { if let Some(conf) = Config::get_socks() {
return Ok(Stream::Tcp( return Ok(Stream::Tcp(
FramedStream::connect(target, local, &conf, ms_timeout).await?, FramedStream::connect(target, local, &conf, ms_timeout).await?,
@ -139,8 +136,7 @@ pub async fn connect_tcp_local<
if local_addr.is_ipv6() && target_addr.is_ipv4() { if local_addr.is_ipv6() && target_addr.is_ipv4() {
let resolved_target = query_nip_io(target_addr).await?; let resolved_target = query_nip_io(target_addr).await?;
return Ok(Stream::Tcp( return Ok(Stream::Tcp(
FramedStream::new(resolved_target, Some(local_addr), ms_timeout) FramedStream::new(resolved_target, Some(local_addr), ms_timeout).await?,
.await?,
)); ));
} }
} }

13
src/websocket.rs
View File

@ -93,7 +93,7 @@ impl WsFramedStream {
pub async fn from_tcp_stream(stream: TcpStream, addr: SocketAddr) -> ResultType<Self> { pub async fn from_tcp_stream(stream: TcpStream, addr: SocketAddr) -> ResultType<Self> {
let ws_stream = let ws_stream =
WebSocketStream::from_raw_socket(MaybeTlsStream::Plain(stream), Role::Server, None) WebSocketStream::from_raw_socket(MaybeTlsStream::Plain(stream), Role::Client, None)
.await; .await;
Ok(Self { Ok(Self {
@ -169,6 +169,7 @@ impl WsFramedStream {
#[inline] #[inline]
pub async fn next(&mut self) -> Option<Result<BytesMut, Error>> { pub async fn next(&mut self) -> Option<Result<BytesMut, Error>> {
log::info!("test");
loop { loop {
match self.stream.next().await? { match self.stream.next().await? {
Ok(WsMessage::Binary(data)) => { Ok(WsMessage::Binary(data)) => {
@ -180,6 +181,16 @@ impl WsFramedStream {
} }
return Some(Ok(bytes)); return Some(Ok(bytes));
} }
Ok(WsMessage::Ping(ping)) => {
if let Err(e) = self.stream.send(WsMessage::Pong(ping)).await {
return Some(Err(Error::new(
ErrorKind::Other,
format!("Failed to send pong: {}", e),
)));
}
continue;
}
Ok(WsMessage::Close(_)) => return None,
Ok(_) => continue, Ok(_) => continue,
Err(e) => return Some(Err(Error::new(ErrorKind::Other, e))), Err(e) => return Some(Err(Error::new(ErrorKind::Other, e))),
} }