From 7a8a22c04a9c24ba3909d18f8b8803bd2b110e92 Mon Sep 17 00:00:00 2001
From: Sayan Nandan <nandansayan@outlook.com>
Date: Mon, 6 Sep 2021 19:50:16 -0700
Subject: [PATCH] Prevent usage of lists as a key type

---
 server/src/protocol/responses.rs |  2 ++
 server/src/queryengine/parser.rs |  6 ++++++
 server/src/queryengine/tests.rs  | 25 +++++++++++++++++++++++++
 3 files changed, 33 insertions(+)

diff --git a/server/src/protocol/responses.rs b/server/src/protocol/responses.rs
index 6efbd426..47c3073d 100644
--- a/server/src/protocol/responses.rs
+++ b/server/src/protocol/responses.rs
@@ -99,6 +99,8 @@ pub mod groups {
     pub const UNKNOWN_PROPERTY: &[u8] = "!16\nunknown-property\n".as_bytes();
     /// The keyspace is not empty and hence cannot be removed
     pub const KEYSPACE_NOT_EMPTY: &[u8] = "!18\nkeyspace-not-empty\n".as_bytes();
+    /// Bad type supplied in a DDL query for the key
+    pub const BAD_TYPE_FOR_KEY: &[u8] = "!16\nbad-type-for-key\n".as_bytes();
 }
 
 pub mod full_responses {
diff --git a/server/src/queryengine/parser.rs b/server/src/queryengine/parser.rs
index 1fd1096c..1f7e1369 100644
--- a/server/src/queryengine/parser.rs
+++ b/server/src/queryengine/parser.rs
@@ -135,6 +135,8 @@ pub(super) fn parse_table_args(
         (BINSTR, LIST_STR) => 5,
         (STR, LIST_BINSTR) => 6,
         (STR, LIST_STR) => 7,
+        // KVExt bad keytypes (we can't use lists as keys for obvious reasons)
+        (LIST_STR, _) | (LIST_BINSTR, _) => return Err(responses::groups::BAD_TYPE_FOR_KEY),
         _ => return Err(responses::groups::UNKNOWN_DATA_TYPE),
     };
     Ok((
@@ -225,14 +227,18 @@ pub(super) fn parse_table_args_test(
     let key_ty = key_ty.as_bytes();
     let val_ty = val_ty.as_bytes();
     let model_code: u8 = match (key_ty, val_ty) {
+        // pure KVE
         (BINSTR, BINSTR) => 0,
         (BINSTR, STR) => 1,
         (STR, STR) => 2,
         (STR, BINSTR) => 3,
+        // KVExt listmap
         (BINSTR, LIST_BINSTR) => 4,
         (BINSTR, LIST_STR) => 5,
         (STR, LIST_BINSTR) => 6,
         (STR, LIST_STR) => 7,
+        // KVExt bad keytypes (we can't use lists as keys for obvious reasons)
+        (LIST_STR, _) | (LIST_BINSTR, _) => return Err(responses::groups::BAD_TYPE_FOR_KEY),
         _ => return Err(responses::groups::UNKNOWN_DATA_TYPE),
     };
     Ok((
diff --git a/server/src/queryengine/tests.rs b/server/src/queryengine/tests.rs
index 25028b23..17b51312 100644
--- a/server/src/queryengine/tests.rs
+++ b/server/src/queryengine/tests.rs
@@ -281,6 +281,31 @@ mod parser_ddl_tests {
             responses::groups::TOO_MANY_ARGUMENTS
         );
     }
+
+    #[test]
+    fn test_bad_key_type() {
+        let mut it = bi!("myverycooltbl", "keymap(list<str>, str)");
+        assert_eq!(
+            parse_table_args_test(&mut it).unwrap_err(),
+            responses::groups::BAD_TYPE_FOR_KEY
+        );
+        let mut it = bi!("myverycooltbl", "keymap(list<binstr>, binstr)");
+        assert_eq!(
+            parse_table_args_test(&mut it).unwrap_err(),
+            responses::groups::BAD_TYPE_FOR_KEY
+        );
+        // for consistency checks
+        let mut it = bi!("myverycooltbl", "keymap(list<str>, binstr)");
+        assert_eq!(
+            parse_table_args_test(&mut it).unwrap_err(),
+            responses::groups::BAD_TYPE_FOR_KEY
+        );
+        let mut it = bi!("myverycooltbl", "keymap(list<binstr>, str)");
+        assert_eq!(
+            parse_table_args_test(&mut it).unwrap_err(),
+            responses::groups::BAD_TYPE_FOR_KEY
+        );
+    }
 }
 
 mod entity_parser_tests {