Support reject setup.

Author: jjeffcaii

README.md

@@ -33,7 +33,9 @@ async fn main() -> Result<(), Box<dyn Error>> {
         .transport(&addr)
         .acceptor(|setup, _socket| {
             info!("accept setup: {:?}", setup);
-            Box::new(EchoRSocket)
+            Ok(Box::new(EchoRSocket))
+            // Or you can reject setup
+            // Err(From::from("SETUP_NOT_ALLOW"))
         })
         .serve()
         .await

examples/echo/main.rs

@@ -16,7 +16,9 @@ async fn main() -> Result<(), Box<dyn Error>> {
         .transport(&addr)
         .acceptor(|setup, _socket| {
             info!("accept setup: {:?}", setup);
-            Box::new(EchoRSocket)
+            Ok(Box::new(EchoRSocket))
+            // Or you can reject setup
+            // Err(From::from("SETUP_NOT_ALLOW"))
         })
         .on_start(|| info!("+++++++ echo server started! +++++++"))
         .serve()

examples/proxy/main.rs

@@ -16,15 +16,15 @@ async fn main() -> Result<(), Box<dyn Error>> {
     RSocketFactory::receive()
         .acceptor(|setup, _sending_socket| {
             info!("incoming socket: setup={:?}", setup);
-            Box::new(block_on(async move {
+            Ok(Box::new(block_on(async move {
                 RSocketFactory::connect()
                     .acceptor(|| Box::new(EchoRSocket))
                     .setup(Payload::from("I'm Rust!"))
                     .transport("tcp://127.0.0.1:7878")
                     .start()
                     .await
                     .unwrap()
-            }))
+            })))
         })
         .transport("tcp://127.0.0.1:7979")
         .serve()

src/spi.rs

@@ -57,7 +57,7 @@ impl RSocket for EchoRSocket {
     }
 }
 
-pub struct EmptyRSocket;
+pub(crate) struct EmptyRSocket;
 
 impl EmptyRSocket {
     fn must_failed(&self) -> RSocketError {

src/transport/socket.rs

@@ -1,6 +1,6 @@
 use super::misc::{self, Counter, StreamID};
 use super::spi::*;
-use crate::errors::{ErrorKind, RSocketError, ERR_APPLICATION};
+use crate::errors::{self, ErrorKind, RSocketError};
 use crate::frame::{self, Body, Frame};
 use crate::payload::{Payload, SetupPayload};
 use crate::result::RSocketResult;
@@ -103,7 +103,17 @@ impl DuplexSocket {
             let flag = msg.get_flag();
             misc::debug_frame(false, &msg);
             match msg.get_body() {
-                Body::Setup(v) => self.on_setup(&acceptor, sid, flag, SetupPayload::from(v)),
+                Body::Setup(v) => {
+                    if let Err(e) = self.on_setup(&acceptor, sid, flag, SetupPayload::from(v)) {
+                        let errmsg = format!("{}", e);
+                        let sending = frame::Error::builder(0, 0)
+                            .set_code(errors::ERR_REJECT_SETUP)
+                            .set_data(Bytes::from(errmsg))
+                            .build();
+                        self.tx.send(sending).unwrap();
+                        return;
+                    }
+                }
                 Body::Resume(v) => {
                     // TODO: support resume
                 }
@@ -225,16 +235,28 @@ impl DuplexSocket {
     }
 
     #[inline]
-    fn on_setup(&self, acceptor: &Acceptor, sid: u32, flag: u16, setup: SetupPayload) {
+    fn on_setup(
+        &self,
+        acceptor: &Acceptor,
+        sid: u32,
+        flag: u16,
+        setup: SetupPayload,
+    ) -> Result<(), Box<dyn Error>> {
         match acceptor {
             Acceptor::Simple(gen) => {
                 self.responder.set(gen());
+                Ok(())
             }
-            Acceptor::Generate(gen) => {
-                self.responder.set(gen(setup, Box::new(self.clone())));
-            }
+            Acceptor::Generate(gen) => match gen(setup, Box::new(self.clone())) {
+                Ok(it) => {
+                    self.responder.set(it);
+                    Ok(())
+                }
+                Err(e) => Err(e),
+            },
             Acceptor::Empty() => {
                 self.responder.set(Box::new(EmptyRSocket));
+                Ok(())
             }
         }
     }
@@ -277,7 +299,7 @@ impl DuplexSocket {
                     bu.build()
                 }
                 Err(e) => frame::Error::builder(sid, 0)
-                    .set_code(ERR_APPLICATION)
+                    .set_code(errors::ERR_APPLICATION)
                     .set_data(Bytes::from("TODO: should be error details"))
                     .build(),
             };

src/transport/spi.rs

@@ -1,7 +1,9 @@
 use crate::frame::Frame;
 use crate::payload::SetupPayload;
 use crate::spi::RSocket;
+use std::error::Error;
 use std::future::Future;
+use std::result::Result;
 use std::sync::Arc;
 use tokio::sync::{mpsc, oneshot};
 
@@ -33,7 +35,8 @@ impl Transport {
     }
 }
 
-type FnAcceptorWithSetup = fn(SetupPayload, Box<dyn RSocket>) -> Box<dyn RSocket>;
+pub type FnAcceptorWithSetup =
+    fn(SetupPayload, Box<dyn RSocket>) -> Result<Box<dyn RSocket>, Box<dyn Error>>;
 
 pub(crate) enum Acceptor {
     Simple(Arc<fn() -> Box<dyn RSocket>>),

src/x/server.rs

@@ -3,19 +3,19 @@ use crate::errors::RSocketError;
 use crate::frame::{self, Frame};
 use crate::payload::SetupPayload;
 use crate::spi::{EmptyRSocket, RSocket};
-use crate::transport::{Acceptor, DuplexSocket};
+use crate::transport::{Acceptor, DuplexSocket, FnAcceptorWithSetup};
 use std::error::Error;
 use std::net::SocketAddr;
+use std::result::Result;
 use std::sync::Arc;
 use tokio::net::TcpListener;
 use tokio::sync::mpsc;
 
-type FnSetup = fn(SetupPayload, Box<dyn RSocket>) -> Box<dyn RSocket>;
 type FnStart = fn();
 
 pub struct ServerBuilder {
     uri: Option<String>,
-    on_setup: FnSetup,
+    on_setup: FnAcceptorWithSetup,
     start_handler: Option<FnStart>,
 }
 
@@ -34,7 +34,7 @@ impl ServerBuilder {
         }
     }
 
-    pub fn acceptor(mut self, handler: FnSetup) -> Self {
+    pub fn acceptor(mut self, handler: FnAcceptorWithSetup) -> Self {
         self.on_setup = handler;
         self
     }
@@ -64,7 +64,7 @@ impl ServerBuilder {
     #[inline]
     async fn serve_tcp(
         addr: SocketAddr,
-        on_setup: FnSetup,
+        on_setup: FnAcceptorWithSetup,
         on_start: Option<FnStart>,
     ) -> Result<(), Box<dyn Error>> {
         let mut listener = TcpListener::bind(&addr).await.unwrap();
@@ -91,6 +91,9 @@ impl ServerBuilder {
 }
 
 #[inline]
-fn on_setup_noop(_setup: SetupPayload, _socket: Box<dyn RSocket>) -> Box<dyn RSocket> {
-    Box::new(EmptyRSocket)
+fn on_setup_noop(
+    _setup: SetupPayload,
+    _socket: Box<dyn RSocket>,
+) -> Result<Box<dyn RSocket>, Box<dyn Error>> {
+    Ok(Box::new(EmptyRSocket))
 }