flwr-nightly 1.18.0.dev20250416__py3-none-any.whl → 1.18.0.dev20250417__py3-none-any.whl

flwr/common/secure_aggregation/crypto/shamir.py

@@ -15,61 +15,83 @@
 """Shamir's secret sharing."""
 
 
-import pickle
+import os
 from concurrent.futures import ThreadPoolExecutor
-from typing import cast
 
 from Crypto.Protocol.SecretSharing import Shamir
 from Crypto.Util.Padding import pad, unpad
 
 
 def create_shares(secret: bytes, threshold: int, num: int) -> list[bytes]:
-    """Return list of shares (bytes)."""
+    """Return a list of shares (bytes).
+
+    Shares are created from the provided secret using Shamir's secret sharing.
+    """
+    # Shamir's secret sharing requires the secret to be a multiple of 16 bytes
+    # (AES block size). Pad the secret to the next multiple of 16 bytes.
     secret_padded = pad(secret, 16)
-    secret_padded_chunk = [
-        (threshold, num, secret_padded[i : i + 16])
-        for i in range(0, len(secret_padded), 16)
-    ]
-    share_list: list[list[tuple[int, bytes]]] = [[] for _ in range(num)]
+    chunks = [secret_padded[i : i + 16] for i in range(0, len(secret_padded), 16)]
+
+    # The share list should contain shares of the secret, and each share consists of:
+    # <4 bytes of index><share of chunk1><share of chunk2>...<share of chunkN>
+    share_list: list[bytearray] = [bytearray() for _ in range(num)]
 
-    with ThreadPoolExecutor(max_workers=10) as executor:
+    # Create shares for each chunk in parallel
+    max_workers = min(len(chunks), os.cpu_count() or 1)
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
         for chunk_shares in executor.map(
-            lambda arg: _shamir_split(*arg), secret_padded_chunk
+            lambda chunk: _shamir_split(threshold, num, chunk), chunks
         ):
             for idx, share in chunk_shares:
-                # Index in `chunk_shares` starts from 1
-                share_list[idx - 1].append((idx, share))
+                # Initialize the share with the index if it is empty
+                if not share_list[idx - 1]:
+                    share_list[idx - 1] += idx.to_bytes(4, "little", signed=False)
 
-    return [pickle.dumps(shares) for shares in share_list]
+                # Append the share to the bytes
+                share_list[idx - 1] += share
+
+    return [bytes(share) for share in share_list]
 
 
 def _shamir_split(threshold: int, num: int, chunk: bytes) -> list[tuple[int, bytes]]:
+    """Create shares for a chunk using Shamir's secret sharing.
+
+    Each share is a tuple (index, share_bytes), where share_bytes is 16 bytes long.
+    """
     return Shamir.split(threshold, num, chunk, ssss=False)
 
 
-# Reconstructing secret with PyCryptodome
 def combine_shares(share_list: list[bytes]) -> bytes:
-    """Reconstruct secret from shares."""
-    unpickled_share_list: list[list[tuple[int, bytes]]] = [
-        cast(list[tuple[int, bytes]], pickle.loads(share)) for share in share_list
-    ]
+    """Reconstruct the secret from a list of shares."""
+    # Compute the number of chunks
+    # Each share contains 4 bytes of index and 16 bytes of share for each chunk
+    chunk_num = (len(share_list[0]) - 4) >> 4
 
-    chunk_num = len(unpickled_share_list[0])
     secret_padded = bytearray(0)
-    chunk_shares_list: list[list[tuple[int, bytes]]] = []
-    for i in range(chunk_num):
-        chunk_shares: list[tuple[int, bytes]] = []
-        for share in unpickled_share_list:
-            chunk_shares.append(share[i])
-        chunk_shares_list.append(chunk_shares)
-
-    with ThreadPoolExecutor(max_workers=10) as executor:
+    chunk_shares_list: list[list[tuple[int, bytes]]] = [[] for _ in range(chunk_num)]
+
+    # Split shares into chunks
+    for share in share_list:
+        # The first 4 bytes are the index
+        index = int.from_bytes(share[:4], "little", signed=False)
+        for i in range(chunk_num):
+            start = (i << 4) + 4
+            chunk_shares_list[i].append((index, share[start : start + 16]))
+
+    # Combine shares for each chunk in parallel
+    max_workers = min(chunk_num, os.cpu_count() or 1)
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
         for chunk in executor.map(_shamir_combine, chunk_shares_list):
             secret_padded += chunk
 
-    secret = unpad(secret_padded, 16)
-    return bytes(secret)
+    try:
+        secret = unpad(bytes(secret_padded), 16)
+    except ValueError:
+        # If unpadding fails, it means the shares are not valid
+        raise ValueError("Failed to combine shares") from None
+    return secret
 
 
 def _shamir_combine(shares: list[tuple[int, bytes]]) -> bytes:
+    """Reconstruct a chunk from shares using Shamir's secret sharing."""
     return Shamir.combine(shares, ssss=False)

{flwr_nightly-1.18.0.dev20250416.dist-info → flwr_nightly-1.18.0.dev20250417.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flwr-nightly
-Version: 1.18.0.dev20250416
+Version: 1.18.0.dev20250417
 Summary: Flower: A Friendly Federated AI Framework
 Home-page: https://flower.ai
 License: Apache-2.0

{flwr_nightly-1.18.0.dev20250416.dist-info → flwr_nightly-1.18.0.dev20250417.dist-info}/RECORD

@@ -145,7 +145,7 @@ flwr/common/recorddict_compat.py,sha256=Znn1xRGiqLpPPgviVqyb-GPTM-pCK6tpnEmhWSXa
 flwr/common/retry_invoker.py,sha256=T6puUH3nCxdRzQHeanyr-0nTxhRiS1TH07rmef9vuLQ,14482
 flwr/common/secure_aggregation/__init__.py,sha256=MgW6uHGhyFLBAYQqa1Vzs5n2Gc0d4yEw1_NmerFir70,731
 flwr/common/secure_aggregation/crypto/__init__.py,sha256=5E4q4-Fw0CNz4tLah_QHj7m_rDeM4ucHcFlPWB_Na3Q,738
-flwr/common/secure_aggregation/crypto/shamir.py,sha256=U_WzU4ghYdNr6nPPURZPHWi0AWjeZW1-di--1BVy1L4,2779
+flwr/common/secure_aggregation/crypto/shamir.py,sha256=N8pPa5cEksowNoAqfFm5SP3IuxuVi9GGMa3JOtPniQY,3954
 flwr/common/secure_aggregation/crypto/symmetric_encryption.py,sha256=H6Rlpr4i-VAwOMNdyfH33uDpXNsquiK1gKrOixtKqek,5333
 flwr/common/secure_aggregation/ndarrays_arithmetic.py,sha256=TrggOlizlny3V2KS7-3mpDr33En8UmW9oJcxcS_6oh0,3013
 flwr/common/secure_aggregation/quantization.py,sha256=ssFZpiRyj9ltIh0Ai3vGkDqWFO4SoqgoD1mDU9XqMEM,2400
@@ -326,7 +326,7 @@ flwr/superexec/exec_servicer.py,sha256=Z0YYfs6eNPhqn8rY0x_R04XgR2mKFpggt07IH0EhU
 flwr/superexec/exec_user_auth_interceptor.py,sha256=iqygALkOMBUu_s_R9G0mFThZA7HTUzuXCLgxLCefiwI,4440
 flwr/superexec/executor.py,sha256=M5ucqSE53jfRtuCNf59WFLqQvA1Mln4741TySeZE7qQ,3112
 flwr/superexec/simulation.py,sha256=j6YwUvBN7EQ09ID7MYOCVZ70PGbuyBy8f9bXU0EszEM,4088
-flwr_nightly-1.18.0.dev20250416.dist-info/METADATA,sha256=gaWOX_Z91ciYX-c5Q_Ih-BtNhA2rLyIb8ENPCQH90fE,15868
-flwr_nightly-1.18.0.dev20250416.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-flwr_nightly-1.18.0.dev20250416.dist-info/entry_points.txt,sha256=2-1L-GNKhwGw2_7_RoH55vHw2SIHjdAQy3HAVAWl9PY,374
-flwr_nightly-1.18.0.dev20250416.dist-info/RECORD,,
+flwr_nightly-1.18.0.dev20250417.dist-info/METADATA,sha256=OS2TSt_sdwWENvHBSddnb5tfZ6fgU-9rtq7piRLFv44,15868
+flwr_nightly-1.18.0.dev20250417.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+flwr_nightly-1.18.0.dev20250417.dist-info/entry_points.txt,sha256=2-1L-GNKhwGw2_7_RoH55vHw2SIHjdAQy3HAVAWl9PY,374
+flwr_nightly-1.18.0.dev20250417.dist-info/RECORD,,