fix(compiler): increases deserialization limit for values

zama-ai · Apr 18, 2024 · ba1d049 · ba1d049
1 parent f66deda
commit ba1d049
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/...ilers/concrete-compiler/compiler/include/concretelang/Bindings/Python/CompilerAPIModule.h b/...ilers/concrete-compiler/compiler/include/concretelang/Bindings/Python/CompilerAPIModule.h
@@ -6,12 +6,15 @@
 #ifndef CONCRETELANG_BINDINGS_PYTHON_COMPILER_API_MODULE_H
 #define CONCRETELANG_BINDINGS_PYTHON_COMPILER_API_MODULE_H
 
+#include <capnp/message.h>
 #include <pybind11/pybind11.h>
 
 namespace mlir {
 namespace concretelang {
 namespace python {
 
+inline constexpr capnp::ReaderOptions DESER_OPTIONS = {7000000000, 64};
+
 void populateCompilerAPISubmodule(pybind11::module &m);
 
 } // namespace python

diff --git a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp
@@ -304,7 +304,7 @@ concretelang::clientlib::EvaluationKeys
 evaluationKeysUnserialize(const std::string &buffer) {
   auto serverKeysetProto = Message<concreteprotocol::ServerKeyset>();
   auto maybeError = serverKeysetProto.readBinaryFromString(
-      buffer, capnp::ReaderOptions{7000000000, 64});
+      buffer, mlir::concretelang::python::DESER_OPTIONS);
   if (maybeError.has_failure()) {
     throw std::runtime_error("Failed to deserialize server keyset." +
                              maybeError.as_failure().error().mesg);
@@ -329,7 +329,7 @@ std::unique_ptr<concretelang::clientlib::KeySet>
 keySetUnserialize(const std::string &buffer) {
   auto keysetProto = Message<concreteprotocol::Keyset>();
   auto maybeError = keysetProto.readBinaryFromString(
-      buffer, capnp::ReaderOptions{7000000000, 64});
+      buffer, mlir::concretelang::python::DESER_OPTIONS);
   if (maybeError.has_failure()) {
     throw std::runtime_error("Failed to deserialize keyset." +
                              maybeError.as_failure().error().mesg);
@@ -351,7 +351,10 @@ std::string keySetSerialize(concretelang::clientlib::KeySet &keySet) {
 concretelang::clientlib::SharedScalarOrTensorData
 valueUnserialize(const std::string &buffer) {
   auto inner = TransportValue();
-  if (inner.readBinaryFromString(buffer).has_failure()) {
+  if (inner
+          .readBinaryFromString(buffer,
+                                mlir::concretelang::python::DESER_OPTIONS)
+          .has_failure()) {
     throw std::runtime_error("Failed to deserialize Value");
   }
   return {inner};