starpc 0.41.2 → 0.43.1

package/srpc/client.rs

@@ -0,0 +1,356 @@
+//! Client implementation for starpc.
+//!
+//! This module provides the client-side API for making RPC calls,
+//! supporting unary, client streaming, server streaming, and
+//! bidirectional streaming patterns.
+
+use async_trait::async_trait;
+use bytes::Bytes;
+use prost::Message;
+use std::sync::Arc;
+
+use crate::error::{Error, Result};
+use crate::rpc::{ClientRpc, PacketWriter};
+use crate::stream::{Context, Stream, StreamExt};
+use crate::transport::create_packet_channel;
+
+/// Receiver for incoming packets.
+pub type PacketReceiver = tokio::sync::mpsc::Receiver<crate::proto::Packet>;
+
+/// Trait for opening streams to a remote server.
+///
+/// Implementations of this trait provide the transport layer for RPC calls.
+/// Each call to `open_stream` should return a new, independent stream.
+#[async_trait]
+pub trait OpenStream: Send + Sync {
+    /// Opens a new stream and returns a packet writer and a receiver for incoming packets.
+    async fn open_stream(&self) -> Result<(Arc<dyn PacketWriter>, PacketReceiver)>;
+}
+
+/// Client trait for making RPC calls.
+///
+/// This trait defines the core client operations matching the Go Client interface.
+#[async_trait]
+pub trait Client: Send + Sync {
+    /// Executes a unary RPC call.
+    ///
+    /// Sends the input message and waits for a single response.
+    async fn exec_call<I, O>(&self, service: &str, method: &str, input: &I) -> Result<O>
+    where
+        I: Message + Send + Sync,
+        O: Message + Default;
+
+    /// Opens a new stream for a streaming RPC.
+    ///
+    /// # Arguments
+    /// * `service` - The service ID
+    /// * `method` - The method ID
+    /// * `first_msg` - Optional initial message data
+    async fn new_stream(
+        &self,
+        service: &str,
+        method: &str,
+        first_msg: Option<&[u8]>,
+    ) -> Result<Box<dyn Stream>>;
+}
+
+/// Boxed Client trait object.
+pub type BoxClient = Box<dyn Client>;
+
+/// Standard starpc client implementation.
+///
+/// This is the primary client implementation that works with any transport
+/// implementing the `OpenStream` trait.
+pub struct SrpcClient<T: OpenStream> {
+    /// The stream opener.
+    opener: T,
+}
+
+impl<T: OpenStream> SrpcClient<T> {
+    /// Creates a new client with the given stream opener.
+    pub fn new(opener: T) -> Self {
+        Self { opener }
+    }
+}
+
+#[async_trait]
+impl<T: OpenStream + 'static> Client for SrpcClient<T> {
+    async fn exec_call<I, O>(&self, service: &str, method: &str, input: &I) -> Result<O>
+    where
+        I: Message + Send + Sync,
+        O: Message + Default,
+    {
+        // Marshal the input.
+        let input_data = input.encode_to_vec();
+
+        // Open a stream.
+        let (writer, mut receiver) = self.opener.open_stream().await?;
+
+        // Create the client RPC.
+        let ctx = Context::new();
+        let rpc = Arc::new(ClientRpc::new(
+            ctx.clone(),
+            service.to_string(),
+            method.to_string(),
+            writer,
+        ));
+
+        // Start the RPC with the input data.
+        rpc.start(Some(Bytes::from(input_data))).await?;
+
+        // Spawn a task to handle incoming packets.
+        let rpc_clone = rpc.clone();
+        let packet_handler = tokio::spawn(async move {
+            while let Some(packet) = receiver.recv().await {
+                if rpc_clone.handle_packet(packet).await.is_err() {
+                    break;
+                }
+            }
+            let _ = rpc_clone.handle_stream_close(None).await;
+        });
+
+        // Send close to indicate we're done sending.
+        rpc.close_send().await?;
+
+        // Receive the response.
+        let output: O = rpc.msg_recv().await?;
+
+        // Wait for the RPC to complete properly (receive any trailing packets).
+        // This ensures we process any completion/error packets from the server
+        // before cleaning up.
+        let _ = rpc.wait().await;
+
+        // Close the RPC to signal completion.
+        let _ = rpc.close().await;
+
+        // Clean up the packet handler.
+        packet_handler.abort();
+
+        Ok(output)
+    }
+
+    async fn new_stream(
+        &self,
+        service: &str,
+        method: &str,
+        first_msg: Option<&[u8]>,
+    ) -> Result<Box<dyn Stream>> {
+        // Open a stream.
+        let (writer, mut receiver) = self.opener.open_stream().await?;
+
+        // Create the client RPC.
+        let ctx = Context::new();
+        let rpc = Arc::new(ClientRpc::new(
+            ctx.clone(),
+            service.to_string(),
+            method.to_string(),
+            writer,
+        ));
+
+        // Start the RPC.
+        let first_data = first_msg.map(|d| Bytes::from(d.to_vec()));
+        rpc.start(first_data).await?;
+
+        // Spawn a task to handle incoming packets.
+        let rpc_clone = rpc.clone();
+        let packet_handler = tokio::spawn(async move {
+            while let Some(packet) = receiver.recv().await {
+                if rpc_clone.handle_packet(packet).await.is_err() {
+                    break;
+                }
+            }
+            let _ = rpc_clone.handle_stream_close(None).await;
+        });
+
+        // Return a stream wrapper that provides the Stream interface.
+        Ok(Box::new(ClientStream {
+            rpc,
+            packet_handler: tokio::sync::Mutex::new(Some(packet_handler)),
+        }))
+    }
+}
+
+/// Stream wrapper for client-side streaming.
+struct ClientStream {
+    rpc: Arc<ClientRpc>,
+    /// Handle to the background packet handler task.
+    /// Aborted when the stream is closed.
+    packet_handler: tokio::sync::Mutex<Option<tokio::task::JoinHandle<()>>>,
+}
+
+#[async_trait]
+impl Stream for ClientStream {
+    fn context(&self) -> &Context {
+        self.rpc.context()
+    }
+
+    async fn send_bytes(&self, data: Bytes) -> Result<()> {
+        self.rpc.send_bytes(data).await
+    }
+
+    async fn recv_bytes(&self) -> Result<Bytes> {
+        self.rpc.recv_bytes().await
+    }
+
+    async fn close_send(&self) -> Result<()> {
+        self.rpc.close_send().await
+    }
+
+    async fn close(&self) -> Result<()> {
+        let _ = self.rpc.close().await;
+        // Abort the background packet handler to ensure cleanup.
+        if let Some(handle) = self.packet_handler.lock().await.take() {
+            handle.abort();
+        }
+        Ok(())
+    }
+}
+
+/// Transport-based stream openers.
+///
+/// This module provides `OpenStream` implementations for common transport types.
+pub mod transport {
+    use super::*;
+    use std::sync::Mutex;
+    use tokio::io::{AsyncRead, AsyncWrite};
+
+    /// A simple stream opener over a single connection.
+    ///
+    /// This opener consumes a single transport connection on the first call
+    /// to `open_stream`. Subsequent calls will fail with `StreamClosed`.
+    ///
+    /// For multiplexed connections (multiple concurrent streams), use yamux
+    /// or similar multiplexing protocols.
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// use tokio::net::TcpStream;
+    /// use starpc::client::transport::SingleStreamOpener;
+    ///
+    /// let stream = TcpStream::connect("127.0.0.1:8080").await?;
+    /// let opener = SingleStreamOpener::new(stream);
+    /// let client = SrpcClient::new(opener);
+    /// ```
+    pub struct SingleStreamOpener<T> {
+        inner: Mutex<Option<T>>,
+    }
+
+    impl<T: AsyncRead + AsyncWrite + Send + Unpin + 'static> SingleStreamOpener<T> {
+        /// Creates a new single stream opener.
+        pub fn new(transport: T) -> Self {
+            Self {
+                inner: Mutex::new(Some(transport)),
+            }
+        }
+    }
+
+    #[async_trait]
+    impl<T: AsyncRead + AsyncWrite + Send + Unpin + 'static> OpenStream for SingleStreamOpener<T> {
+        async fn open_stream(&self) -> Result<(Arc<dyn PacketWriter>, PacketReceiver)> {
+            let transport = self
+                .inner
+                .lock()
+                .unwrap()
+                .take()
+                .ok_or(Error::StreamClosed)?;
+
+            let (read_half, write_half) = tokio::io::split(transport);
+            Ok(create_packet_channel(read_half, write_half))
+        }
+    }
+
+    /// Re-export TransportPacketWriter for direct use.
+    pub use crate::transport::TransportPacketWriter;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::atomic::{AtomicBool, Ordering};
+    use std::sync::Mutex;
+
+    struct MockWriter {
+        closed: AtomicBool,
+    }
+
+    impl MockWriter {
+        fn new() -> Self {
+            Self {
+                closed: AtomicBool::new(false),
+            }
+        }
+    }
+
+    #[async_trait]
+    impl PacketWriter for MockWriter {
+        async fn write_packet(&self, _packet: crate::proto::Packet) -> Result<()> {
+            Ok(())
+        }
+
+        async fn close(&self) -> Result<()> {
+            self.closed.store(true, Ordering::SeqCst);
+            Ok(())
+        }
+    }
+
+    struct MockOpener {
+        writer: Arc<MockWriter>,
+        receiver: Mutex<Option<PacketReceiver>>,
+    }
+
+    impl MockOpener {
+        fn new() -> (Self, tokio::sync::mpsc::Sender<crate::proto::Packet>) {
+            let (tx, rx) = tokio::sync::mpsc::channel(32);
+            (
+                Self {
+                    writer: Arc::new(MockWriter::new()),
+                    receiver: Mutex::new(Some(rx)),
+                },
+                tx,
+            )
+        }
+    }
+
+    #[async_trait]
+    impl OpenStream for MockOpener {
+        async fn open_stream(&self) -> Result<(Arc<dyn PacketWriter>, PacketReceiver)> {
+            let rx = self
+                .receiver
+                .lock()
+                .unwrap()
+                .take()
+                .ok_or(Error::StreamClosed)?;
+            Ok((self.writer.clone(), rx))
+        }
+    }
+
+    #[tokio::test]
+    async fn test_client_new_stream() {
+        let (opener, _tx) = MockOpener::new();
+        let client = SrpcClient::new(opener);
+
+        let stream = client
+            .new_stream("test.Service", "TestMethod", Some(b"hello"))
+            .await
+            .unwrap();
+
+        assert!(!stream.context().is_cancelled());
+    }
+
+    #[tokio::test]
+    async fn test_single_stream_opener_only_once() {
+        use tokio::io::duplex;
+
+        let (client_stream, _server_stream) = duplex(1024);
+        let opener = transport::SingleStreamOpener::new(client_stream);
+
+        // First open should succeed
+        let result1 = opener.open_stream().await;
+        assert!(result1.is_ok());
+
+        // Second open should fail
+        let result2 = opener.open_stream().await;
+        assert!(matches!(result2, Err(Error::StreamClosed)));
+    }
+}

package/srpc/codec.rs

@@ -0,0 +1,225 @@
+//! Length-prefixed packet codec for starpc wire format.
+//!
+//! Wire format: 4-byte little-endian u32 length prefix + protobuf-encoded packet.
+
+use bytes::{Buf, BufMut, BytesMut};
+use prost::Message;
+use tokio_util::codec::{Decoder, Encoder};
+
+use crate::error::{Error, Result};
+use crate::proto::Packet;
+
+/// Maximum message size (10MB, same as Go implementation).
+pub const MAX_MESSAGE_SIZE: usize = 10_000_000;
+
+/// Length of the size prefix in bytes.
+const SIZE_PREFIX_LEN: usize = 4;
+
+/// Codec for encoding and decoding starpc packets with length-prefix framing.
+#[derive(Debug, Default, Clone)]
+pub struct PacketCodec;
+
+impl PacketCodec {
+    /// Creates a new packet codec.
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+impl Decoder for PacketCodec {
+    type Item = Packet;
+    type Error = Error;
+
+    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>> {
+        // Need at least the size prefix.
+        if src.len() < SIZE_PREFIX_LEN {
+            return Ok(None);
+        }
+
+        // Read the length prefix (little-endian u32).
+        let mut size_bytes = [0u8; SIZE_PREFIX_LEN];
+        size_bytes.copy_from_slice(&src[..SIZE_PREFIX_LEN]);
+        let msg_size = u32::from_le_bytes(size_bytes) as usize;
+
+        // Validate message size.
+        if msg_size == 0 {
+            return Err(Error::MessageSizeZero);
+        }
+        if msg_size > MAX_MESSAGE_SIZE {
+            return Err(Error::MessageTooLarge(msg_size, MAX_MESSAGE_SIZE));
+        }
+
+        // Check if we have the complete message.
+        let total_size = SIZE_PREFIX_LEN + msg_size;
+        if src.len() < total_size {
+            // Reserve capacity for the remaining data.
+            src.reserve(total_size - src.len());
+            return Ok(None);
+        }
+
+        // Consume the length prefix.
+        src.advance(SIZE_PREFIX_LEN);
+
+        // Decode the packet.
+        let packet_bytes = src.split_to(msg_size);
+        let packet = Packet::decode(&packet_bytes[..])?;
+
+        Ok(Some(packet))
+    }
+}
+
+impl Encoder<Packet> for PacketCodec {
+    type Error = Error;
+
+    fn encode(&mut self, item: Packet, dst: &mut BytesMut) -> Result<()> {
+        // Calculate the encoded size.
+        let msg_size = item.encoded_len();
+
+        // Validate message size.
+        if msg_size > MAX_MESSAGE_SIZE {
+            return Err(Error::MessageTooLarge(msg_size, MAX_MESSAGE_SIZE));
+        }
+
+        // Reserve space for the length prefix and message.
+        dst.reserve(SIZE_PREFIX_LEN + msg_size);
+
+        // Write the length prefix (little-endian u32).
+        dst.put_u32_le(msg_size as u32);
+
+        // Encode the packet directly into the buffer.
+        item.encode(dst)?;
+
+        Ok(())
+    }
+}
+
+/// Encode a packet to bytes with length prefix.
+pub fn encode_packet(packet: &Packet) -> Result<Vec<u8>> {
+    let msg_size = packet.encoded_len();
+    if msg_size > MAX_MESSAGE_SIZE {
+        return Err(Error::MessageTooLarge(msg_size, MAX_MESSAGE_SIZE));
+    }
+
+    let mut buf = Vec::with_capacity(SIZE_PREFIX_LEN + msg_size);
+    buf.extend_from_slice(&(msg_size as u32).to_le_bytes());
+    packet.encode(&mut buf)?;
+
+    Ok(buf)
+}
+
+/// Decode a packet from bytes (without length prefix).
+pub fn decode_packet(bytes: &[u8]) -> Result<Packet> {
+    Ok(Packet::decode(bytes)?)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::proto::{packet::Body, CallData, CallStart};
+
+    #[test]
+    fn test_codec_roundtrip_call_start() {
+        let mut codec = PacketCodec::new();
+        let mut buf = BytesMut::new();
+
+        let packet = Packet {
+            body: Some(Body::CallStart(CallStart {
+                rpc_service: "test.Service".into(),
+                rpc_method: "TestMethod".into(),
+                data: vec![1, 2, 3],
+                data_is_zero: false,
+            })),
+        };
+
+        codec.encode(packet.clone(), &mut buf).unwrap();
+
+        let decoded = codec.decode(&mut buf).unwrap().unwrap();
+        assert_eq!(decoded, packet);
+    }
+
+    #[test]
+    fn test_codec_roundtrip_call_data() {
+        let mut codec = PacketCodec::new();
+        let mut buf = BytesMut::new();
+
+        let packet = Packet {
+            body: Some(Body::CallData(CallData {
+                data: vec![4, 5, 6],
+                data_is_zero: false,
+                complete: true,
+                error: String::new(),
+            })),
+        };
+
+        codec.encode(packet.clone(), &mut buf).unwrap();
+
+        let decoded = codec.decode(&mut buf).unwrap().unwrap();
+        assert_eq!(decoded, packet);
+    }
+
+    #[test]
+    fn test_codec_roundtrip_call_cancel() {
+        let mut codec = PacketCodec::new();
+        let mut buf = BytesMut::new();
+
+        let packet = Packet {
+            body: Some(Body::CallCancel(true)),
+        };
+
+        codec.encode(packet.clone(), &mut buf).unwrap();
+
+        let decoded = codec.decode(&mut buf).unwrap().unwrap();
+        assert_eq!(decoded, packet);
+    }
+
+    #[test]
+    fn test_codec_partial_read() {
+        let mut codec = PacketCodec::new();
+        let mut buf = BytesMut::new();
+
+        let packet = Packet {
+            body: Some(Body::CallData(CallData {
+                data: vec![1, 2, 3, 4, 5],
+                data_is_zero: false,
+                complete: false,
+                error: String::new(),
+            })),
+        };
+
+        // Encode the packet.
+        codec.encode(packet.clone(), &mut buf).unwrap();
+
+        // Split the buffer to simulate partial read.
+        let full_buf = buf.clone();
+        buf.truncate(3); // Only the first 3 bytes.
+
+        // Should return None (need more data).
+        assert!(codec.decode(&mut buf).unwrap().is_none());
+
+        // Add the rest of the data.
+        buf.extend_from_slice(&full_buf[3..]);
+
+        // Now it should decode.
+        let decoded = codec.decode(&mut buf).unwrap().unwrap();
+        assert_eq!(decoded, packet);
+    }
+
+    #[test]
+    fn test_codec_message_too_large() {
+        let mut codec = PacketCodec::new();
+        let mut buf = BytesMut::new();
+
+        // Create a message that exceeds the maximum size.
+        let packet = Packet {
+            body: Some(Body::CallData(CallData {
+                data: vec![0u8; MAX_MESSAGE_SIZE + 1],
+                data_is_zero: false,
+                complete: false,
+                error: String::new(),
+            })),
+        };
+
+        let result = codec.encode(packet, &mut buf);
+        assert!(matches!(result, Err(Error::MessageTooLarge(_, _))));
+    }
+}