From 5f676903739fe25bf333d1bd9f7f6c786d5fd8b6 Mon Sep 17 00:00:00 2001 From: yanghong Date: Sun, 23 Jul 2023 09:57:20 +0800 Subject: [PATCH] feat: allow config service only on client --- src/client.rs | 20 +++++++++++++++----- src/config.rs | 7 +++++++ src/protocol.rs | 34 +++++++++++++++++++++++++++++++++- src/server.rs | 34 +++++++++++++++++++++++++--------- 4 files changed, 80 insertions(+), 15 deletions(-) diff --git a/src/client.rs b/src/client.rs index 95a5a749..190a3a71 100644 --- a/src/client.rs +++ b/src/client.rs @@ -441,11 +441,21 @@ impl ControlChannel { // Read ack debug!("Reading ack"); - match read_ack(&mut conn).await? { - Ack::Ok => {} - v => { - return Err(anyhow!("{}", v)) - .with_context(|| format!("Authentication failed: {}", self.service.name)); + for _ in 0..2 { + match read_ack(&mut conn).await? { + Ack::Ok => break, + Ack::RequireServiceConfig => { + debug!("Sending client service config"); + let s = toml::to_string(&self.service).unwrap(); + let buf = s.as_bytes(); + conn.write_u32(buf.len() as u32).await?; + conn.write_all(&buf).await?; + conn.flush().await?; + } + v => { + return Err(anyhow!("{}", v)) + .with_context(|| format!("Authentication failed: {}", self.service.name)); + } } } diff --git a/src/config.rs b/src/config.rs index ca85fc20..36abd7c5 100644 --- a/src/config.rs +++ b/src/config.rs @@ -63,6 +63,7 @@ pub struct ClientServiceConfig { #[serde(skip)] pub name: String, pub local_addr: String, + pub recommend_blind_addr: Option, pub token: Option, pub nodelay: Option, pub retry_interval: Option, @@ -214,11 +215,17 @@ fn default_heartbeat_interval() -> u64 { DEFAULT_HEARTBEAT_INTERVAL_SECS } +fn default_accept_client_recommend_service() -> bool { + false +} + #[derive(Debug, Serialize, Deserialize, Default, PartialEq, Eq, Clone)] #[serde(deny_unknown_fields)] pub struct ServerConfig { pub bind_addr: String, pub default_token: Option, + #[serde(default = "default_accept_client_recommend_service")] + pub accept_client_recommend_service: bool, pub services: HashMap, #[serde(default)] pub transport: TransportConfig, diff --git a/src/protocol.rs b/src/protocol.rs index 577c7323..1abc3ce3 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -1,6 +1,6 @@ pub const HASH_WIDTH_IN_BYTES: usize = 32; -use anyhow::{bail, Context, Result}; +use anyhow::{bail, Context, Result, anyhow}; use bytes::{Bytes, BytesMut}; use lazy_static::lazy_static; use serde::{Deserialize, Serialize}; @@ -8,6 +8,8 @@ use std::net::SocketAddr; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tracing::trace; +use crate::config::{ClientServiceConfig, ServerServiceConfig}; + type ProtocolVersion = u8; const _PROTO_V0: u8 = 0u8; const PROTO_V1: u8 = 1u8; @@ -30,6 +32,7 @@ pub enum Ack { Ok, ServiceNotExist, AuthFailed, + RequireServiceConfig, } impl std::fmt::Display for Ack { @@ -41,6 +44,7 @@ impl std::fmt::Display for Ack { Ack::Ok => "Ok", Ack::ServiceNotExist => "Service not exist", Ack::AuthFailed => "Incorrect token", + Ack::RequireServiceConfig => "Try to use service config defined in client", } ) } @@ -207,6 +211,34 @@ pub async fn read_hello(conn: &mut T) -> Resu Ok(hello) } +pub async fn read_server_service_config_from_client(conn: &mut T) -> Result { + conn.write_all(&bincode::serialize(&Ack::RequireServiceConfig).unwrap()) + .await?; + conn.flush().await?; + + let n = conn.read_u32() + .await + .with_context(|| "Failed to read client service config")? as usize; + let mut buf = vec![0u8; n]; + conn.read_exact(&mut buf) + .await + .with_context(|| "Failed to read client service config")?; + + let config: ClientServiceConfig = toml::from_str(&String::from_utf8(buf)?[..]).with_context(|| "Failed to parse the config")?; + Ok( + ServerServiceConfig{ + bind_addr: match config.recommend_blind_addr { + Some(bind_addr) => bind_addr, + None => return Err(anyhow!(format!("Expect 'recommend_blind_addr' in {}", config.name))), + }, + service_type: config.service_type, + name: config.name, + nodelay: config.nodelay, + token: config.token, + } + ) +} + pub async fn read_auth(conn: &mut T) -> Result { let mut buf = vec![0u8; PACKET_LEN.auth]; conn.read_exact(&mut buf) diff --git a/src/server.rs b/src/server.rs index a36e3c21..2d8810a4 100644 --- a/src/server.rs +++ b/src/server.rs @@ -6,7 +6,7 @@ use crate::multi_map::MultiMap; use crate::protocol::Hello::{ControlChannelHello, DataChannelHello}; use crate::protocol::{ self, read_auth, read_hello, Ack, ControlChannelCmd, DataChannelCmd, Hello, UdpTraffic, - HASH_WIDTH_IN_BYTES, + HASH_WIDTH_IN_BYTES, read_server_service_config_from_client, }; use crate::transport::{SocketOpts, TcpTransport, Transport}; use anyhow::{anyhow, bail, Context, Result}; @@ -297,16 +297,34 @@ async fn do_control_channel_handshake( .await?; conn.flush().await?; + + // Read auth + let protocol::Auth(d) = read_auth(&mut conn).await?; + // Lookup the service let service_config = match services.read().await.get(&service_digest) { - Some(v) => v, + Some(v) => v.clone(), None => { - conn.write_all(&bincode::serialize(&Ack::ServiceNotExist).unwrap()) - .await?; - bail!("No such a service {}", hex::encode(service_digest)); + let op = match server_config.accept_client_recommend_service { + true => { + match read_server_service_config_from_client(&mut conn).await { // Send ACK::RequireServiceConfig + Ok(config) => Some(config), + Err(_) => None, + } + }, + false => None, + }; + + match op { + Some(config) => config, + None => { + conn.write_all(&bincode::serialize(&Ack::ServiceNotExist).unwrap()) + .await?; + bail!("No such a service {}", hex::encode(service_digest)); + } + } } - } - .to_owned(); + }; let service_name = &service_config.name; @@ -314,8 +332,6 @@ async fn do_control_channel_handshake( let mut concat = Vec::from(service_config.token.as_ref().unwrap().as_bytes()); concat.append(&mut nonce); - // Read auth - let protocol::Auth(d) = read_auth(&mut conn).await?; // Validate let session_key = protocol::digest(&concat);