From 8d10d6d6ecd25a8316dc5a65d2b914738bad95b0 Mon Sep 17 00:00:00 2001
From: bit-aloo <sshourya17@gmail.com>
Date: Sat, 6 Jul 2024 04:21:02 -0400
Subject: [PATCH] Add flag check at setup connection

---
 roles/jd-server/src/lib/job_declarator/mod.rs | 70 +++++++++++++++----
 .../job-declarator-mock-invalid-flag.json     |  5 +-
 2 files changed, 57 insertions(+), 18 deletions(-)

diff --git a/roles/jd-server/src/lib/job_declarator/mod.rs b/roles/jd-server/src/lib/job_declarator/mod.rs
index 34d9e66de4..ff017d3d7c 100644
--- a/roles/jd-server/src/lib/job_declarator/mod.rs
+++ b/roles/jd-server/src/lib/job_declarator/mod.rs
@@ -3,12 +3,15 @@ use super::{error::JdsError, mempool::JDsMempool, status, Configuration, EitherF
 use async_channel::{Receiver, Sender};
 use binary_sv2::{B0255, U256};
 use codec_sv2::{HandshakeRole, Responder};
+use core::panic;
 use error_handling::handle_result;
 use key_utils::{Secp256k1PublicKey, Secp256k1SecretKey, SignatureService};
 use network_helpers_sv2::noise_connection_tokio::Connection;
 use nohash_hasher::BuildNoHashHasher;
 use roles_logic_sv2::{
-    common_messages_sv2::SetupConnectionSuccess,
+    common_messages_sv2::{
+        Protocol, SetupConnection, SetupConnectionError, SetupConnectionSuccess,
+    },
     handlers::job_declaration::{ParseClientJobDeclarationMessages, SendTo},
     job_declaration_sv2::{DeclareMiningJob, SubmitSolutionJd},
     parsers::{JobDeclaration, PoolMessages as JdsMessages},
@@ -450,25 +453,62 @@ impl JobDeclarator {
             if let Ok((receiver, sender, _, _)) =
                 Connection::new(stream, HandshakeRole::Responder(responder)).await
             {
-                let setup_message_from_proxy_jd = receiver.recv().await.unwrap();
-                info!(
-                    "Setup connection message from proxy: {:?}",
-                    setup_message_from_proxy_jd
-                );
-
-                let setup_connection_success_to_proxy = SetupConnectionSuccess {
-                    used_version: 2,
-                    // Setup flags for async_mining_allowed
-                    flags: 0b_0000_0000_0000_0000_0000_0000_0000_0001,
+                let address = addr.unwrap();
+                let mut incoming = match receiver.recv().await {
+                    Ok(EitherFrame::Sv2(s)) => {
+                        debug!("Got sv2 message: {:?}", s);
+                        s
+                    }
+                    Ok(EitherFrame::HandShake(s)) => {
+                        error!(
+                            "Got unexpected handshake message from upstream: {:?} at {}",
+                            s, address
+                        );
+                        continue;
+                    }
+                    Err(e) => {
+                        error!("Error receiving message: {:?}", e);
+                        continue;
+                    }
                 };
-                let sv2_frame: StdFrame =
-                    JdsMessages::Common(setup_connection_success_to_proxy.into())
+
+                let payload = incoming.payload();
+                let flag = payload[5] as u32;
+
+                let is_valid =
+                    SetupConnection::check_flags(Protocol::JobDeclarationProtocol, flag, 1);
+
+                let sv2_frame: StdFrame = if is_valid {
+                    let success_message = SetupConnectionSuccess {
+                        used_version: 2,
+                        flags: 0b_0000_0000_0000_0000_0000_0000_0000_0001,
+                    };
+                    info!("Sending success message for proxy");
+                    JdsMessages::Common(success_message.into())
+                        .try_into()
+                        .unwrap()
+                } else {
+                    let error_message = SetupConnectionError {
+                        flags: flag,
+                        error_code: "unsupported-feature-flags"
+                            .to_string()
+                            .into_bytes()
+                            .try_into()
+                            .unwrap(),
+                    };
+                    info!("Sending error message for proxy");
+                    JdsMessages::Common(error_message.into())
                         .try_into()
-                        .unwrap();
+                        .unwrap()
+                };
+
                 let sv2_frame = sv2_frame.into();
-                info!("Sending success message for proxy");
                 sender.send(sv2_frame).await.unwrap();
 
+                if !is_valid {
+                    continue;
+                }
+
                 let jddownstream = Arc::new(Mutex::new(JobDeclaratorDownstream::new(
                     receiver.clone(),
                     sender.clone(),
diff --git a/test/message-generator/mock/job-declarator-mock-invalid-flag.json b/test/message-generator/mock/job-declarator-mock-invalid-flag.json
index a14f0554bd..1091fbeb61 100644
--- a/test/message-generator/mock/job-declarator-mock-invalid-flag.json
+++ b/test/message-generator/mock/job-declarator-mock-invalid-flag.json
@@ -19,11 +19,10 @@
             "results": [
                 {
                     "type": "match_message_type",
-                    "value": "0x01"
+                    "value": "0x02"
                 }
-
             ],
-            "actiondoc": "This action sends SetupConnection and checks that .Success"
+            "actiondoc": "This action sends SetupConnection and checks that .Error"
        }
     ],
     "setup_commands": [],