ZOOKEEPER-3426: C client: Consider encoded length before completing handshake

Suhas Dantkale noticed that the C client, which uses non-blocking sockets, could end up reading only the first 40 bytes of a 41-byte connection response and concluding that its handshake was complete, leaving one byte in the pipe.

Similarly (though even less likely), the library could be left hanging waiting for a 41th byte after having received a fragmented 40-byte response from an old server.

This patch makes the logic resistant to fragmentation by considering the server-provided length encoded in the response packet.

Author: Damien Diederen <dd@crosstwine.com>

Reviewers: Mate Szalay-Beko <symat@apache.org>, Enrico Olivelli <eolivelli@apache.org>

Closes #1576 from ztzg/ZOOKEEPER-3426-c-client-short-handshake

4 files changed, 148 insertions, 4 deletions

diff --git a/zookeeper-client/zookeeper-client-c/src/zookeeper.c b/zookeeper-client/zookeeper-client-c/src/zookeeper.c
index 2b881cb9f..b5d690f35 100644
--- a/zookeeper-client/zookeeper-client-c/src/zookeeper.c
+++ b/zookeeper-client/zookeeper-client-c/src/zookeeper.c
@@ -1866,7 +1866,24 @@ static int recv_buffer(zhandle_t *zh, buffer_list_t *buff)
         /* dirty hack to make new client work against old server
          * old server sends 40 bytes to finish connection handshake,
          * while we're expecting 41 (1 byte for read-only mode data) */
-        if (buff == &zh->primer_buffer && rc == buff->len - 1) ++rc;
+        if (rc > 0 && buff == &zh->primer_buffer) {
+            /* primer_buffer's curr_offset starts at 4 (see prime_connection) */
+            int avail = buff->curr_offset - sizeof(buff->len) + rc;
+
+            /* exactly 40 bytes (out of 41 expected) collected? */
+            if (avail == buff->len - 1) {
+                int32_t reply_len;
+
+                /* extract length of ConnectResponse (+ 1-byte flag?) */
+                memcpy(&reply_len, buff->buffer, sizeof(reply_len));
+                reply_len = ntohl(reply_len);
+
+                /* if 1-byte flag was not sent, fake it (value 0) */
+                if ((int)(reply_len + sizeof(reply_len)) == buff->len - 1) {
+                    ++rc;
+                }
+            }
+        }
 
         switch(rc) {
         case 0:
diff --git a/zookeeper-client/zookeeper-client-c/tests/TestOperations.cc b/zookeeper-client/zookeeper-client-c/tests/TestOperations.cc
index ed8e9f460..ddb9533e8 100644
--- a/zookeeper-client/zookeeper-client-c/tests/TestOperations.cc
+++ b/zookeeper-client/zookeeper-client-c/tests/TestOperations.cc
@@ -36,6 +36,10 @@ class Zookeeper_operations : public CPPUNIT_NS::TestFixture
     CPPUNIT_TEST(testCloseWhileInProgressFromCompletion);
     CPPUNIT_TEST(testCloseWhileMultiInProgressFromMain);
     CPPUNIT_TEST(testCloseWhileMultiInProgressFromCompletion);
+    CPPUNIT_TEST(testConnectResponseFull);
+    CPPUNIT_TEST(testConnectResponseNoReadOnlyFlag);
+    CPPUNIT_TEST(testConnectResponseSplitAtReadOnlyFlag);
+    CPPUNIT_TEST(testConnectResponseNoReadOnlyFlagSplit);
 #else    
     CPPUNIT_TEST(testAsyncWatcher1);
     CPPUNIT_TEST(testAsyncGetOperation);
@@ -710,6 +714,126 @@ public:
         CPPUNIT_ASSERT_EQUAL((int)ZCLOSING,res2.rc_);
     }
 
+    void testConnectResponseFull()
+    {
+        CloseFinally guard(&zh);
+        Mock_socket socketMock;
+        HandshakeResponse hsResponse;
+        std::string hsResponseData = hsResponse.toString();
+
+        CPPUNIT_ASSERT_EQUAL(hsResponseData.length(), static_cast<size_t>(41));
+
+        zh = zookeeper_init("localhost:2121", watcher, 10000, NULL, NULL, 0);
+        CPPUNIT_ASSERT(zh!=0);
+
+        int rc, fd, interest;
+        timeval tv;
+
+        rc = zookeeper_interest(zh, &fd, &interest, &tv);
+        CPPUNIT_ASSERT_EQUAL(static_cast<int>(ZOK), rc);
+
+        socketMock.recvReturnBuffer = hsResponseData;
+        rc = zookeeper_process(zh, interest);
+        CPPUNIT_ASSERT_EQUAL(static_cast<int>(ZOK), rc);
+
+        CPPUNIT_ASSERT_EQUAL(ZOO_CONNECTED_STATE, static_cast<int>(zh->state));
+    }
+
+    void testConnectResponseNoReadOnlyFlag()
+    {
+        CloseFinally guard(&zh);
+        Mock_socket socketMock;
+        HandshakeResponse hsResponse;
+
+        hsResponse.omitReadOnly = true;
+
+        std::string hsResponseData = hsResponse.toString();
+
+        CPPUNIT_ASSERT_EQUAL(hsResponseData.length(), static_cast<size_t>(40));
+
+        zh = zookeeper_init("localhost:2121", watcher, 10000, NULL, NULL, 0);
+        CPPUNIT_ASSERT(zh!=0);
+
+        int rc, fd, interest;
+        timeval tv;
+
+        rc = zookeeper_interest(zh, &fd, &interest, &tv);
+        CPPUNIT_ASSERT_EQUAL(static_cast<int>(ZOK), rc);
+
+        socketMock.recvReturnBuffer = hsResponseData;
+        rc = zookeeper_process(zh, interest);
+        CPPUNIT_ASSERT_EQUAL(static_cast<int>(ZOK), rc);
+
+        CPPUNIT_ASSERT_EQUAL(ZOO_CONNECTED_STATE, static_cast<int>(zh->state));
+    }
+
+    void testConnectResponseSplitAtReadOnlyFlag()
+    {
+        CloseFinally guard(&zh);
+        Mock_socket socketMock;
+        HandshakeResponse hsResponse;
+        std::string hsResponseData = hsResponse.toString();
+
+        CPPUNIT_ASSERT_EQUAL(hsResponseData.length(), static_cast<size_t>(41));
+
+        zh = zookeeper_init("localhost:2121", watcher, 10000, NULL, NULL, 0);
+        CPPUNIT_ASSERT(zh!=0);
+
+        int rc, fd, interest;
+        timeval tv;
+
+        rc = zookeeper_interest(zh, &fd, &interest, &tv);
+        CPPUNIT_ASSERT_EQUAL(static_cast<int>(ZOK), rc);
+
+        socketMock.recvReturnBuffer = hsResponseData.substr(0, 40);
+        rc = zookeeper_process(zh, interest);
+        // Response not complete.
+        CPPUNIT_ASSERT_EQUAL(static_cast<int>(ZNOTHING), rc);
+
+        CPPUNIT_ASSERT_EQUAL(ZOO_ASSOCIATING_STATE, static_cast<int>(zh->state));
+
+        socketMock.recvReturnBuffer = hsResponseData.substr(40);
+        rc = zookeeper_process(zh, interest);
+        CPPUNIT_ASSERT_EQUAL(static_cast<int>(ZOK), rc);
+
+        CPPUNIT_ASSERT_EQUAL(ZOO_CONNECTED_STATE, static_cast<int>(zh->state));
+    }
+
+    void testConnectResponseNoReadOnlyFlagSplit()
+    {
+        CloseFinally guard(&zh);
+        Mock_socket socketMock;
+        HandshakeResponse hsResponse;
+
+        hsResponse.omitReadOnly = true;
+
+        std::string hsResponseData = hsResponse.toString();
+
+        CPPUNIT_ASSERT_EQUAL(hsResponseData.length(), static_cast<size_t>(40));
+
+        zh = zookeeper_init("localhost:2121", watcher, 10000, NULL, NULL, 0);
+        CPPUNIT_ASSERT(zh!=0);
+
+        int rc, fd, interest;
+        timeval tv;
+
+        rc = zookeeper_interest(zh, &fd, &interest, &tv);
+        CPPUNIT_ASSERT_EQUAL(static_cast<int>(ZOK), rc);
+
+        socketMock.recvReturnBuffer = hsResponseData.substr(0, 20);
+        rc = zookeeper_process(zh, interest);
+        // Response not complete.
+        CPPUNIT_ASSERT_EQUAL(static_cast<int>(ZNOTHING), rc);
+
+        CPPUNIT_ASSERT_EQUAL(ZOO_ASSOCIATING_STATE, static_cast<int>(zh->state));
+
+        socketMock.recvReturnBuffer = hsResponseData.substr(20);
+        rc = zookeeper_process(zh, interest);
+        CPPUNIT_ASSERT_EQUAL(static_cast<int>(ZOK), rc);
+
+        CPPUNIT_ASSERT_EQUAL(ZOO_CONNECTED_STATE, static_cast<int>(zh->state));
+    }
+
 #else   
     class TestGetDataJob: public TestJob{
     public:
diff --git a/zookeeper-client/zookeeper-client-c/tests/ZKMocks.cc b/zookeeper-client/zookeeper-client-c/tests/ZKMocks.cc
index d26c295f2..54c0819d9 100644
--- a/zookeeper-client/zookeeper-client-c/tests/ZKMocks.cc
+++ b/zookeeper-client/zookeeper-client-c/tests/ZKMocks.cc
@@ -344,9 +344,11 @@ string HandshakeResponse::toString() const {
     tmp=htonl(passwd_len);
     buf.append((char*)&tmp,sizeof(tmp));
     buf.append(passwd,sizeof(passwd));
-    buf.append(&readOnly,sizeof(readOnly));
+    if (!omitReadOnly) {
+        buf.append(&readOnly,sizeof(readOnly));
+    }
     // finally set the buffer length
-    tmp=htonl(buf.size()+sizeof(tmp));
+    tmp=htonl(buf.size());
     buf.insert(0,(char*)&tmp, sizeof(tmp));
     return buf;
 }
diff --git a/zookeeper-client/zookeeper-client-c/tests/ZKMocks.h b/zookeeper-client/zookeeper-client-c/tests/ZKMocks.h
index 2717ded30..48eddabbc 100644
--- a/zookeeper-client/zookeeper-client-c/tests/ZKMocks.h
+++ b/zookeeper-client/zookeeper-client-c/tests/ZKMocks.h
@@ -305,7 +305,7 @@ class HandshakeResponse: public Response
 public:
     HandshakeResponse(int64_t sessId=1):
         protocolVersion(1),timeOut(10000),sessionId(sessId),
-        passwd_len(sizeof(passwd)),readOnly(0)
+        passwd_len(sizeof(passwd)),readOnly(0),omitReadOnly(false)
     {
         memcpy(passwd,"1234567890123456",sizeof(passwd));
     }
@@ -315,6 +315,7 @@ public:
     int32_t passwd_len;
     char passwd[16];
     char readOnly;
+    bool omitReadOnly;
     virtual std::string toString() const ;
 };