From ba1d049d95ce7f628222b0ff5bb9ce59cfb423dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alexandre=20P=C3=A9r=C3=A9?= Date: Tue, 16 Apr 2024 11:18:41 +0200 Subject: [PATCH] fix(compiler): increases deserialization limit for values --- .../concretelang/Bindings/Python/CompilerAPIModule.h | 3 +++ .../compiler/lib/Bindings/Python/CompilerAPIModule.cpp | 9 ++++++--- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Bindings/Python/CompilerAPIModule.h b/compilers/concrete-compiler/compiler/include/concretelang/Bindings/Python/CompilerAPIModule.h index c0fee68fe9..7dd65bf955 100644 --- a/compilers/concrete-compiler/compiler/include/concretelang/Bindings/Python/CompilerAPIModule.h +++ b/compilers/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 #include 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 index 5388396ff1..f5d9734756 100644 --- 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(); 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 keySetUnserialize(const std::string &buffer) { auto keysetProto = Message(); 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};