From e9df58512a55e4b18b4630cfae8832fc1a2e3ce9 Mon Sep 17 00:00:00 2001
From: Xiaojun Huang <xiaojun.huang@intel.com>
Date: Tue, 28 Mar 2023 15:13:58 +0800
Subject: [PATCH] Add pk serialization test and fix the std::bad_alloc issue
 (#66)

* pubkey: Fix serialization std::bad_alloc issue.
* unittest: Add pubkey serialization test case.
Signed-off-by: Huang, Xiaojun <xiaojun.huang@intel.com>
---
 ipcl/include/ipcl/pub_key.hpp | 15 ++++++++++-----
 ipcl/pub_key.cpp              |  1 -
 test/CMakeLists.txt           |  1 +
 test/test_serialization.cpp   | 25 +++++++++++++++++++++++++
 4 files changed, 36 insertions(+), 6 deletions(-)
 create mode 100644 test/test_serialization.cpp

diff --git a/ipcl/include/ipcl/pub_key.hpp b/ipcl/include/ipcl/pub_key.hpp
index f399c40..e524ea8 100644
--- a/ipcl/include/ipcl/pub_key.hpp
+++ b/ipcl/include/ipcl/pub_key.hpp
@@ -148,14 +148,19 @@ class PublicKey {
     ar(::cereal::make_nvp("enable_DJN", enable_DJN));
     ar(::cereal::make_nvp("randbits", randbits));
 
-    std::vector<Ipp32u> n_v, hs_v;
-    ar(::cereal::make_nvp("n", n_v));
-    ar(::cereal::make_nvp("hs", hs_v));
+    int bn_len = bits / 32;
+    std::vector<Ipp32u> n_v(bn_len, 0);
+    std::vector<Ipp32u> hs_v(bn_len * 2, 0);
+    BigNumber n(n_v.data(), bn_len);
+    BigNumber hs(hs_v.data(), bn_len * 2);
+
+    ar(::cereal::make_nvp("n", n));
+    ar(::cereal::make_nvp("hs", hs));
 
     if (enable_DJN)
-      create(n_v.data(), bits, hs_v.data(), randbits);
+      create(n, bits, hs, randbits);
     else
-      create(n_v.data(), bits);
+      create(n, bits);
   }
 
   bool m_isInitialized = false;
diff --git a/ipcl/pub_key.cpp b/ipcl/pub_key.cpp
index d5b2c88..09f0d2f 100644
--- a/ipcl/pub_key.cpp
+++ b/ipcl/pub_key.cpp
@@ -152,7 +152,6 @@ void PublicKey::create(const BigNumber& n, int bits, bool enableDJN_) {
   }
   m_testv = false;
   m_isInitialized = true;
-  std::cout << "create complete" << std::endl;
 }
 
 void PublicKey::create(const BigNumber& n, int bits, const BigNumber& hs,
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 4d7550d..c49535b 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -6,6 +6,7 @@ set(IPCL_UNITTEST_SRC
   main.cpp
   test_cryptography.cpp
   test_ops.cpp
+  test_serialization.cpp
 )
 
 add_executable(unittest_ipcl ${IPCL_UNITTEST_SRC})
diff --git a/test/test_serialization.cpp b/test/test_serialization.cpp
new file mode 100644
index 0000000..df4ed83
--- /dev/null
+++ b/test/test_serialization.cpp
@@ -0,0 +1,25 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+#include "gtest/gtest.h"
+#include "ipcl/ipcl.hpp"
+
+TEST(SerialTest, PublicKeyTest) {
+  const int key_bits = 2048;
+  ipcl::KeyPair key = ipcl::generateKeypair(key_bits);
+  ipcl::PublicKey exp_pk = key.pub_key;
+  ipcl::PrivateKey sk = key.priv_key;
+  std::vector<Ipp32u> vec(key_bits / 32, 0);
+  BigNumber bn(vec.data(), key_bits / 32);
+  ipcl::PublicKey ret_pk(bn, key_bits);
+
+  std::ostringstream os;
+  ipcl::serializer::serialize(os, exp_pk);
+  std::istringstream is(os.str());
+  ipcl::serializer::deserialize(is, ret_pk);
+
+  ipcl::PlainText pt(123);
+  ipcl::CipherText ct = ret_pk.encrypt(pt);
+  ipcl::PlainText dt = sk.decrypt(ct);
+  EXPECT_EQ(pt.getElement(0), dt.getElement(0));
+}