From 34d1cc178fb21e1e7e0d5c7e4cedeca1abfa6da1 Mon Sep 17 00:00:00 2001
From: Picnoir <picnoir@alternativebit.fr>
Date: Thu, 21 Mar 2024 09:27:08 +0100
Subject: feat(users/picnoir/tvix-daemon): implement full handshake

Implementing the full connection handshake.

The integration test is a bit naive, but there's not much to test yet.

Tested this against cpp nix. We reach the stage where cppnix sends the
opcode.

Change-Id: I98322832848ee5b048f22105731b0adeb44b2ce0
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11227
Reviewed-by: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
---
 users/picnoir/tvix-daemon/src/main.rs | 60 +++++++++++++++++++++++++++++++++--
 1 file changed, 58 insertions(+), 2 deletions(-)

diff --git a/users/picnoir/tvix-daemon/src/main.rs b/users/picnoir/tvix-daemon/src/main.rs
index d41369e1b1..87947cee31 100644
--- a/users/picnoir/tvix-daemon/src/main.rs
+++ b/users/picnoir/tvix-daemon/src/main.rs
@@ -2,7 +2,7 @@ use anyhow::anyhow;
 use clap::Parser;
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio_listener::{self, SystemOptions, UserOptions};
-use tracing::{error, info, instrument};
+use tracing::{debug, error, info, instrument};
 
 use nix_compat::wire::primitive;
 
@@ -57,12 +57,17 @@ where
 /// Performs the initial handshake. During the handshake, the client
 /// will first send a magic u64, to which the daemon needs to respond
 /// with another magic u64.
+///
+/// We then retrieve the client version, and discard a bunch of now
+/// obsolete data.
+#[instrument()]
 async fn perform_init_handshake<'a, R: 'a>(mut conn: &'a mut R) -> anyhow::Result<()>
 where
-    &'a mut R: AsyncReadExt + AsyncWriteExt + Unpin,
+    &'a mut R: AsyncReadExt + AsyncWriteExt + Unpin + std::fmt::Debug,
 {
     let mut magic_hello = vec![0; 8];
     conn.read(&mut magic_hello).await?;
+    debug!("Hello read");
     if magic_hello != primitive::MAGIC_HELLO {
         Err(anyhow!(
             "Invalid client hello received: {:?}, expected {:?}",
@@ -71,10 +76,51 @@ where
         ))
     } else {
         conn.write(&primitive::MAGIC_HELLO_RESPONSE).await?;
+        conn.write(&primitive::PROTOCOL_VERSION).await?;
+        conn.flush().await?;
+        debug!("Hello responded");
+        let client_version = primitive::read_u32(&mut conn).await?;
+        debug!("Version read");
+        if client_version < 0x10a {
+            return Err(anyhow!("The nix client version is too old"));
+        }
+        let protocol_minor = client_version & 0x00ff;
+        let protocol_major = client_version & 0xff00;
+        debug!(client.version = %client_version, client.minor = %protocol_minor, client.major = %protocol_major);
+        if protocol_minor >= 14 {
+            debug!("read cpu affinity");
+            // Obsolete CPU affinity.
+            let read_affinity = primitive::read_u32(&mut conn).await?;
+            if read_affinity != 0 {
+                skip_8_bytes(&mut conn).await?;
+            };
+        }
+        if protocol_minor >= 11 {
+            // Obsolete reserveSpace
+            debug!("read reservespace");
+            skip_8_bytes(&mut conn).await?;
+        }
+        if protocol_minor >= 33 {
+            // Nix version. We're plain lying, we're not Nix, but eh…
+            // Setting it to the 2.3 lineage. Not 100% sure this is a
+            // good idea.
+            debug!("write version");
+            // Plain str padded to 64 bits.
+            conn.write(&"2.3.17\0\0".as_bytes()).await?;
+        }
         Ok(())
     }
 }
 
+async fn skip_8_bytes<R>(conn: &mut R) -> anyhow::Result<()>
+where
+    R: AsyncReadExt + Unpin + std::fmt::Debug,
+{
+    let mut _discard_buffer = [0; 8];
+    conn.read_exact(&mut _discard_buffer).await?;
+    Ok(())
+}
+
 #[cfg(test)]
 mod integration_tests {
     use nix_compat::wire::primitive;
@@ -83,6 +129,16 @@ mod integration_tests {
         let mut test_conn = tokio_test::io::Builder::new()
             .read(&primitive::MAGIC_HELLO)
             .write(&primitive::MAGIC_HELLO_RESPONSE)
+            .write(&primitive::PROTOCOL_VERSION)
+            // Let's say the client is in sync with the daemon
+            // protocol-wise
+            .read(&primitive::PROTOCOL_VERSION)
+            // cpu affinity
+            .read(&vec![0; 8])
+            // reservespace
+            .read(&vec![0; 8])
+            // version
+            .write(&"2.3.17\0\0".as_bytes())
             .build();
         crate::worker(&mut test_conn).await;
     }
-- 
cgit 1.4.1