CCBC-1610: fix collection id encoding

* use packet flags only to decide whether to decode collection id
* fix header length calculation when overriding collection id
* prefer standalone network buffer when overriding collection id

Change-Id: Ib1c94d1dcca1ea6bf9b51eb8437c1f718abf8741
Reviewed-on: https://review.couchbase.org/c/libcouchbase/+/196764
Reviewed-by: Brett Lawson <brett19@gmail.com>
Tested-by: Build Bot <build@couchbase.com>

Author: avsej

src/handler.cc

@@ -271,7 +271,7 @@ void init_resp(lcb_INSTANCE *instance, mc_PIPELINE *pipeline, const MemcachedRes
     resp->cookie = const_cast<void *>(MCREQ_PKT_COOKIE(req));
     const char *key = nullptr;
     size_t key_len = 0;
-    mcreq_get_key(instance, req, &key, &key_len);
+    mcreq_get_key(req, &key, &key_len);
     if (key != nullptr) {
         resp->ctx.key.assign(key, key_len);
     }

src/mc/mcreq.c

@@ -376,7 +376,7 @@ void mcreq_wipe_packet(mc_PIPELINE *pipeline, mc_PACKET *packet)
         return;
     }
 
-    if ((packet->flags & MCREQ_F_DETACHED) || IS_STANDALONE_SPAN(&packet->kh_span)) {
+    if (packet->flags & MCREQ_F_DETACHED) {
         free(SPAN_BUFFER(&packet->u_value.single));
     } else {
         netbuf_mblock_release(&pipeline->nbmgr, &packet->u_value.single);
@@ -651,34 +651,37 @@ void mcreq_set_cid(mc_PIPELINE *pipeline, mc_PACKET *packet, uint32_t cid)
     char *key = header_and_key + header_size;
 
     // parse old collection id and determine its length
-    uint32_t old_collection_id;
-    int old_collection_id_length = leb128_decode((uint8_t *)key, key_length, &old_collection_id);
+    uint32_t old_collection_id = 0;
+    int old_collection_id_length = 0;
+    if ((packet->flags & MCREQ_F_HASCID) != 0) {
+        old_collection_id_length = leb128_decode((uint8_t *)key, key_length, &old_collection_id);
+    }
 
     // encode new collection id
     uint8_t collection_id[5] = {0};
     int collection_id_length = leb128_encode(cid, collection_id);
 
     // fix field lengths in the packet
     int diff = collection_id_length - old_collection_id_length;
-    size_t new_header_and_key_size = old_span.size + diff;
     req.request.bodylen = htonl(ntohl(req.request.bodylen) + diff);
-    size_t new_klen = key_length + diff;
+    size_t new_key_length = key_length + diff;
     if (req.request.magic == PROTOCOL_BINARY_AREQ) {
-        req.request.keylen = (new_klen << 8) | (flexible_extras_length & 0xff);
+        req.request.keylen = (new_key_length << 8) | (flexible_extras_length & 0xff);
     } else {
-        req.request.keylen = htons(new_klen);
+        req.request.keylen = htons(new_key_length);
     }
 
     // copy old header fields, with only collection id updated
-    netbuf_mblock_reserve(&pipeline->nbmgr, &packet->kh_span);
-    char *new_header_and_key = SPAN_BUFFER(&packet->kh_span);
+    char *new_header_and_key = malloc(old_span.size + diff);
+    CREATE_STANDALONE_SPAN(&packet->kh_span, new_header_and_key, old_span.size + diff);
+
     const char *ptr = header_and_key;
     memcpy(new_header_and_key, ptr, header_size);
+    // update header with new length values
     memcpy(new_header_and_key, req.bytes, sizeof(req.bytes));
-    ptr += header_size + old_collection_id_length;
     memcpy(new_header_and_key + header_size, collection_id, collection_id_length);
-    memcpy(new_header_and_key + header_size + collection_id_length, ptr,
-           new_header_and_key_size - collection_id_length - header_size);
+    ptr += header_size + old_collection_id_length;
+    memcpy(new_header_and_key + header_size + collection_id_length, ptr, key_length - old_collection_id_length);
 
     // deallocate the old span
     if (IS_STANDALONE_SPAN(&old_span)) {
@@ -689,6 +692,7 @@ void mcreq_set_cid(mc_PIPELINE *pipeline, mc_PACKET *packet, uint32_t cid)
     }
 
     packet->flags |= MCREQ_F_HASCID;
+    packet->flags &= ~MCREQ_F_KEY_NOCOPY;
 }
 
 uint32_t mcreq_get_cid(lcb_INSTANCE *instance, const mc_PACKET *packet, int *cid_set)
@@ -725,7 +729,7 @@ uint32_t mcreq_get_cid(lcb_INSTANCE *instance, const mc_PACKET *packet, int *cid
     return 0;
 }
 
-void mcreq_get_key(lcb_INSTANCE *instance, const mc_PACKET *packet, const char **key, size_t *nkey)
+void mcreq_get_key(const mc_PACKET *packet, const char **key, size_t *nkey)
 {
     uint8_t ffext = 0;
     uint16_t nk = 0;
@@ -743,7 +747,7 @@ void mcreq_get_key(lcb_INSTANCE *instance, const mc_PACKET *packet, const char *
         nk = ntohs(req.request.keylen);
     }
     k = kh + sizeof(req) + req.request.extlen + ffext;
-    if ((packet->flags & MCREQ_F_NOCID) == 0 && instance && LCBT_SETTING(instance, use_collections)) {
+    if ((packet->flags & MCREQ_F_HASCID) != 0) {
         ncid = leb128_decode((uint8_t *)k, nk, &cid);
         (void)cid;
     }

src/mc/mcreq.h

@@ -688,7 +688,7 @@ lcb_STATUS mcreq_basic_packet(mc_CMDQUEUE *queue, const lcb_KEYBUF *key, uint32_
  * @param[out] key
  * @param[out] nkey
  */
-void mcreq_get_key(lcb_INSTANCE *instance, const mc_PACKET *packet, const char **key, size_t *nkey);
+void mcreq_get_key(const mc_PACKET *packet, const char **key, size_t *nkey);
 
 /** @brief Returns the size of the key, in bytes */
 uint16_t mcreq_get_key_size(protocol_binary_request_header *hdr);

src/mcserver/mcserver.cc

@@ -290,7 +290,7 @@ bool Server::handle_unknown_collection(MemcachedResponse &resp, mc_PACKET *oldpk
     std::string name = instance->collcache->id_to_name(cid);
 
     packet_wrapper wrapper;
-    mcreq_get_key(instance, oldpkt, (const char **)&wrapper.key.contig.bytes, &wrapper.key.contig.nbytes);
+    mcreq_get_key(oldpkt, (const char **)&wrapper.key.contig.bytes, &wrapper.key.contig.nbytes);
 
     lcb_log(LOGARGS_T(WARN), LOGFMT "UNKNOWN_COLLECTION. Packet=%p (M=0x%x, S=%u, OP=0x%x), CID=%u, CNAME=%s",
             LOGID_T(), (void *)oldpkt, (int)req.request.magic, oldpkt->opaque, (int)req.request.opcode, (unsigned)cid,

src/newconfig.cc

@@ -214,7 +214,7 @@ static int iterwipe_cb(mc_CMDQUEUE *cq, mc_PIPELINE *oldpl, mc_PACKET *oldpkt, v
 
         /* XXX: We ignore hashkey. This is going away soon, and is probably
          * better than simply failing the items. */
-        mcreq_get_key(instance, oldpkt, &key, &nkey);
+        mcreq_get_key(oldpkt, &key, &nkey);
         lcbvb_map_key(cq->config, key, nkey, &tmpid, &newix);
     }
 

tests/mc/t_alloc.cc

@@ -144,7 +144,7 @@ TEST_F(McAlloc, testKeyAlloc)
     const char *key;
     size_t nkey;
     // Get back the key we just placed inside the header
-    mcreq_get_key(nullptr, packet, &key, &nkey);
+    mcreq_get_key(packet, &key, &nkey);
     ASSERT_EQ(5, nkey);
     ASSERT_EQ(0, memcmp(key, "Hello", 5));
 

tests/mc/t_forward.cc

@@ -250,7 +250,7 @@ TEST_F(McFwd, testNoMap)
     // Get the key
     const char *key;
     size_t nkey;
-    mcreq_get_key(NULL, pkt_tmp, &key, &nkey);
+    mcreq_get_key(pkt_tmp, &key, &nkey);
     ASSERT_EQ(0, nkey);
 
     // Ensure we have no vBucket stamping