Merge pull request #374 from weibocom/dev_kv_fixbuf_20231017

update handshake response
weibocom · Nov 6, 2023 · 5ec1ba7 · 5ec1ba7
2 parents 0cb1ed3 + 9d5da8f
commit 5ec1ba7
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 4 deletions.
diff --git a/ds/src/mem/ring_slice.rs b/ds/src/mem/ring_slice.rs
@@ -68,13 +68,23 @@ impl RingSlice {
             mask,
         }
     }
+    /// 注意：本方法用vec的capacity作为cap，方便ManuallyDrop结构的恢复及正确drop
     #[inline(always)]
     pub fn from_vec(data: &Vec<u8>) -> Self {
         let mut mem: RingSlice = data.as_slice().into();
         // 这里面的cap是真实的cap
         mem.cap = data.capacity() as u32;
         mem
     }
+
+    /// 注意：对于Vec请使用from_vec，本方法直接用slice的长度作为cap，对ManuallyDrop结构无法友好支持
+    #[inline(always)]
+    pub fn from_slice(data: &[u8]) -> Self {
+        let mut mem: RingSlice = data.into();
+        // 这里面的cap是真实的cap
+        mem.cap = data.len() as u32;
+        mem
+    }
     #[inline(always)]
     pub fn slice(&self, offset: usize, len: usize) -> RingSlice {
         self.sub_slice(offset, len)

diff --git a/protocol/src/kv/common/packets/mod.rs b/protocol/src/kv/common/packets/mod.rs
@@ -1787,7 +1787,7 @@ impl HandshakeResponse {
         db_name: Option<RingSlice>,
         auth_plugin: Option<AuthPlugin>,
         mut capabilities: CapabilityFlags,
-        connect_attributes: Option<HashMap<String, String>>,
+        connect_attributes: Option<&HashMap<String, String>>,
     ) -> Self {
         let scramble_buf =
             if capabilities.contains(CapabilityFlags::CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA) {
@@ -1839,8 +1839,8 @@ impl HandshakeResponse {
                     // .map(|(k, v)| (RawBytes::new(k.into_bytes()), RawBytes::new(v.into_bytes())))
                     .map(|(k, v)| {
                         (
-                            RawBytes::new(RingSlice::from_vec(&k.into_bytes())),
-                            RawBytes::new(RingSlice::from_vec(&v.into_bytes())),
+                            RawBytes::new(RingSlice::from_slice(k.as_bytes())),
+                            RawBytes::new(RingSlice::from_slice(v.as_bytes())),
                         )
                     })
                     .collect()

diff --git a/protocol/src/kv/rsppacket.rs b/protocol/src/kv/rsppacket.rs
@@ -383,6 +383,7 @@ impl<'a, S: crate::Stream> ResponsePacket<'a, S> {
         let client = self.client.as_ref().unwrap();
         let user = client.get_user().unwrap_or_default().as_bytes().to_vec();
         let db_name = client.get_db_name().unwrap_or_default().as_bytes().to_vec();
+        let conn_attrs = client.connect_attrs();
 
         let handshake_response = HandshakeResponse::new(
             scramble_buf,
@@ -393,7 +394,7 @@ impl<'a, S: crate::Stream> ResponsePacket<'a, S> {
             Some(RingSlice::from_vec(&db_name)),
             Some(auth_plugin.clone()),
             client.capability_flags,
-            Some(client.connect_attrs()),
+            Some(&conn_attrs),
         );
         log::debug!("+++ kv handshake rsp: {}", handshake_response);
         let mut buf: Vec<u8> = Vec::with_capacity(256);