fix: add connection states for net stream

Other changes:
1. Add operation validation based on states
2. Gracefully handle exceptions and cleanup

Author: AkMo3

libp2p/network/stream/net_stream.py

@@ -1,3 +1,12 @@
+from enum import (
+    Enum,
+)
+from typing import (
+    Optional,
+)
+
+import trio
+
 from libp2p.abc import (
     IMuxedStream,
     INetStream,
@@ -19,18 +28,42 @@
 )
 
 
-# TODO: Handle exceptions from `muxed_stream`
-# TODO: Add stream state
-#   - Reference: https://github.com/libp2p/go-libp2p-swarm/blob/99831444e78c8f23c9335c17d8f7c700ba25ca14/swarm_stream.go  # noqa: E501
+class StreamState(Enum):
+    """NetStream States"""
+
+    OPEN = "open"
+    CLOSE_READ = "close_read"
+    CLOSE_WRITE = "close_write"
+    CLOSE_BOTH = "close_both"
+    RESET = "reset"
+
+
 class NetStream(INetStream):
+    """Class representing NetStream Handler"""
+
     muxed_stream: IMuxedStream
-    protocol_id: TProtocol | None
+    protocol_id: Optional[TProtocol]
+    __stream_state: StreamState
+
+    def __init__(
+        self, muxed_stream: IMuxedStream, nursery: Optional[trio.Nursery] = None
+    ) -> None:
+        super().__init__()
 
-    def __init__(self, muxed_stream: IMuxedStream) -> None:
         self.muxed_stream = muxed_stream
         self.muxed_conn = muxed_stream.muxed_conn
         self.protocol_id = None
 
+        # For background tasks
+        self._nursery = nursery
+
+        # State management
+        self.__stream_state = StreamState.OPEN
+        self._state_lock = trio.Lock()
+
+        # For notification handling
+        self._notify_lock = trio.Lock()
+
     def get_protocol(self) -> TProtocol | None:
         """
         :return: protocol id that stream runs on
@@ -43,42 +76,168 @@ def set_protocol(self, protocol_id: TProtocol) -> None:
         """
         self.protocol_id = protocol_id
 
-    async def read(self, n: int | None = None) -> bytes:
+    @property
+    async def state(self) -> StreamState:
+        """Get current stream state."""
+        async with self._state_lock:
+            return self.__stream_state
+
+    async def read(self, n: Optional[int] = None) -> bytes:
         """
         Read from stream.
 
         :param n: number of bytes to read
         :return: bytes of input
         """
+        async with self._state_lock:
+            if self.__stream_state in [
+                StreamState.CLOSE_READ,
+                StreamState.CLOSE_BOTH,
+            ]:
+                raise StreamClosed("Stream is closed for reading")
+
+            if self.__stream_state == StreamState.RESET:
+                raise StreamReset("Stream is reset, cannot be used to read")
+
         try:
-            return await self.muxed_stream.read(n)
+            data = await self.muxed_stream.read(n)
+            return data
         except MuxedStreamEOF as error:
+            async with self._state_lock:
+                if self.__stream_state == StreamState.CLOSE_WRITE:
+                    self.__stream_state = StreamState.CLOSE_BOTH
+                    await self._remove()
+                elif self.__stream_state == StreamState.OPEN:
+                    self.__stream_state = StreamState.CLOSE_READ
             raise StreamEOF() from error
         except MuxedStreamReset as error:
+            async with self._state_lock:
+                if self.__stream_state in [
+                    StreamState.OPEN,
+                    StreamState.CLOSE_READ,
+                    StreamState.CLOSE_WRITE,
+                ]:
+                    self.__stream_state = StreamState.RESET
+                    await self._remove()
             raise StreamReset() from error
 
     async def write(self, data: bytes) -> None:
         """
         Write to stream.
 
-        :return: number of bytes written
+        :param data: bytes to write
         """
+        async with self._state_lock:
+            if self.__stream_state in [
+                StreamState.CLOSE_WRITE,
+                StreamState.CLOSE_BOTH,
+                StreamState.RESET,
+            ]:
+                raise StreamClosed("Stream is closed for writing")
+
         try:
             await self.muxed_stream.write(data)
         except (MuxedStreamClosed, MuxedStreamError) as error:
+            async with self._state_lock:
+                if self.__stream_state == StreamState.OPEN:
+                    self.__stream_state = StreamState.CLOSE_WRITE
+                elif self.__stream_state == StreamState.CLOSE_READ:
+                    self.__stream_state = StreamState.CLOSE_BOTH
+                    await self._remove()
             raise StreamClosed() from error
 
     async def close(self) -> None:
-        """Close stream."""
+        """Close stream for writing."""
+        async with self._state_lock:
+            if self.__stream_state in [
+                StreamState.CLOSE_BOTH,
+                StreamState.RESET,
+                StreamState.CLOSE_WRITE,
+            ]:
+                return
+
         await self.muxed_stream.close()
 
+        async with self._state_lock:
+            if self.__stream_state == StreamState.CLOSE_READ:
+                self.__stream_state = StreamState.CLOSE_BOTH
+                await self._remove()
+            elif self.__stream_state == StreamState.OPEN:
+                self.__stream_state = StreamState.CLOSE_WRITE
+
     async def reset(self) -> None:
+        """Reset stream, closing both ends."""
+        async with self._state_lock:
+            if self.__stream_state == StreamState.RESET:
+                return
+
         await self.muxed_stream.reset()
 
-    def get_remote_address(self) -> tuple[str, int] | None:
+        async with self._state_lock:
+            if self.__stream_state in [
+                StreamState.OPEN,
+                StreamState.CLOSE_READ,
+                StreamState.CLOSE_WRITE,
+            ]:
+                self.__stream_state = StreamState.RESET
+                await self._remove()
+
+    async def _remove(self) -> None:
+        """
+        Remove stream from connection and notify listeners.
+        This is called when the stream is fully closed or reset.
+        """
+        if hasattr(self.muxed_conn, "remove_stream"):
+            remove_stream = getattr(self.muxed_conn, "remove_stream")
+            await remove_stream(self)
+
+        # Notify in background using Trio nursery if available
+        if self._nursery:
+            self._nursery.start_soon(self._notify_closed)
+        else:
+            await self._notify_closed()
+
+    async def _notify_closed(self) -> None:
+        """
+        Notify all listeners that the stream has been closed.
+        This runs in a separate task to avoid blocking the main flow.
+        """
+        async with self._notify_lock:
+            if hasattr(self.muxed_conn, "swarm"):
+                swarm = getattr(self.muxed_conn, "swarm")
+
+                if hasattr(swarm, "notify_all"):
+                    await swarm.notify_all(
+                        lambda notifiee: notifiee.closed_stream(swarm, self)
+                    )
+
+                if hasattr(swarm, "refs") and hasattr(swarm.refs, "done"):
+                    swarm.refs.done()
+
+    def get_remote_address(self) -> Optional[tuple[str, int]]:
         """Delegate to the underlying muxed stream."""
         return self.muxed_stream.get_remote_address()
 
-    # TODO: `remove`: Called by close and write when the stream is in specific states.
-    #   It notifies `ClosedStream` after `SwarmConn.remove_stream` is called.
-    # Reference: https://github.com/libp2p/go-libp2p-swarm/blob/99831444e78c8f23c9335c17d8f7c700ba25ca14/swarm_stream.go  # noqa: E501
+    def is_closed(self) -> bool:
+        """Check if stream is closed."""
+        return self.__stream_state in [StreamState.CLOSE_BOTH, StreamState.RESET]
+
+    def is_readable(self) -> bool:
+        """Check if stream is readable."""
+        return self.__stream_state not in [
+            StreamState.CLOSE_READ,
+            StreamState.CLOSE_BOTH,
+            StreamState.RESET,
+        ]
+
+    def is_writable(self) -> bool:
+        """Check if stream is writable."""
+        return self.__stream_state not in [
+            StreamState.CLOSE_WRITE,
+            StreamState.CLOSE_BOTH,
+            StreamState.RESET,
+        ]
+
+    def __str__(self) -> str:
+        """String representation of the stream."""
+        return f"<NetStream[{self.__stream_state.value}] protocol={self.protocol_id}>"