Network Socket Extensions

def patch_socket():
    """
    Patch the Python socket class with BSON transmission methods.
    
    Adds the following methods to socket objects:
    - recvbytes(bytes_needed, sock_buf=None): Read exact number of bytes
    - recvobj(): Read complete BSON object  
    - sendobj(obj): Send BSON object atomically
    
    Note: Must be called before creating socket objects to enable BSON methods
    """

import socket
import bson

# Patch socket class first
bson.patch_socket()

# Create socket with BSON capabilities
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

# Now socket has BSON methods available
# sock.sendobj({"message": "hello"})
# obj = sock.recvobj()

def recvbytes(self, bytes_needed, sock_buf=None):
    """
    Atomically read exact number of bytes from socket.
    
    This method is added to socket objects by patch_socket().
    
    Parameters:
    - bytes_needed: int, exact number of bytes to read
    - sock_buf: BytesIO buffer to append to (optional, creates new if None)
    
    Returns:
    BytesIO: Buffer containing exactly bytes_needed bytes, or None if socket closed
    
    Raises:
    socket.error: On network errors
    """

import socket
import bson
from io import BytesIO

bson.patch_socket()

# Server side example
server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_sock.bind(('localhost', 8080))
server_sock.listen(1)

conn, addr = server_sock.accept()

# Read exactly 100 bytes
data_buffer = conn.recvbytes(100)
if data_buffer is not None:
    raw_data = data_buffer.getvalue()
    print(f"Received {len(raw_data)} bytes")
else:
    print("Connection closed by client")

# Reuse buffer for multiple reads
buffer = BytesIO()
header = conn.recvbytes(4, buffer)  # Read 4-byte header
if header is not None:
    payload = conn.recvbytes(96, header)  # Read 96 more bytes into same buffer
    if payload is not None:
        total_data = payload.getvalue()  # Now contains 100 bytes total

def recvobj(self):
    """
    Atomically receive a complete BSON object from socket.
    
    This method is added to socket objects by patch_socket().
    Handles BSON message framing automatically by reading length header first.
    
    Returns:
    dict: Deserialized BSON object, or None if socket closed
    
    Raises:
    socket.error: On network errors
    ValueError: If BSON data is malformed
    """

import socket
import bson

bson.patch_socket()

# Server receiving BSON objects
server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_sock.bind(('localhost', 8080))
server_sock.listen(1)

conn, addr = server_sock.accept()

while True:
    # Receive BSON object
    obj = conn.recvobj()
    
    if obj is None:
        print("Client disconnected")
        break
        
    print(f"Received: {obj}")
    
    # Process the received object
    if obj.get('command') == 'quit':
        break

conn.close()

def sendobj(self, obj):
    """
    Atomically send a BSON object over socket.
    
    This method is added to socket objects by patch_socket().
    Automatically serializes object to BSON and sends with proper framing.
    
    Parameters:
    - obj: dict or BSONCoding object to send
    
    Raises:
    socket.error: On network errors
    UnknownSerializerError: If object cannot be serialized
    """

import socket
import bson

bson.patch_socket()

# Client sending BSON objects
client_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
client_sock.connect(('localhost', 8080))

# Send various BSON objects
messages = [
    {"type": "greeting", "message": "Hello server!"},
    {"type": "data", "values": [1, 2, 3, 4, 5]},
    {"type": "user", "name": "Alice", "age": 30},
    {"command": "quit"}
]

for msg in messages:
    client_sock.sendobj(msg)
    print(f"Sent: {msg}")

client_sock.close()

# Server (server.py)
import socket
import threading
import bson

bson.patch_socket()

def handle_client(conn, addr):
    print(f"Connected by {addr}")
    
    while True:
        try:
            # Receive BSON object
            obj = conn.recvobj()
            
            if obj is None:
                print(f"Client {addr} disconnected")
                break
                
            print(f"Received from {addr}: {obj}")
            
            # Echo back with server timestamp
            import time
            response = {
                "echo": obj,
                "server_time": time.time(),
                "status": "received"
            }
            conn.sendobj(response)
            
        except Exception as e:
            print(f"Error handling {addr}: {e}")
            break
    
    conn.close()

# Start server
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind(('localhost', 8080))
server.listen(5)
print("Server listening on port 8080")

while True:
    conn, addr = server.accept()
    thread = threading.Thread(target=handle_client, args=(conn, addr))
    thread.start()

# Client (client.py)  
import socket
import bson
import time

bson.patch_socket()

client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
client.connect(('localhost', 8080))

# Send test messages
test_data = [
    {"message": "Hello", "id": 1},
    {"data": [10, 20, 30], "type": "array"},
    {"user": {"name": "Bob", "role": "admin"}, "nested": True}
]

for data in test_data:
    # Send object
    client.sendobj(data)
    
    # Receive response
    response = client.recvobj()
    if response:
        print(f"Server response: {response}")
    
    time.sleep(1)

client.close()

obj = sock.recvobj()
if obj is None:
    # Connection closed cleanly
    print("Peer disconnected")
else:
    # Process received object
    handle_message(obj)