rustdesk/hbb_common
1
0
Fork 0
mirror of https://github.com/rustdesk/hbb_common.git synced 2025-07-01 15:36:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

[bug fix] add logic

Browse Source
This commit is contained in:
YinMo19 2025-04-24 21:25:00 +08:00
parent 7110634c09
commit 608eb5983f
2 changed files with 4 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/tcp.rs
View File

@ -25,7 +25,7 @@ pub trait TcpStreamTrait: AsyncRead + AsyncWrite + Unpin {}
pub struct DynTcpStream(pub(crate) Box<dyn TcpStreamTrait + Send + Sync>);
#[derive(Clone)]
pub struct Encrypt(Key, u64, u64);
pub struct Encrypt(pub Key, pub u64, pub u64);
pub struct FramedStream(
pub(crate) Framed<DynTcpStream, BytesCodec>,

51
src/websocket.rs
View File

@ -1,3 +1,4 @@
use crate::tcp::Encrypt;
use crate::{
config::Socks5Server,
protobuf::Message,
@ -18,9 +19,6 @@ use tokio_tungstenite::{
use tungstenite::client::IntoClientRequest;
use tungstenite::protocol::Role;
#[derive(Clone)]
pub struct Encrypt(Key, u64, u64);
pub struct WsFramedStream {
stream: WebSocketStream<MaybeTlsStream<TcpStream>>,
addr: SocketAddr,
@ -176,19 +174,10 @@ impl WsFramedStream {
log::info!("Received binary data ({} bytes)", data.len());
let mut bytes = BytesMut::from(&data[..]);
if let Some(key) = self.encrypt.as_mut() {
log::debug!("Decrypting data with seq: {}", key.2);
match key.dec(&mut bytes) {
Ok(_) => {
log::debug!("Decryption successful");
return Some(Ok(bytes));
}
Err(e) => {
log::error!("Decryption failed: {}", e);
return Some(Err(e));
}
if let Err(err) = key.dec(&mut bytes) {
return Some(Err(err));
}
}
log::error!("not encrypt set.");
return Some(Ok(bytes));
}
WsMessage::Text(text) => {
@ -218,37 +207,3 @@ impl WsFramedStream {
}
}
}
impl Encrypt {
pub fn new(key: Key) -> Self {
Self(key, 0, 0)
}
pub fn dec(&mut self, bytes: &mut BytesMut) -> Result<(), Error> {
if bytes.len() <= 1 {
return Ok(());
}
self.2 += 1;
let nonce = get_nonce(self.2);
match secretbox::open(bytes, &nonce, &self.0) {
Ok(res) => {
bytes.clear();
bytes.put_slice(&res);
Ok(())
}
Err(()) => Err(Error::new(ErrorKind::Other, "decryption error")),
}
}
pub fn enc(&mut self, data: &[u8]) -> Vec<u8> {
self.1 += 1;
let nonce = get_nonce(self.1);
secretbox::seal(data, &nonce, &self.0)
}
}
fn get_nonce(seqnum: u64) -> Nonce {
let mut nonce = Nonce([0u8; secretbox::NONCEBYTES]);
nonce.0[..std::mem::size_of_val(&seqnum)].copy_from_slice(&seqnum.to_le_bytes());
nonce
}