Merge pull request #34 from orbisai0security/fix-tcp-auth-token-v003

Author: vickiegpt

include/tcp_communication.h

@@ -18,10 +18,24 @@ struct TCPRequest {
     uint64_t size;
     uint64_t timestamp;
     uint8_t host_id;
+    uint64_t auth_token;
     uint64_t virtual_addr;
     uint8_t data[TCP_CACHELINE_SIZE];
 } __attribute__((packed));
 
+// Returns a FNV-1a hash of CXL_CLUSTER_SECRET env var, or 0 if unset.
+static inline uint64_t tcp_get_auth_token() {
+    const char *secret = std::getenv("CXL_CLUSTER_SECRET");
+    if (!secret || secret[0] == '\0')
+        return 0;
+    uint64_t hash = 14695981039346656037ULL;
+    for (const char *p = secret; *p; ++p) {
+        hash ^= static_cast<uint8_t>(*p);
+        hash *= 1099511628211ULL;
+    }
+    return hash;
+}
+
 struct TCPResponse {
     uint8_t status;
     uint64_t latency_ns;

src/distributed_server.cpp

@@ -1616,6 +1616,7 @@ bool DistributedTCPTransport::send_message(uint32_t dst_node, const dist_message
     tcp_msg.request.size = sizeof(dist_message_t);
     tcp_msg.request.timestamp = msg.header.timestamp;
     tcp_msg.request.host_id = static_cast<uint8_t>(local_node_id_);
+    tcp_msg.request.auth_token = tcp_get_auth_token();
 
     // Serialize the dist_message into the data field (truncated to fit)
     size_t copy_size = std::min(sizeof(tcp_msg.request.data), sizeof(dist_message_t));

src/tcp_communication.cpp

@@ -167,13 +167,19 @@ int TCPServer::accept_connection() {
 }
 
 void TCPServer::handle_client() {
+    const uint64_t expected_token = tcp_get_auth_token();
     while (running_ && connected_) {
         TCPMessage recv_msg, send_msg;
 
         if (receive_message(recv_msg) < 0) {
             break;
         }
 
+        if (expected_token != 0 && recv_msg.request.auth_token != expected_token) {
+            std::cerr << "TCP auth failed: invalid token, closing connection" << std::endl;
+            break;
+        }
+
         if (message_handler_) {
             message_handler_(recv_msg, send_msg);
         } else {