From 50dbc034090614d004d097c7a45b0a28a3bbb80b Mon Sep 17 00:00:00 2001
From: Keuin <keuinx@gmail.com>
Date: Tue, 5 Sep 2023 01:52:56 +0800
Subject: feature: 0-rtt connection phase extension

---
 src/server.rs | 30 ++++++++++++++++++++++++------
 1 file changed, 24 insertions(+), 6 deletions(-)

(limited to 'src/server.rs')

diff --git a/src/server.rs b/src/server.rs
index 968bfeb..4e42105 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -5,11 +5,13 @@ use tokio::net::{TcpSocket, TcpStream};
 use tracing::{error, info};
 
 use crate::config;
+use crate::protocol::{ServiceName, PREFIX};
 
 pub struct SharedConfig {
     pub max_srv_name_length: usize,
     pub services: Vec<config::Service>,
     pub allow_help: bool,
+    pub no_ack_extension: bool,
 }
 
 const ERR_SERVICE_NAME_NOT_FOUND: &str = "-Invalid service name\r\n";
@@ -19,7 +21,13 @@ const SERVICE_HELP: &str = "HELP";
 
 pub async fn handle_client(cfg: &SharedConfig, mut sock: TcpStream, addr: SocketAddr) {
     info!("New client: {}", addr);
-    let cap = cfg.max_srv_name_length + 2;
+    let cap = cfg.max_srv_name_length
+        + 2
+        + if cfg.no_ack_extension {
+            PREFIX.len()
+        } else {
+            0
+        };
     let mut service_name_buf = vec![0_u8; cap];
     let mut i = 0;
     let mut ok = false;
@@ -64,23 +72,33 @@ pub async fn handle_client(cfg: &SharedConfig, mut sock: TcpStream, addr: Socket
         }
         return;
     }
+    let wanted_srv_name = if cfg.no_ack_extension {
+        ServiceName::from(wanted_srv_name)
+    } else {
+        ServiceName {
+            service_name: wanted_srv_name,
+            no_ack: false,
+        }
+    };
     let mut srv = None;
     for s in &cfg.services {
-        if s.name == wanted_srv_name {
+        if s.name == wanted_srv_name.service_name {
             srv = Some(s);
             break;
         }
     }
     if srv == None {
-        error!("Service name not found: {}", wanted_srv_name);
+        error!("Service name not found: {}", wanted_srv_name.to_string());
         let _ = sock.write_all(ERR_SERVICE_NAME_NOT_FOUND.as_bytes()).await;
         return;
     }
     let srv = srv.unwrap();
     info!("Client selected service: {}", srv.name);
-    if let Err(why) = sock.write_all(OK.as_bytes()).await {
-        error!("Failed to send command to client: {}", why);
-        return;
+    if !wanted_srv_name.no_ack {
+        if let Err(why) = sock.write_all(OK.as_bytes()).await {
+            error!("Failed to send command to client: {}", why);
+            return;
+        }
     }
     let remote = match srv.addr.to_socket_addrs() {
         Ok(mut addrs) => match addrs.next() {
-- 
cgit v1.2.3