# [grain] Add checks for parameters passed to grain.experimental.index_shuffle that will throw an exception that can by caught by the Python interpreter

Currently when an invalid `rounds` parameter is passed, an assert fails which crashes the interpreter. If index > max_index, an infinite loop may result.

PiperOrigin-RevId: 838873254

Author: Grain Team

grain/_src/python/experimental/index_shuffle/index_shuffle.cc

@@ -22,19 +22,18 @@ limitations under the License.
 // https://eprint.iacr.org/2013/404
 // and following recommendations in
 // https://nsacyber.github.io/simon-speck/implementations/ImplementationGuide1.1.pdf.
-// However we use a single fixed key size and support arbtitary block sizes.
-// Further we fixed the number of rounds in the Feistel structuro to be always
+// However we use a single fixed key size and support arbitrary block sizes.
+// Further we fixed the number of rounds in the Feistel structure to be always
 // 4. This reduces the computational cost and still gives good shuffle behavior.
 //
-// Warning: Given the modifications descripted above this implementation should
-// not be used for application that require cryptograhic secure RNGs.
+// Warning: Given the modifications description above this implementation should
+// not be used for application that require cryptographic secure RNGs.
 
 #include "grain/_src/python/experimental/index_shuffle/index_shuffle.h"
 
 #include <assert.h>
 
 #include <algorithm>
-#include <array>
 #include <bitset>
 #include <cmath>
 #include <cstdint>
@@ -120,6 +119,8 @@ uint64_t index_shuffle(const uint64_t index, const uint64_t max_index,
   assert(block_size > 0 && block_size % 2 == 0 && block_size <= 64);
   // At least 4 rounds and number of rounds must be even.
   assert(rounds >= 4 && rounds % 2 == 0);
+  // Assert the index is bounded by [0, max_index].
+  assert(index >= 0 && index <= max_index);
 #define HANDLE_BLOCK_SIZE(B) \
   case B:                    \
     return impl::index_shuffle<B>(index, max_index, seed, rounds);

grain/_src/python/experimental/index_shuffle/python/index_shuffle_module.cc

@@ -1,5 +1,7 @@
 #include <pybind11/pybind11.h>
 
+#include <cstdint>
+
 #include "grain/_src/python/experimental/index_shuffle/index_shuffle.h"
 
 namespace py = pybind11;
@@ -9,7 +11,17 @@ PYBIND11_MODULE(index_shuffle_module, m) {
       "Returns the position of `index` in a permutation of [0, ..., "
       "max_index].";
   m.doc() = kDoc;
-  m.def("index_shuffle", &::grain::random::index_shuffle, kDoc,
-        py::arg("index"), py::arg("max_index"), py::arg("seed"),
-        py::arg("rounds"));
+  m.def(
+      "index_shuffle",
+      [](uint64_t index, uint64_t max_index, uint32_t seed, uint32_t rounds) {
+        if (rounds < 4 || rounds % 2 != 0) {
+          throw py::value_error("rounds must be an even integer >= 4");
+        }
+        if (index < 0 || index > max_index) {
+          throw py::value_error("index must be in [0, max_index]");
+        }
+        return grain::random::index_shuffle(index, max_index, seed, rounds);
+      },
+      kDoc, py::arg("index"), py::arg("max_index"), py::arg("seed"),
+      py::arg("rounds"));
 }