SERVER-32889 Mongo shell injection function for hashing shard key

4 files changed, 126 insertions, 1 deletions

diff --git a/src/mongo/SConscript b/src/mongo/SConscript
index beef6b2a7ac..6d6b8dfbdb0 100644
--- a/src/mongo/SConscript
+++ b/src/mongo/SConscript
@@ -455,7 +455,6 @@ if not has_option('noshell') and usemozjs:
                     "shell/mk_wcwidth.cpp",
                     "shell/mongo-server.cpp",
                     "shell/shell_options.cpp",
-                    "shell/shell_options_init.cpp",
                     "shell/shell_utils.cpp",
                     "shell/shell_utils_extended.cpp",
                     "shell/shell_utils_launcher.cpp",
@@ -509,6 +508,7 @@ if not has_option('noshell') and usemozjs:
         [
             "shell/dbshell.cpp",
             "shell/mongodbcr.cpp",
+            "shell/shell_options_init.cpp",
         ] + env.WindowsResourceFile("shell/shell.rc"),
         LIBDEPS=[
             "$BUILD_DIR/third_party/shim_pcrecpp",
diff --git a/src/mongo/dbtests/SConscript b/src/mongo/dbtests/SConscript
index 0257196aafb..2bc24a869ad 100644
--- a/src/mongo/dbtests/SConscript
+++ b/src/mongo/dbtests/SConscript
@@ -145,6 +145,7 @@ dbtest = env.Program(
         "$BUILD_DIR/mongo/util/net/network",
         "$BUILD_DIR/mongo/util/progress_meter",
         "$BUILD_DIR/mongo/util/version_impl",
+        "$BUILD_DIR/mongo/shell_core",
         "mocklib",
         "testframework",
     ],
diff --git a/src/mongo/dbtests/jstests.cpp b/src/mongo/dbtests/jstests.cpp
index 04bb4c2ad1b..e928f54b1da 100644
--- a/src/mongo/dbtests/jstests.cpp
+++ b/src/mongo/dbtests/jstests.cpp
@@ -39,10 +39,12 @@
 #include "mongo/base/parse_number.h"
 #include "mongo/db/client.h"
 #include "mongo/db/dbdirectclient.h"
+#include "mongo/db/hasher.h"
 #include "mongo/db/json.h"
 #include "mongo/dbtests/dbtests.h"
 #include "mongo/platform/decimal128.h"
 #include "mongo/scripting/engine.h"
+#include "mongo/shell/shell_utils.h"
 #include "mongo/util/concurrency/thread_name.h"
 #include "mongo/util/log.h"
 #include "mongo/util/timer.h"
@@ -2447,6 +2449,89 @@ public:
     }
 };
 
+class ConvertShardKeyToHashed {
+public:
+    void check(shared_ptr<Scope> s, const mongo::BSONObj& o) {
+        s->setObject("o", o, true);
+        s->invoke("return convertShardKeyToHashed(o);", 0, 0);
+        const auto scopeShardKey = s->getNumber("__returnValue");
+
+        // Wrapping to form a proper element
+        const auto wrapO = BSON("" << o);
+        const auto e = wrapO[""];
+        const auto trueShardKey =
+            mongo::BSONElementHasher::hash64(e, mongo::BSONElementHasher::DEFAULT_HASH_SEED);
+
+        ASSERT_EQUALS(scopeShardKey, trueShardKey);
+    }
+
+    void checkWithSeed(shared_ptr<Scope> s, const mongo::BSONObj& o, int seed) {
+        s->setObject("o", o, true);
+        s->setNumber("seed", seed);
+        s->invoke("return convertShardKeyToHashed(o, seed);", 0, 0);
+        const auto scopeShardKey = s->getNumber("__returnValue");
+
+        // Wrapping to form a proper element
+        const auto wrapO = BSON("" << o);
+        const auto e = wrapO[""];
+        const auto trueShardKey = mongo::BSONElementHasher::hash64(e, seed);
+
+        ASSERT_EQUALS(scopeShardKey, trueShardKey);
+    }
+
+    void checkNoArgs(shared_ptr<Scope> s) {
+        s->invoke("return convertShardKeyToHashed();", 0, 0);
+    }
+
+    void checkWithExtraArg(shared_ptr<Scope> s, const mongo::BSONObj& o, int seed) {
+        s->setObject("o", o, true);
+        s->setNumber("seed", seed);
+        s->invoke("return convertShardKeyToHashed(o, seed, 1);", 0, 0);
+    }
+
+    void checkWithBadSeed(shared_ptr<Scope> s, const mongo::BSONObj& o) {
+        s->setObject("o", o, true);
+        s->setString("seed", "sunflower");
+        s->invoke("return convertShardKeyToHashed(o, seed);", 0, 0);
+    }
+
+    void run() {
+        shared_ptr<Scope> s(getGlobalScriptEngine()->newScope());
+        shell_utils::installShellUtils(*s);
+
+        // Check a few elementary objects
+        check(s, BSON("" << 1));
+        check(s, BSON("" << 10.0));
+        check(s,
+              BSON(""
+                   << "Shardy"));
+        check(s, BSON("" << BSON_ARRAY(1 << 2 << 3)));
+        check(s, BSON("" << mongo::jstNULL));
+        check(s, BSON("" << mongo::BSONObj()));
+        check(s,
+              BSON("A" << 1 << "B"
+                       << "Shardy"));
+
+        // Check a few different seeds
+        checkWithSeed(s,
+                      BSON(""
+                           << "Shardy"),
+                      mongo::BSONElementHasher::DEFAULT_HASH_SEED);
+        checkWithSeed(s,
+                      BSON(""
+                           << "Shardy"),
+                      0);
+        checkWithSeed(s,
+                      BSON(""
+                           << "Shardy"),
+                      -1);
+
+        ASSERT_THROWS(checkNoArgs(s), mongo::DBException);
+        ASSERT_THROWS(checkWithExtraArg(s, BSON("" << 10.0), 0), mongo::DBException);
+        ASSERT_THROWS(checkWithBadSeed(s, BSON("" << 1)), mongo::DBException);
+    }
+};
+
 class All : public Suite {
 public:
     All() : Suite("js") {}
@@ -2505,6 +2590,7 @@ public:
         add<ErrorCodeFromInvoke>();
         add<ErrorWithSidecarFromInvoke>();
         add<RequiresOwnedObjects>();
+        add<ConvertShardKeyToHashed>();
 
         add<RoundTripTests::DBRefTest>();
         add<RoundTripTests::DBPointerTest>();
diff --git a/src/mongo/shell/shell_utils.cpp b/src/mongo/shell/shell_utils.cpp
index c13f6291636..962e4ef2a43 100644
--- a/src/mongo/shell/shell_utils.cpp
+++ b/src/mongo/shell/shell_utils.cpp
@@ -35,6 +35,7 @@
 
 #include "mongo/client/dbclientinterface.h"
 #include "mongo/client/replica_set_monitor.h"
+#include "mongo/db/hasher.h"
 #include "mongo/platform/random.h"
 #include "mongo/scripting/engine.h"
 #include "mongo/shell/bench.h"
@@ -185,6 +186,42 @@ BSONObj computeSHA256Block(const BSONObj& a, void* data) {
     return bob.obj();
 }
 
+/**
+ * This function computes a hash value for a document.
+ * Specifically, this is the same hash function that is used to form a hashed index,
+ * and thus used to generate shard keys for a collection.
+ *
+ * e.g.
+ * > // For a given collection prepared like so:
+ * > use mydb
+ * > db.mycollection.createIndex({ x: "hashed" })
+ * > sh.shardCollection("mydb.mycollection", { x: "hashed" })
+ * > // And a sample object like so:
+ * > var obj = { x: "Whatever key", y: 2, z: 10.0 }
+ * > // The hashed value of the shard key can be acquired from the shard key-value pair like so:
+ * > convertShardKeyToHashed({x: "Whatever key"})
+ */
+BSONObj convertShardKeyToHashed(const BSONObj& a, void* data) {
+    const auto& objEl = a[0];
+
+    uassert(10151,
+            "convertShardKeyToHashed accepts either 1 or 2 arguments",
+            a.nFields() >= 1 && a.nFields() <= 2);
+
+    // It looks like the seed is always default right now.
+    // But no reason not to allow for the future
+    auto seed = BSONElementHasher::DEFAULT_HASH_SEED;
+    if (a.nFields() > 1) {
+        auto seedEl = a[1];
+
+        uassert(10159, "convertShardKeyToHashed seed value should be a number", seedEl.isNumber());
+        seed = seedEl.numberInt();
+    }
+
+    auto key = BSONElementHasher::hash64(objEl, seed);
+    return BSON("" << key);
+}
+
 BSONObj replMonitorStats(const BSONObj& a, void* data) {
     uassert(17134,
             "replMonitorStats requires a single string argument (the ReplSet name)",
@@ -238,6 +275,7 @@ void installShellUtils(Scope& scope) {
     scope.injectNative("interpreterVersion", interpreterVersion);
     scope.injectNative("getBuildInfo", getBuildInfo);
     scope.injectNative("computeSHA256Block", computeSHA256Block);
+    scope.injectNative("convertShardKeyToHashed", convertShardKeyToHashed);
     scope.injectNative("fileExists", fileExistsJS);
 
 #ifndef MONGO_SAFE_SHELL