diff --git a/src/client.rs b/src/client.rs index de28e2c..078793e 100644 --- a/src/client.rs +++ b/src/client.rs @@ -14,7 +14,7 @@ use crate::types::{Template, Submission, SubmissionResponse}; use crate::new_job::NewJobConsumer; use crate::submission::{SubmissionProvider, SubmissionResponseHandler}; use crate::insecure::SkipServerVerification; - +use crate::types::CURRENT_VERSION; #[derive(Debug)] pub struct QuiverClient { @@ -45,62 +45,94 @@ impl QuiverClient { } } + async fn perform_version_handshake(&self) -> Result<()> { + info!("Negotiating protocol version..."); + let (mut send_version, mut recv_version) = self.conn.open_bi().await?; + + let ver_bytes = bincode::serialize(&CURRENT_VERSION)?; + send_version.write_u32(ver_bytes.len() as u32).await?; + send_version.write_all(&ver_bytes).await?; + send_version.finish()?; + + // read length-prefixed "ok" + let ok_len = recv_version.read_u32().await?; + let mut ok_buf = vec![0; ok_len as usize]; + recv_version.read_exact(&mut ok_buf).await?; + let response = String::from_utf8(ok_buf)?; + + if response != "ok" { + return Err(anyhow::anyhow!("Protocol negotiation failed: {}", response)); + } + info!("Protocol version compatible with server."); + Ok(()) + } + async fn serve(&mut self) -> Result<()> { - // --- Authentication Transaction --- + // 1) Version handshake + self.perform_version_handshake().await?; + + // 2) Authentication info!("authenticating..."); let (mut send_auth, mut recv_auth) = self.conn.open_bi().await?; - send_auth.write_all(self.key.as_bytes()).await?; + let api_key_bytes = self.key.as_bytes(); + send_auth.write_u32(api_key_bytes.len() as u32).await?; + send_auth.write_all(api_key_bytes).await?; send_auth.finish()?; - let auth_res = String::from_utf8(recv_auth.read_to_end(50).await?)?; + + // read length-prefixed auth reply + let n = recv_auth.read_u32().await?; + let mut buf = vec![0; n as usize]; + recv_auth.read_exact(&mut buf).await?; + let auth_res = String::from_utf8(buf)?; if auth_res != "authenticated" { return Err(anyhow::anyhow!("Authentication failed: {}", auth_res)); } - // --- Device Info Transaction --- + // 3) Device info info!("sending device info..."); let (mut send_device, mut recv_device) = self.conn.open_bi().await?; - send_device.write_all(&bincode::serialize(&self.device_info)?).await?; + let device_info_bytes = bincode::serialize(&self.device_info)?; + send_device.write_u32(device_info_bytes.len() as u32).await?; + send_device.write_all(&device_info_bytes).await?; send_device.finish()?; - let device_res = String::from_utf8(recv_device.read_to_end(50).await?)?; + + // read length-prefixed device-info reply + let n = recv_device.read_u32().await?; + let mut buf = vec![0; n as usize]; + recv_device.read_exact(&mut buf).await?; + let device_res = String::from_utf8(buf)?; if device_res != "accepted" { return Err(anyhow::anyhow!("Device info rejected: {}", device_res)); } info!("Client authenticated and ready for work."); - // Spawn the background task to receive new jobs. + // --- Your existing loop unchanged --- let mut job_handle = tokio::spawn(receive_jobs(self.conn.clone(), self.new_job_consumer.clone())); - loop { tokio::select! { - // Use a biased select to ensure we always check for connection - // closure and task failure before waiting on a new submission. biased; - // Branch 1: The connection dies for any reason. reason = self.conn.closed() => { error!("Connection closed: {:?}", reason); - job_handle.abort(); // Stop the background job stream. + job_handle.abort(); return Err(anyhow::anyhow!("Connection closed by server or network")); }, - // Branch 2: The background job receiver task fails or panics. res = &mut job_handle => { - match res { + match res { Ok(Ok(_)) => error!("Job stream closed unexpectedly."), Ok(Err(e)) => error!("Job stream failed: {}", e), Err(e) => error!("Job stream task panicked: {}", e), - } - return Err(anyhow::anyhow!("Job stream failed, disconnecting.")); + } + return Err(anyhow::anyhow!("Job stream failed, disconnecting.")); }, - // Branch 3: The submission provider (miner) has a new proof to send. submission_result = self.submission_provider.submit() => { match submission_result { Ok(submission) => { let conn = self.conn.clone(); let handler = self.submission_response_handler.clone(); - // Spawn a task to handle this specific submission. tokio::spawn(async move { if let Err(e) = handle_one_submission(conn, submission, handler).await { error!("Failed to process submission: {}", e); @@ -109,8 +141,8 @@ impl QuiverClient { } Err(e) => { error!("Submission provider failed critically: {}", e); - job_handle.abort(); // Clean up background task. - return Err(e); // Propagate critical errors. + job_handle.abort(); + return Err(e); } } }, diff --git a/src/server.rs b/src/server.rs index ae0c2d7..ffefe29 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,4 +1,4 @@ -use tracing::{info, error}; +use tracing::{info, warn, error}; use anyhow::Result; use quinn::{ServerConfig, Endpoint, Connection, TransportConfig}; use std::net::SocketAddr; @@ -14,6 +14,7 @@ use crate::device_info::{DeviceInfo, DeviceInfoUpdater}; use crate::new_job::NewJobProvider; use crate::types::{Template, Submission, SubmissionResponse, AccountInformation}; use crate::submission::SubmissionConsumer; +use crate::types::{CURRENT_VERSION, ProtocolVersion}; struct QuiverInstance { conn: Connection, @@ -45,54 +46,205 @@ impl QuiverInstance { submission_consumer, } } + + const TOKEN_PREFIX: u32 = 0x6e6f636b; // nock - async fn handle_auth_stream(&mut self) -> Result { - let Ok((mut send, mut recv)) = self.conn.accept_bi().await else { - return Err(anyhow::anyhow!("connection closed before auth")); - }; + pub async fn serve(&mut self) -> Result<()> { + // First bidi stream: either legacy auth (v0) or version handshake (v1+) + let (mut send0, mut recv0) = self.conn.accept_bi().await?; + + // Peek first 4 bytes as u32 (legacy "nock" or length of version blob) + let first_u32 = recv0.read_u32().await; + + match first_u32 { + Ok(x) if x == Self::TOKEN_PREFIX => { + // v0 legacy path + self.serve_v1(send0, recv0).await + } + Ok(len) => { + // v1+ path (len = size of ProtocolVersion) + self.serve_v2(send0, recv0, len).await + } + Err(e) => { + error!("Failed to read initial handshake bytes: {}", e); + Err(e.into()) + } + } + } - let api_key = String::from_utf8(recv.read_to_end(50).await?)?; + /// v0 (legacy) protocol: + /// - First stream carries the API key (raw, no length prefix), starting with "nock" + /// - Reply on same stream with raw "authenticated"/"unauthorized" + /// - Device info stream: raw bincode (no length), raw "accepted"/"rejected" + async fn serve_v1( + &mut self, + mut send0: quinn::SendStream, + mut recv0: quinn::RecvStream, + ) -> Result<()> { + warn!("Legacy authentication, no version handshake. Please update your miner!"); + // rebuild the full token: + let mut key = Self::TOKEN_PREFIX.to_be_bytes().to_vec(); + key.extend(recv0.read_to_end(50 - 4).await?); + let api_key = String::from_utf8(key)?; + // Authenticate & reply on same stream match self.authenticator.authenticate(&api_key).await { - Ok((account_information, guard)) => { - self.account_information = Some(account_information); + Ok((acct, guard)) => { + self.account_information = Some(acct); self.connection_guard = Some(guard); self.api_key = Some(api_key); - send.write_all(b"authenticated").await?; - send.finish()?; - Ok(true) + send0.write_all(b"authenticated").await?; + send0.finish()?; } - Err(_) => { - send.write_all(b"rejected").await?; - send.finish()?; - Ok(false) + Err(e) => { + let _ = send0.write_all(b"unauthorized").await; + let _ = send0.finish(); + return Err(e); } } + + // Device info (legacy) + if !self.handle_device_info_stream(false).await? { + return Err(anyhow::anyhow!("Device info rejected")); + } + + self.after_ready_then_work().await } - async fn handle_device_info_stream(&mut self) -> Result { + /// v2 protocol: + /// - First stream is version handshake (len-prefixed bincode), reply "ok" (len-prefixed) + /// - Second stream carries API key (len-prefixed), reply len-prefixed "authenticated"/"unauthorized" + /// - Device info stream: len-prefixed bincode, len-prefixed "accepted"/"rejected" + async fn serve_v2( + &mut self, + mut send0: quinn::SendStream, + mut recv0: quinn::RecvStream, + ver_len: u32, + ) -> Result<()> { + // Read version struct + if ver_len == 0 || ver_len >= 100 { + return Err(anyhow::anyhow!("Invalid protocol handshake length: {}", ver_len)); + } + let mut ver_buf = vec![0; ver_len as usize]; + recv0.read_exact(&mut ver_buf).await + .map_err(|_| anyhow::anyhow!("Malformed protocol handshake. Please update your miner."))?; + let client_v: ProtocolVersion = bincode::deserialize(&ver_buf) + .map_err(|_| anyhow::anyhow!("Malformed protocol handshake."))?; + if client_v.major != CURRENT_VERSION.major { + return Err(anyhow::anyhow!("Incompatible protocol version")); + } + info!("Handshake successful: v{}.{}.{}", client_v.major, client_v.minor, client_v.patch); + // Send len-prefixed "ok" and finish handshake stream + Self::lp_write(&mut send0, b"ok").await?; + send0.finish()?; + + // Auth stream + let (mut send_auth, mut recv_auth) = self.conn.accept_bi().await?; + let key = Self::lp_read(&mut recv_auth, 256).await?; // cap 256B + let api_key = String::from_utf8(key)?; + + match self.authenticator.authenticate(&api_key).await { + Ok((acct, guard)) => { + self.account_information = Some(acct); + self.connection_guard = Some(guard); + self.api_key = Some(api_key); + Self::lp_write(&mut send_auth, b"authenticated").await?; + send_auth.finish()?; + } + Err(e) => { + let _ = Self::lp_write(&mut send_auth, b"unauthorized").await; + let _ = send_auth.finish(); + return Err(e); + } + } + + // Device info v2 + if !self.handle_device_info_stream(true).await? { + return Err(anyhow::anyhow!("Device info rejected")); + } + + self.after_ready_then_work().await + } + + async fn after_ready_then_work(&mut self) -> Result<()> { + let truncated_api_key = self + .api_key + .as_ref() + .unwrap() + .chars() + .rev() + .take(8) + .collect::(); + + info!( + "{:?} device key {:?} put to work", + self.account_information.as_ref().unwrap().user_uuid, + truncated_api_key + ); + + self.handle_work().await + } + + /// len_prefixed: v2, expect/read/write length-prefixed frames + /// !len_prefixed: v1, legacy raw frames + async fn handle_device_info_stream(&mut self, len_prefixed: bool) -> Result { let Ok((mut send, mut recv)) = self.conn.accept_bi().await else { return Err(anyhow::anyhow!("connection closed before device info")); }; - let device_info = bincode::deserialize::(&recv.read_to_end(1024).await?)?; + // Read device info bytes + let device_info_bytes = if len_prefixed { + Self::lp_read(&mut recv, 1_000_000).await? // 1MB cap + } else { + // legacy raw bincode, small bound + recv.read_to_end(4096).await? + }; + + let device_info = bincode::deserialize::(&device_info_bytes)?; let Some(api_key) = self.api_key.clone() else { return Err(anyhow::anyhow!("no api key")); }; + // update & reply in matching format match self.device_info_updater.update_device_info(&device_info, &api_key).await { Ok(_) => { - send.write_all(b"accepted").await?; + if len_prefixed { + Self::lp_write(&mut send, b"accepted").await?; + } else { + send.write_all(b"accepted").await?; + } send.finish()?; Ok(true) } Err(_) => { - send.write_all(b"rejected").await?; + if len_prefixed { + Self::lp_write(&mut send, b"rejected").await?; + } else { + send.write_all(b"rejected").await?; + } send.finish()?; Ok(false) } } } + + // helpers for v2 length-prefixed messages + async fn lp_read(recv: &mut quinn::RecvStream, max_len: u32) -> Result> { + let n = recv.read_u32().await?; + if n == 0 || n > max_len { + return Err(anyhow::anyhow!("frame size out of bounds: {}", n)); + } + let mut buf = vec![0; n as usize]; + recv.read_exact(&mut buf).await?; + Ok(buf) + } + + async fn lp_write(send: &mut quinn::SendStream, bytes: &[u8]) -> Result<()> { + send.write_u32(bytes.len() as u32).await?; + send.write_all(bytes).await?; + Ok(()) + } + async fn handle_work(&mut self) -> Result<()> { // New job pusher let new_job_provider = self.new_job_provider.clone(); @@ -146,26 +298,6 @@ impl QuiverInstance { } Ok(()) } - - async fn serve(&mut self) -> Result<()> { - // The first stream MUST be for authentication. - if !self.handle_auth_stream().await? { - // If auth fails, we close the connection immediately. - return Err(anyhow::anyhow!("Authentication failed")); - } - - // After auth, receive device info. - if !self.handle_device_info_stream().await? { - return Err(anyhow::anyhow!("Device info rejected")); - } - - // take the last 8 characters of the api key - let truncated_api_key = self.api_key.as_ref().unwrap().chars().rev().take(8).collect::(); - info!("{:?} put to work with key {:?}", self.account_information.as_ref().unwrap().user_uuid, truncated_api_key); - - // Ready for work. - self.handle_work().await - } } pub async fn run( @@ -183,7 +315,7 @@ pub async fn run( let mut transport = TransportConfig::default(); transport.keep_alive_interval(Some(Duration::from_secs(2))); - transport.max_idle_timeout(Some(Duration::from_secs(5).try_into()?)); + transport.max_idle_timeout(Some(Duration::from_secs(60).try_into()?)); transport.max_concurrent_bidi_streams(quinn::VarInt::from_u32(50)); transport.max_concurrent_uni_streams(quinn::VarInt::from_u32(50)); transport.stream_receive_window(quinn::VarInt::from_u32(10_000_000)); diff --git a/src/types.rs b/src/types.rs index 22a49e9..959119a 100644 --- a/src/types.rs +++ b/src/types.rs @@ -4,6 +4,15 @@ use bytes::Bytes; use nockapp::noun::slab::NounSlab; use uuid::Uuid; +pub const CURRENT_VERSION: ProtocolVersion = ProtocolVersion { major: 0, minor: 2, patch: 0 }; + +#[derive(Serialize, Deserialize, Debug, PartialEq, Eq)] +pub struct ProtocolVersion { + pub major: u16, + pub minor: u16, + pub patch: u16, +} + #[derive(Serialize, Deserialize, Debug, Clone)] pub struct AccountInformation { pub user_uuid: Uuid,