mirror of https://github.com/rapiz1/rathole.git
refactor: refactor handle_connection
This commit is contained in:
parent
8f3bf5c7c7
commit
f4b7e600bc
|
@ -129,6 +129,23 @@ async fn handle_connection(
|
|||
let hello = read_hello(&mut conn).await?;
|
||||
match hello {
|
||||
ControlChannelHello(_, service_digest) => {
|
||||
do_control_channel_handshake(conn, addr, services, control_channels, service_digest)
|
||||
.await?;
|
||||
}
|
||||
DataChannelHello(_, nonce) => {
|
||||
do_data_channel_handshake(conn, control_channels, nonce).await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn do_control_channel_handshake(
|
||||
mut conn: TcpStream,
|
||||
addr: SocketAddr,
|
||||
services: Arc<RwLock<HashMap<ServiceDigest, ServerServiceConfig>>>,
|
||||
control_channels: Arc<RwLock<ControlChannelMap>>,
|
||||
service_digest: ServiceDigest,
|
||||
) -> Result<()> {
|
||||
info!("New control channel incomming from {}", addr);
|
||||
|
||||
// Generate a nonce
|
||||
|
@ -198,8 +215,15 @@ async fn handle_connection(
|
|||
// Drop the old handle
|
||||
let _ = h.insert(service_digest, session_key, handle);
|
||||
}
|
||||
}
|
||||
DataChannelHello(_, nonce) => {
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn do_data_channel_handshake(
|
||||
conn: TcpStream,
|
||||
control_channels: Arc<RwLock<ControlChannelMap>>,
|
||||
nonce: Nonce,
|
||||
) -> Result<()> {
|
||||
// Validate
|
||||
let control_channels_guard = control_channels.read().await;
|
||||
match control_channels_guard.get2(&nonce) {
|
||||
|
@ -215,8 +239,6 @@ async fn handle_connection(
|
|||
warn!("Data channel has incorrect nonce");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
|
Loading…
Reference in New Issue