Merge pull request #658 from AkMo3/main

fix: add connection state for net stream and gracefully handle failure

Author: seetadev

examples/doc-examples/example_net_stream.py

@@ -0,0 +1,263 @@
+"""
+Enhanced NetStream Example for py-libp2p with State Management
+
+This example demonstrates the new NetStream features including:
+- State tracking and transitions
+- Proper error handling and validation
+- Resource cleanup and event notifications
+- Thread-safe operations with Trio locks
+
+Based on the standard echo demo but enhanced to show NetStream state management.
+"""
+
+import argparse
+import random
+import secrets
+
+import multiaddr
+import trio
+
+from libp2p import (
+    new_host,
+)
+from libp2p.crypto.secp256k1 import (
+    create_new_key_pair,
+)
+from libp2p.custom_types import (
+    TProtocol,
+)
+from libp2p.network.stream.exceptions import (
+    StreamClosed,
+    StreamEOF,
+    StreamReset,
+)
+from libp2p.network.stream.net_stream import (
+    NetStream,
+    StreamState,
+)
+from libp2p.peer.peerinfo import (
+    info_from_p2p_addr,
+)
+
+PROTOCOL_ID = TProtocol("/echo/1.0.0")
+
+
+async def enhanced_echo_handler(stream: NetStream) -> None:
+    """
+    Enhanced echo handler that demonstrates NetStream state management.
+    """
+    print(f"New connection established: {stream}")
+    print(f"Initial stream state: {await stream.state}")
+
+    try:
+        # Verify stream is in expected initial state
+        assert await stream.state == StreamState.OPEN
+        assert await stream.is_readable()
+        assert await stream.is_writable()
+        print("✓ Stream initialized in OPEN state")
+
+        # Read incoming data with proper state checking
+        print("Waiting for client data...")
+
+        while await stream.is_readable():
+            try:
+                # Read data from client
+                data = await stream.read(1024)
+                if not data:
+                    print("Received empty data, client may have closed")
+                    break
+
+                print(f"Received: {data.decode('utf-8').strip()}")
+
+                # Check if we can still write before echoing
+                if await stream.is_writable():
+                    await stream.write(data)
+                    print(f"Echoed: {data.decode('utf-8').strip()}")
+                else:
+                    print("Cannot echo - stream not writable")
+                    break
+
+            except StreamEOF:
+                print("Client closed their write side (EOF)")
+                break
+            except StreamReset:
+                print("Stream was reset by client")
+                return
+            except StreamClosed as e:
+                print(f"Stream operation failed: {e}")
+                break
+
+        # Demonstrate graceful closure
+        current_state = await stream.state
+        print(f"Current state before close: {current_state}")
+
+        if current_state not in [StreamState.CLOSE_BOTH, StreamState.RESET]:
+            await stream.close()
+            print("Server closed write side")
+
+        final_state = await stream.state
+        print(f"Final stream state: {final_state}")
+
+    except Exception as e:
+        print(f"Handler error: {e}")
+        # Reset stream on unexpected errors
+        if await stream.state not in [StreamState.RESET, StreamState.CLOSE_BOTH]:
+            await stream.reset()
+            print("Stream reset due to error")
+
+
+async def enhanced_client_demo(stream: NetStream) -> None:
+    """
+    Enhanced client that demonstrates various NetStream state scenarios.
+    """
+    print(f"Client stream established: {stream}")
+    print(f"Initial state: {await stream.state}")
+
+    try:
+        # Verify initial state
+        assert await stream.state == StreamState.OPEN
+        print("✓ Client stream in OPEN state")
+
+        # Scenario 1: Normal communication
+        message = b"Hello from enhanced NetStream client!\n"
+
+        if await stream.is_writable():
+            await stream.write(message)
+            print(f"Sent: {message.decode('utf-8').strip()}")
+        else:
+            print("Cannot write - stream not writable")
+            return
+
+        # Close write side to signal EOF to server
+        await stream.close()
+        print("Client closed write side")
+
+        # Verify state transition
+        state_after_close = await stream.state
+        print(f"State after close: {state_after_close}")
+        assert state_after_close == StreamState.CLOSE_WRITE
+        assert await stream.is_readable()  # Should still be readable
+        assert not await stream.is_writable()  # Should not be writable
+
+        # Try to write (should fail)
+        try:
+            await stream.write(b"This should fail")
+            print("ERROR: Write succeeded when it should have failed!")
+        except StreamClosed as e:
+            print(f"✓ Expected error when writing to closed stream: {e}")
+
+        # Read the echo response
+        if await stream.is_readable():
+            try:
+                response = await stream.read()
+                print(f"Received echo: {response.decode('utf-8').strip()}")
+            except StreamEOF:
+                print("Server closed their write side")
+            except StreamReset:
+                print("Stream was reset")
+
+        # Check final state
+        final_state = await stream.state
+        print(f"Final client state: {final_state}")
+
+    except Exception as e:
+        print(f"Client error: {e}")
+        # Reset on error
+        await stream.reset()
+        print("Client reset stream due to error")
+
+
+async def run_enhanced_demo(
+    port: int, destination: str, seed: int | None = None
+) -> None:
+    """
+    Run enhanced echo demo with NetStream state management.
+    """
+    listen_addr = multiaddr.Multiaddr(f"/ip4/0.0.0.0/tcp/{port}")
+
+    # Generate or use provided key
+    if seed:
+        random.seed(seed)
+        secret_number = random.getrandbits(32 * 8)
+        secret = secret_number.to_bytes(length=32, byteorder="big")
+    else:
+        secret = secrets.token_bytes(32)
+
+    host = new_host(key_pair=create_new_key_pair(secret))
+
+    async with host.run(listen_addrs=[listen_addr]):
+        print(f"Host ID: {host.get_id().to_string()}")
+        print("=" * 60)
+
+        if not destination:  # Server mode
+            print("🖥️  ENHANCED ECHO SERVER MODE")
+            print("=" * 60)
+
+            # type: ignore: Stream is type of NetStream
+            host.set_stream_handler(PROTOCOL_ID, enhanced_echo_handler)
+
+            print(
+                "Run client from another console:\n"
+                f"python3 example_net_stream.py "
+                f"-d {host.get_addrs()[0]}\n"
+            )
+            print("Waiting for connections...")
+            print("Press Ctrl+C to stop server")
+            await trio.sleep_forever()
+
+        else:  # Client mode
+            print("📱 ENHANCED ECHO CLIENT MODE")
+            print("=" * 60)
+
+            # Connect to server
+            maddr = multiaddr.Multiaddr(destination)
+            info = info_from_p2p_addr(maddr)
+            await host.connect(info)
+            print(f"Connected to server: {info.peer_id.pretty()}")
+
+            # Create stream and run enhanced demo
+            stream = await host.new_stream(info.peer_id, [PROTOCOL_ID])
+            if isinstance(stream, NetStream):
+                await enhanced_client_demo(stream)
+
+            print("\n" + "=" * 60)
+            print("CLIENT DEMO COMPLETE")
+
+
+def main() -> None:
+    example_maddr = (
+        "/ip4/127.0.0.1/tcp/8000/p2p/QmQn4SwGkDZKkUEpBRBvTmheQycxAHJUNmVEnjA2v1qe8Q"
+    )
+
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("-p", "--port", default=0, type=int, help="source port number")
+    parser.add_argument(
+        "-d",
+        "--destination",
+        type=str,
+        help=f"destination multiaddr string, e.g. {example_maddr}",
+    )
+    parser.add_argument(
+        "-s",
+        "--seed",
+        type=int,
+        help="seed for deterministic peer ID generation",
+    )
+    parser.add_argument(
+        "--demo-states", action="store_true", help="run state transition demo only"
+    )
+
+    args = parser.parse_args()
+
+    try:
+        trio.run(run_enhanced_demo, args.port, args.destination, args.seed)
+    except KeyboardInterrupt:
+        print("\n👋 Demo interrupted by user")
+    except Exception as e:
+        print(f"❌ Demo failed: {e}")
+
+
+if __name__ == "__main__":
+    main()