Repository: secretflow/spu (language: python)

Issue:
[Bug][Critical] Mistaking the source/dst copy in the PackedB2A implementation
### Issue Type

Currentness/Accuracy

### Modules Involved

MPC protocol

### Have you reproduced the bug with SPU HEAD?

Yes

### Have you searched existing issues?

Yes

### SPU Version

74c6d54e5bff9e5d35ab5a73dda98cd9ccf0bfc8

### OS Platform and Distribution

MacOS

### Python Version

3.9

### Compiler Version

_No response_

### Current Behavior?

Basically the copying is wrong here.

https://github.com/secretflow/spu/blob/main/libspu/mpc/cheetah/ot/basic_ot_prot.cc#L133

UTs are passed due to the random masking `r = 0` now which is also unsecure. 

### Standalone code to reproduce the issue

```Python
The 1st argument in `SseTranspose` should be the output.
```


### Relevant log output

_No response_
[Bug][Critical] Mistaking the source/dst copy in the PackedB2A implementation
### Issue Type

Currentness/Accuracy

### Modules Involved

MPC protocol

### Have you reproduced the bug with SPU HEAD?

Yes

### Have you searched existing issues?

Yes

### SPU Version

74c6d54e5bff9e5d35ab5a73dda98cd9ccf0bfc8

### OS Platform and Distribution

MacOS

### Python Version

3.9

### Compiler Version

_No response_

### Current Behavior?

Basically the copying is wrong here.

https://github.com/secretflow/spu/blob/main/libspu/mpc/cheetah/ot/basic_ot_prot.cc#L133

UTs are passed due to the random masking `r = 0` now which is also unsecure. 

### Standalone code to reproduce the issue

```Python
The 1st argument in `SseTranspose` should be the output.
```


### Relevant log output

_No response_


Discussion / hints:
Thank Jinjin Zhou for finding my awful mistake :(.
Thank Jinjin Zhou for finding my awful mistake :(.

Fix patch:
diff --git a/libspu/mpc/cheetah/ot/basic_ot_prot.cc b/libspu/mpc/cheetah/ot/basic_ot_prot.cc
index f5a03549..34d05a1c 100644
--- a/libspu/mpc/cheetah/ot/basic_ot_prot.cc
+++ b/libspu/mpc/cheetah/ot/basic_ot_prot.cc
@@ -68,122 +68,89 @@ NdArrayRef BasicOTProtocols::B2A(const NdArrayRef &inp) {
   return PackedB2A(inp);
 }
 
+// Convert the packed boolean shares to arithmetic share
+// Input x in Z2k is the packed of b-bits for 1 <= b <= k.
+//       That is x0, x1, ..., x{b-1}
+// Output y in Z2k such that y = \sum_i x{i}*2^i mod 2^k
+//
+// Ref: The ABY paper https://encrypto.de/papers/DSZ15.pdf Section E
 NdArrayRef BasicOTProtocols::PackedB2A(const NdArrayRef &inp) {
   const auto *share_t = inp.eltype().as<BShrTy>();
   auto field = inp.eltype().as<Ring2k>()->field();
-  const int64_t n = inp.numel();
-  size_t nbits = share_t->nbits() == 0 ? 1 : share_t->nbits();
-  if (n >= 8) {
-    // 8bits-align for a larger input
-    nbits = (nbits + 7) / 8 * 8;
-  }
-  SPU_ENFORCE(nbits > 0 && nbits <= 8 * SizeOf(field));
-
-  auto rand_bits = DISPATCH_ALL_FIELDS(field, [&]() {
-    if ((nbits & 7) or (n * inp.elsize()) & 7) {
-      //  The SseTranspose requires the #rows and #columns is multiple of 8.
-      //  Thus, we call the less efficient RandBits on margin cases.
-      return RandBits(field, {static_cast<int64_t>(n * nbits)});
-    }
+  const int64_t ring_width = SizeOf(field) * 8;
 
-    // More efficient randbits that ultilize collapse COTs.
-    int64_t B = nbits;
-    auto r = ring_randbit(field, {n * B}).as(makeType<BShrTy>(field, 1));
-    const int64_t numl = r.numel();
+  const int64_t n = inp.numel();
+  const int64_t nbits = share_t->nbits() == 0 ? 1 : share_t->nbits();
+  const int64_t numel = n * nbits;
 
-    NdArrayRef oup = ring_zeros(field, r.shape());
+  NdArrayRef cot_oup = ring_zeros(field, {numel});
+  NdArrayRef arith_oup = ring_zeros(field, inp.shape());
+  DISPATCH_ALL_FIELDS(field, [&]() {
     using u2k = std::make_unsigned<ring2k_t>::type;
-    auto input = NdArrayView<const u2k>(r);
-    auto output = absl::MakeSpan(&oup.at<u2k>(0), numl);
-    SPU_ENFORCE(oup.isCompact());
+    auto input = NdArrayView<const u2k>(inp);
+    auto cot_output = absl::MakeSpan(&cot_oup.at<u2k>(0), cot_oup.numel());
 
     if (Rank() == 0) {
-      std::vector<u2k> corr_data(numl);
-      // NOTE(lwj): Masking to make sure there is only single bit.
-      for (int64_t i = 0; i < numl; ++i) {
-        // corr=-2*xi
-        corr_data[i] = -((input[i] & 1) << 1);
+      std::vector<u2k> corr_data(numel);
+
+      for (int64_t k = 0; k < nbits; ++k) {
+        int64_t i = k * n;
+        auto msk = makeBitsMask<u2k>(ring_width - k);
+        for (int64_t j = 0; j < n; ++j) {
+          // corr[k] = -2*x0_k
+          corr_data[i + j] = -2 * ((input[j] >> k) & 1);
+          corr_data[i + j] &= msk;
+        }
       }
       // Run the multiple COT in the collapse mode.
-      // That is, the i-th COT returns output of `nbits - i` bits.
-      ferret_sender_->SendCAMCC_Collapse(absl::MakeSpan(corr_data), output,
-                                         /*bw*/ nbits, /*num_level*/ nbits);
-      ferret_sender_->Flush();
+      // That is, the k-th COT returns output of `ring_width - k` bits.
+      //
+      // The k-th COT gives the arithmetic share of the k-th bit of the input
+      // according to x_0 ^ x_1 = x_0 + x_1 - 2 * x_0 * x_1
+      ferret_sender_->SendCAMCC_Collapse(absl::MakeSpan(corr_data), cot_output,
+                                         /*bw*/ ring_width,
+                                         /*num_level*/ nbits);
 
-      for (int64_t i = 0; i < numl; ++i) {
-        output[i] = (input[i] & 1) - output[i];
+      ferret_sender_->Flush();
+      for (int64_t k = 0; k < nbits; ++k) {
+        int64_t i = k * n;
+        for (int64_t j = 0; j < n; ++j) {
+          cot_output[i + j] = ((input[j] >> k) & 1) - cot_output[i + j];
+        }
       }
     } else {
-      std::vector<uint8_t> choices(numl);
-      for (int64_t i = 0; i < numl; ++i) {
-        choices[i] = static_cast<uint8_t>(input[i] & 1);
-      }
-      ferret_receiver_->RecvCAMCC_Collapse(absl::MakeSpan(choices), output,
-                                           nbits, nbits);
-
-      for (int64_t i = 0; i < numl; ++i) {
-        output[i] = (input[i] & 1) + output[i];
+      // choice[k] is the k-th bit x1_k
+      std::vector<uint8_t> choices(numel);
+      for (int64_t k = 0; k < nbits; ++k) {
+        int64_t i = k * n;
+        for (int64_t j = 0; j < n; ++j) {
+          choices[i + j] = (input[j] >> k) & 1;
+        }
       }
-    }
-
-    // oup.shape B x (n * T)
-    std::vector<uint8_t> tmp(B * n * inp.elsize());
 
-    // bit matrix transpose
-    SseTranspose(oup.data<uint8_t>(), tmp.data(), B, n * inp.elsize());
+      ferret_receiver_->RecvCAMCC_Collapse(absl::MakeSpan(choices), cot_output,
+                                           ring_width, nbits);
 
-    std::copy_n(tmp.data(), tmp.size(), oup.data<uint8_t>());
-    return oup;
-  });
-
-  // convert the bit form to integer form
-  auto rand = [&](NdArrayRef _bits) {
-    SPU_ENFORCE(_bits.isCompact(), "need compact input");
-    const int64_t n = _bits.numel() / nbits;
-    // init as all 0s.
-    auto iform = ring_zeros(field, inp.shape());
-    DISPATCH_ALL_FIELDS(field, [&]() {
-      auto bits = NdArrayView<const ring2k_t>(_bits);
-      auto digit = NdArrayView<ring2k_t>(iform);
-      for (int64_t i = 0; i < n; ++i) {
-        // LSB is bits[0]; MSB is bits[nbits - 1]
-        // We iterate the bits in reversed order
-        const size_t offset = i * nbits;
-        digit[i] = 0;
-        for (size_t j = nbits; j > 0; --j) {
-          digit[i] = (digit[i] << 1) | (bits[offset + j - 1] & 1);
+      for (int64_t k = 0; k < nbits; ++k) {
+        int64_t i = k * n;
+        for (int64_t j = 0; j < n; ++j) {
+          cot_output[i + j] = ((input[j] >> k) & 1) + cot_output[i + j];
         }
       }
-    });
-    return iform;
-  }(rand_bits);
-
-  // open c = x ^ r
-  auto opened = OpenShare(ring_xor(inp, rand), ReduceOp::XOR, nbits, conn_);
+    }
 
-  // compute c + (1 - 2*c)*<r>
-  NdArrayRef oup = ring_zeros(field, inp.shape());
-  DISPATCH_ALL_FIELDS(field, [&]() {
-    using u2k = std::make_unsigned<ring2k_t>::type;
-    int rank = Rank();
-    auto xr = NdArrayView<const u2k>(rand_bits);
-    auto xc = NdArrayView<const u2k>(opened);
-    auto xo = NdArrayView<ring2k_t>(oup);
-
-    for (int64_t i = 0; i < n; ++i) {
-      const size_t offset = i * nbits;
-      u2k this_elt = xc[i];
-      for (size_t j = 0; j < nbits; ++j, this_elt >>= 1) {
-        u2k c_ij = this_elt & 1;
-        ring2k_t one_bit = (1 - c_ij * 2) * xr[offset + j];
-        if (rank == 0) {
-          one_bit += c_ij;
-        }
-        xo[i] += (one_bit << j);
+    // <x> = \sum_k 2^k * <x_k>
+    // where <x_k> is the arithmetic share of the k-th bit
+    NdArrayView<u2k> arith(arith_oup);
+    for (int64_t k = 0; k < nbits; ++k) {
+      int64_t i = k * n;
+      for (int64_t j = 0; j < n; ++j) {
+        arith[j] += (cot_output[i + j] << k);
       }
     }
   });
-  return oup;
+
+  return arith_oup;
 }
 
 // Math:


Test patch:
diff --git a/libspu/mpc/cheetah/ot/emp/ferret_test.cc b/libspu/mpc/cheetah/ot/emp/ferret_test.cc
index b878bdbe..35353964 100644
--- a/libspu/mpc/cheetah/ot/emp/ferret_test.cc
+++ b/libspu/mpc/cheetah/ot/emp/ferret_test.cc
@@ -205,4 +205,68 @@ TEST_P(FerretCOTTest, ChosenMsgChosenChoice) {
   });
 }
 
+template <typename T>
+T makeMask(int bw) {
+  if (bw == sizeof(T) * 8) {
+    return static_cast<T>(-1);
+  }
+  return (static_cast<T>(1) << bw) - 1;
+}
+
+TEST_P(FerretCOTTest, COT_Collapse) {
+  size_t kWorldSize = 2;
+  int64_t n = 8;
+  auto field = GetParam();
+
+  const auto bw = SizeOf(field) * 8;
+  const int level = bw;
+
+  // generate random choices and correlation
+  const auto _correlation = ring_rand(field, {static_cast<int64_t>(n * level)});
+  const auto N = _correlation.numel();
+
+  NdArrayRef oup1 = ring_zeros(field, _correlation.shape());
+  NdArrayRef oup2 = ring_zeros(field, _correlation.shape());
+
+  std::vector<uint8_t> choices(N, 1);
+
+  DISPATCH_ALL_FIELDS(field, [&]() {
+    using u2k = std::make_unsigned<ring2k_t>::type;
+
+    auto out1_span = absl::MakeSpan(&oup1.at<u2k>(0), N);
+    auto out2_span = absl::MakeSpan(&oup2.at<u2k>(0), N);
+
+    NdArrayView<u2k> correlation(_correlation);
+
+    utils::simulate(kWorldSize, [&](std::shared_ptr<yacl::link::Context> ctx) {
+      auto conn = std::make_shared<Communicator>(ctx);
+      int rank = ctx->Rank();
+
+      EmpFerretOt ferret(conn, rank == 0);
+      if (rank == 0) {
+        ferret.SendCAMCC_Collapse(makeConstSpan(correlation), out1_span, bw,
+                                  level);
+        ferret.Flush();
+
+      } else {
+        ferret.RecvCAMCC_Collapse(absl::MakeSpan(choices), out2_span, bw,
+                                  level);
+      }
+    });
+
+    // Sample-major order
+    //      n ||     n     || n         || .... || n
+    // k=level||k=level - 1||k=level - 2|| ....
+    for (int64_t i = 0; i < N; i += n) {
+      const auto cur_bw = bw - (i / n);
+      const auto mask = makeMask<ring2k_t>(cur_bw);
+      for (int64_t j = 0; j < n; ++j) {
+        ring2k_t c = (-out1_span[i + j] + out2_span[i + j]) & mask;
+        ring2k_t e = (choices[i + j] ? correlation[i + j] : 0) & mask;
+
+        ASSERT_EQ(c, e);
+      }
+    }
+  });
+}
 }  // namespace spu::mpc::cheetah::test
diff --git a/libspu/mpc/cheetah/ot/yacl/ferret_test.cc b/libspu/mpc/cheetah/ot/yacl/ferret_test.cc
index a1eea9e7..1d5822ca 100644
--- a/libspu/mpc/cheetah/ot/yacl/ferret_test.cc
+++ b/libspu/mpc/cheetah/ot/yacl/ferret_test.cc
@@ -210,4 +210,69 @@ TEST_P(FerretCOTTest, ChosenMsgChosenChoice) {
   });
 }
 
+template <typename T>
+T makeMask(int bw) {
+  if (bw == sizeof(T) * 8) {
+    return static_cast<T>(-1);
+  }
+  return (static_cast<T>(1) << bw) - 1;
+}
+
+TEST_P(FerretCOTTest, COT_Collapse) {
+  size_t kWorldSize = 2;
+  int64_t n = 8;
+  auto field = std::get<0>(GetParam());
+  auto use_ss = std::get<1>(GetParam());
+
+  const auto bw = SizeOf(field) * 8;
+  const int level = bw;
+
+  // generate random choices and correlation
+  const auto _correlation = ring_rand(field, {static_cast<int64_t>(n * level)});
+  const auto N = _correlation.numel();
+
+  NdArrayRef oup1 = ring_zeros(field, _correlation.shape());
+  NdArrayRef oup2 = ring_zeros(field, _correlation.shape());
+
+  std::vector<uint8_t> choices(N, 1);
+
+  DISPATCH_ALL_FIELDS(field, [&]() {
+    using u2k = std::make_unsigned<ring2k_t>::type;
+
+    auto out1_span = absl::MakeSpan(&oup1.at<u2k>(0), N);
+    auto out2_span = absl::MakeSpan(&oup2.at<u2k>(0), N);
+
+    NdArrayView<u2k> correlation(_correlation);
+
+    utils::simulate(kWorldSize, [&](std::shared_ptr<yacl::link::Context> ctx) {
+      auto conn = std::make_shared<Communicator>(ctx);
+      int rank = ctx->Rank();
+
+      YaclFerretOt ferret(conn, rank == 0, use_ss);
+      if (rank == 0) {
+        ferret.SendCAMCC_Collapse(makeConstSpan(correlation), out1_span, bw,
+                                  level);
+        ferret.Flush();
+
+      } else {
+        ferret.RecvCAMCC_Collapse(absl::MakeSpan(choices), out2_span, bw,
+                                  level);
+      }
+    });
+
+    // Sample-major order
+    //      n ||     n     || n         || .... || n
+    // k=level||k=level - 1||k=level - 2|| ....
+    for (int64_t i = 0; i < N; i += n) {
+      const auto cur_bw = bw - (i / n);
+      const auto mask = makeMask<ring2k_t>(cur_bw);
+      for (int64_t j = 0; j < n; ++j) {
+        ring2k_t c = (-out1_span[i + j] + out2_span[i + j]) & mask;
+        ring2k_t e = (choices[i + j] ? correlation[i + j] : 0) & mask;
+
+        ASSERT_EQ(c, e);
+      }
+    }
+  });
+}
 }  // namespace spu::mpc::cheetah::test