Fix potential overflow in Intp type (facebookresearch#197)

Summary:
Pull Request resolved: facebookresearch#197

We use int32_t to store 32 bit signed integers. This might cause undesired behavior. To overcome this issue, we add some extra special treatment for these edge cases.

Reviewed By: chualynn

Differential Revision: D36106694

fbshipit-source-id: 978a021d46382857f59f43b454adeb4d7bdcbb40

fbpcf/mpc_std_lib/util/Intp_impl.h

@@ -7,6 +7,7 @@
 
 #pragma once
 #include <climits>
+#include <cmath>
 #include <cstdint>
 #include <exception>
 #include <type_traits>
@@ -54,7 +55,7 @@ class Intp {
   }
 
   Intp<isSigned, width> operator+(const Intp<isSigned, width>& other) const {
-    return Intp<isSigned, width>(round(v_ + other.v_));
+    return Intp<isSigned, width>(add(v_, other.v_));
   }
 
   Intp<isSigned, width> operator-() const {
@@ -70,7 +71,7 @@ class Intp {
   }
 
   Intp<isSigned, width> operator-(const Intp<isSigned, width>& other) const {
-    return Intp<isSigned, width>(round(v_ - other.v_));
+    return Intp<isSigned, width>(subtract(v_, other.v_));
   }
 
  public:
@@ -114,6 +115,42 @@ class Intp {
     }
   }
 
+  static NativeType add(NativeType a, NativeType b) {
+    if constexpr (
+        isSigned &&
+        ((width == 8) || (width == 16) || (width == 32) || (width == 64))) {
+      // special handling is needed only for signed integer with some special
+      // width (e.g. overflow is possible).
+      if (std::signbit(a) == std::signbit(b)) {
+        // the two numbers have the same sign, overflow is possible.
+        // special treatment to prevent overflow
+        return round(uint64_t(a) + uint64_t(b));
+      } else {
+        return round(a + b);
+      }
+    } else {
+      return round(a + b);
+    }
+  }
+
+  static NativeType subtract(NativeType a, NativeType b) {
+    if constexpr (
+        isSigned &&
+        ((width == 8) || (width == 16) || (width == 32) || (width == 64))) {
+      // special handling is needed only for signed integer with some special
+      // width (e.g. overflow is possible).
+      if (std::signbit(a) != std::signbit(b)) {
+        // the two numbers have different sign, overflow is possible.
+        // special treatment to prevent overflow
+        return round(uint64_t(a) - uint64_t(b));
+      } else {
+        return round(a - b);
+      }
+    } else {
+      return round(a - b);
+    }
+  }
+
   NativeType v_;
 };
 

fbpcf/mpc_std_lib/util/test/IntpTest.cpp

@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <memory>
+#include <random>
+
+#include "fbpcf/mpc_std_lib/util/util.h"
+
+namespace fbpcf::mpc_std_lib::util {
+
+TEST(IntpTypeTest, testAdd) {
+  const int8_t width = 32;
+  int64_t largestSigned = std::numeric_limits<int32_t>().max();
+  int64_t smallestSigned = std::numeric_limits<int32_t>().min();
+  std::random_device rd;
+  std::mt19937_64 e(rd());
+  std::uniform_int_distribution<int32_t> dist(smallestSigned, largestSigned);
+  for (int i = 0; i < 1000; i++) {
+    auto v1 = dist(e);
+    auto v2 = dist(e);
+    int32_t v = Intp<true, width>(v1) + Intp<true, width>(v2);
+    int32_t expectedV = (uint64_t)v1 + (uint64_t)v2;
+    EXPECT_EQ(v, expectedV);
+  }
+}
+
+TEST(IntpTypeTest, testSubtract) {
+  const int8_t width = 32;
+  int64_t largestSigned = std::numeric_limits<int32_t>().max();
+  int64_t smallestSigned = std::numeric_limits<int32_t>().min();
+  std::random_device rd;
+  std::mt19937_64 e(rd());
+  std::uniform_int_distribution<int32_t> dist(smallestSigned, largestSigned);
+  for (int i = 0; i < 1000; i++) {
+    auto v1 = dist(e);
+    auto v2 = dist(e);
+    int32_t v = Intp<true, width>(v1) - Intp<true, width>(v2);
+    int32_t expectedV = (uint64_t)v1 - (uint64_t)v2;
+    EXPECT_EQ(v, expectedV);
+  }
+}
+
+} // namespace fbpcf::mpc_std_lib::util