syft-flwr 0.1.7__py3-none-any.whl → 0.2.1__py3-none-any.whl

syft_flwr/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.7"
+__version__ = "0.2.1"
 
 from syft_flwr.bootstrap import bootstrap
 from syft_flwr.run import run

syft_flwr/consts.py

@@ -0,0 +1,2 @@
+# Environment variable to control encryption
+SYFT_FLWR_ENCRYPTION_ENABLED = "SYFT_FLWR_ENCRYPTION_ENABLED"

syft_flwr/flower_client.py

@@ -1,3 +1,4 @@
+import base64
 import sys
 import traceback
 
@@ -5,88 +6,176 @@ from flwr.client import ClientApp
 from flwr.common import Context
 from flwr.common.constant import ErrorCode, MessageType
 from flwr.common.message import Error, Message
+from flwr.common.record import RecordDict
 from loguru import logger
 from syft_event import SyftEvents
 from syft_event.types import Request
+from typing_extensions import Optional, Union
 
-from syft_flwr.flwr_compatibility import RecordDict, create_flwr_message
 from syft_flwr.serde import bytes_to_flower_message, flower_message_to_bytes
+from syft_flwr.utils import create_flwr_message, setup_client
 
 
-def _handle_normal_message(
-    message: Message, client_app: ClientApp, context: Context
-) -> bytes:
-    # Normal message handling
-    logger.info(f"Receive message with metadata: {message.metadata}")
-    reply_message: Message = client_app(message=message, context=context)
-    res_bytes: bytes = flower_message_to_bytes(reply_message)
-    logger.info(f"Reply message size: {len(res_bytes)/2**20} MB")
-    return res_bytes
-
-
-def _create_error_reply(message: Message, error: Error) -> bytes:
-    """Create and return error reply message in bytes."""
-    error_reply: Message = create_flwr_message(
-        content=RecordDict(),
-        reply_to=message,
-        message_type=message.metadata.message_type,
-        src_node_id=message.metadata.dst_node_id,
-        dst_node_id=message.metadata.src_node_id,
-        group_id=message.metadata.group_id,
-        run_id=message.metadata.run_id,
-        error=error,
-    )
-    error_bytes: bytes = flower_message_to_bytes(error_reply)
-    logger.info(f"Error reply message size: {len(error_bytes)/2**20} MB")
-    return error_bytes
+class MessageHandler:
+    """Handles message processing for Flower client."""
 
+    def __init__(
+        self, client_app: ClientApp, context: Context, encryption_enabled: bool
+    ):
+        self.client_app = client_app
+        self.context = context
+        self.encryption_enabled = encryption_enabled
 
-def syftbox_flwr_client(client_app: ClientApp, context: Context, app_name: str):
-    """Run the Flower ClientApp with SyftBox."""
-    syft_flwr_app_name = f"flwr/{app_name}"
-    box = SyftEvents(syft_flwr_app_name)
-    client_email = box.client.email
-    logger.info(f"Started SyftBox Flower Client on: {client_email}")
-    logger.info(f"syft_flwr app name: {syft_flwr_app_name}")
+    def prepare_reply(self, data: bytes) -> Union[str, bytes]:
+        """Prepare reply data based on encryption setting."""
+        if self.encryption_enabled:
+            logger.info(f"🔒 Preparing ENCRYPTED reply, size: {len(data)/2**20:.2f} MB")
+            return base64.b64encode(data).decode("utf-8")
+        else:
+            logger.info(f"📤 Preparing PLAINTEXT reply, size: {len(data)/2**20:.2f} MB")
+            return data
+
+    def process_message(self, message: Message) -> Union[str, bytes]:
+        """Process normal Flower message and return reply."""
+        logger.info(f"Processing message with metadata: {message.metadata}")
+        reply_message = self.client_app(message=message, context=self.context)
+        reply_bytes = flower_message_to_bytes(reply_message)
+        return self.prepare_reply(reply_bytes)
+
+    def create_error_reply(
+        self, message: Optional[Message], error: Error
+    ) -> Union[str, bytes]:
+        """Create error reply message."""
+        error_reply = create_flwr_message(
+            content=RecordDict(),
+            reply_to=message,
+            message_type=message.metadata.message_type if message else MessageType.TASK,
+            dst_node_id=message.metadata.src_node_id if message else 0,
+            group_id=message.metadata.group_id if message else "",
+            error=error,
+        )
+        error_bytes = flower_message_to_bytes(error_reply)
+        logger.info(f"Error reply size: {len(error_bytes)/2**20:.2f} MB")
+        return self.prepare_reply(error_bytes)
+
+
+class RequestProcessor:
+    """Processes incoming requests and handles encryption/decryption."""
+
+    def __init__(
+        self, message_handler: MessageHandler, box: SyftEvents, client_email: str
+    ):
+        self.message_handler = message_handler
+        self.box = box
+        self.client_email = client_email
+
+    def decode_request_body(self, request_body: Union[bytes, str]) -> bytes:
+        """Decode request body, handling base64 if encrypted."""
+        if not self.message_handler.encryption_enabled:
+            return request_body
+
+        try:
+            # Convert to string if bytes
+            if isinstance(request_body, bytes):
+                request_body_str = request_body.decode("utf-8")
+            else:
+                request_body_str = request_body
+            # Decode base64
+            decoded = base64.b64decode(request_body_str)
+            logger.debug("🔓 Decoded base64 message")
+            return decoded
+        except Exception:
+            # Not base64 or decoding failed, use as-is
+            return request_body
+
+    def is_stop_signal(self, message: Message) -> bool:
+        """Check if message is a stop signal."""
+        if message.metadata.message_type != MessageType.SYSTEM:
+            return False
+
+        # Check for stop action in config
+        if "config" in message.content and "action" in message.content["config"]:
+            return message.content["config"]["action"] == "stop"
+
+        # Alternative stop signal format
+        return message.metadata.group_id == "final"
+
+    def process(self, request: Request) -> Optional[Union[str, bytes]]:
+        """Process incoming request and return response."""
+        original_sender = request.headers.get("X-Syft-Original-Sender", "unknown")
+        encryption_status = (
+            "🔐 ENCRYPTED"
+            if self.message_handler.encryption_enabled
+            else "📥 PLAINTEXT"
+        )
 
-    @box.on_request("/messages")
-    def handle_messages(request: Request) -> None:
         logger.info(
-            f"Received request id: {request.id}, size: {len(request.body) / 1024 / 1024} (MB)"
+            f"{encryption_status} Received request from {original_sender}, "
+            f"id: {request.id}, size: {len(request.body) / 1024 / 1024:.2f} MB"
         )
-        message: Message = bytes_to_flower_message(request.body)
+
+        # Parse message
         try:
-            # Handle stop signal
-            if message.metadata.message_type == MessageType.SYSTEM:
-                # Check for stop action in various possible formats
-                is_stop_signal = False
-                if (
-                    "config" in message.content
-                    and "action" in message.content["config"]
-                ):
-                    is_stop_signal = message.content["config"]["action"] == "stop"
-                elif message.metadata.group_id == "final":
-                    # Alternative stop signal format
-                    is_stop_signal = True
-
-                if is_stop_signal:
-                    logger.info(f"Received stop message: {message}")
-                    box._stop_event.set()
-                    return None
-
-            return _handle_normal_message(message, client_app, context)
+            request_body = self.decode_request_body(request.body)
+            message = bytes_to_flower_message(request_body)
+
+            if self.message_handler.encryption_enabled:
+                logger.debug(
+                    f"🔓 Successfully decrypted message from {original_sender}"
+                )
+        except Exception as e:
+            logger.error(
+                f"❌ Failed to deserialize message from {original_sender}: {e}"
+            )
+            error = Error(
+                code=ErrorCode.CLIENT_APP_RAISED_EXCEPTION,
+                reason=f"Message deserialization failed: {e}",
+            )
+            return self.message_handler.create_error_reply(None, error)
+
+        # Handle message
+        try:
+            # Check for stop signal
+            if self.is_stop_signal(message):
+                logger.info("Received stop signal")
+                self.box._stop_event.set()
+                return None
+
+            # Process normal message
+            return self.message_handler.process_message(message)
 
         except Exception as e:
-            error_traceback = traceback.format_exc()
-            error_message = f"Client: '{client_email}'. Error: {str(e)}. Traceback: {error_traceback}"
+            error_message = f"Client: '{self.client_email}'. Error: {str(e)}. Traceback: {traceback.format_exc()}"
             logger.error(error_message)
 
             error = Error(
                 code=ErrorCode.CLIENT_APP_RAISED_EXCEPTION, reason=error_message
             )
-            box._stop_event.set()
-            return _create_error_reply(message, error)
+            self.box._stop_event.set()
+            return self.message_handler.create_error_reply(message, error)
+
+
+def syftbox_flwr_client(client_app: ClientApp, context: Context, app_name: str):
+    """Run the Flower ClientApp with SyftBox."""
+    # Setup
+    client, encryption_enabled, syft_flwr_app_name = setup_client(app_name)
+    box = SyftEvents(app_name=syft_flwr_app_name, client=client)
+
+    logger.info(f"Started SyftBox Flower Client on: {box.client.email}")
+    logger.info(f"syft_flwr app name: {syft_flwr_app_name}")
+
+    # Create handlers
+    message_handler = MessageHandler(client_app, context, encryption_enabled)
+    processor = RequestProcessor(message_handler, box, box.client.email)
+
+    # Register message handler
+    @box.on_request(
+        "/messages", auto_decrypt=encryption_enabled, encrypt_reply=encryption_enabled
+    )
+    def handle_messages(request: Request) -> Optional[Union[str, bytes]]:
+        return processor.process(request)
 
+    # Run
     try:
         box.run_forever()
     except Exception as e:

syft_flwr/flower_server.py

@@ -1,12 +1,13 @@
 import traceback
 from random import randint
 
-from loguru import logger
-
 from flwr.common import Context
 from flwr.server import ServerApp
 from flwr.server.run_serverapp import run as run_server
+from loguru import logger
+
 from syft_flwr.grid import SyftGrid
+from syft_flwr.utils import setup_client
 
 
 def syftbox_flwr_server(
@@ -16,10 +17,17 @@ def syftbox_flwr_server(
     app_name: str,
 ) -> Context:
     """Run the Flower ServerApp with SyftBox."""
-    syft_flwr_app_name = f"flwr/{app_name}"
-    syft_grid = SyftGrid(app_name=syft_flwr_app_name, datasites=datasites)
+    client, _, syft_flwr_app_name = setup_client(app_name)
+
+    # Construct the SyftGrid
+    syft_grid = SyftGrid(
+        app_name=syft_flwr_app_name, datasites=datasites, client=client
+    )
+
+    # Set the run id (random for now)
     run_id = randint(0, 1000)
     syft_grid.set_run(run_id)
+
     logger.info(f"Started SyftBox Flower Server on: {syft_grid._client.email}")
     logger.info(f"syft_flwr app name: {syft_flwr_app_name}")