3131import io .rsocket .internal .ClientSetup ;
3232import io .rsocket .internal .ServerSetup ;
3333import io .rsocket .keepalive .KeepAliveHandler ;
34- import io .rsocket .lease .* ;
35- import io .rsocket .plugins . DuplexConnectionInterceptor ;
36- import io .rsocket .plugins . PluginRegistry ;
37- import io .rsocket .plugins . Plugins ;
38- import io .rsocket .plugins .RSocketInterceptor ;
34+ import io .rsocket .lease .LeaseStats ;
35+ import io .rsocket .lease . Leases ;
36+ import io .rsocket .lease . RequesterLeaseHandler ;
37+ import io .rsocket .lease . ResponderLeaseHandler ;
38+ import io .rsocket .plugins .* ;
3939import io .rsocket .resume .*;
4040import io .rsocket .transport .ClientTransport ;
4141import io .rsocket .transport .ServerTransport ;
4444import io .rsocket .util .MultiSubscriberRSocket ;
4545import java .time .Duration ;
4646import java .util .Objects ;
47- import java .util .function .*;
47+ import java .util .function .BiFunction ;
48+ import java .util .function .Consumer ;
49+ import java .util .function .Function ;
50+ import java .util .function .Supplier ;
4851import reactor .core .publisher .Mono ;
4952
5053/** Factory for creating RSocket clients and servers. */
@@ -93,10 +96,7 @@ default <T extends Closeable> Start<T> transport(ServerTransport<T> transport) {
9396 public static class ClientRSocketFactory implements ClientTransportAcceptor {
9497 private static final String CLIENT_TAG = "client" ;
9598
96- private Supplier <Function <RSocket , RSocket >> acceptor =
97- () -> rSocket -> new AbstractRSocket () {};
98-
99- private BiFunction <ConnectionSetupPayload , RSocket , RSocket > biAcceptor ;
99+ private SocketAcceptor acceptor = (setup , sendingSocket ) -> Mono .just (new AbstractRSocket () {});
100100
101101 private Consumer <Throwable > errorConsumer = Throwable ::printStackTrace ;
102102 private int mtu = 0 ;
@@ -161,6 +161,11 @@ public ClientRSocketFactory addResponderPlugin(RSocketInterceptor interceptor) {
161161 return this ;
162162 }
163163
164+ public ClientRSocketFactory addSocketAcceptorPlugin (SocketAcceptorInterceptor interceptor ) {
165+ plugins .addSocketAcceptorPlugin (interceptor );
166+ return this ;
167+ }
168+
164169 /**
165170 * Deprecated as Keep-Alive is not optional according to spec
166171 *
@@ -268,18 +273,25 @@ public Start<RSocket> transport(Supplier<ClientTransport> transportClient) {
268273 }
269274
270275 public ClientTransportAcceptor acceptor (Function <RSocket , RSocket > acceptor ) {
271- this .acceptor = () -> acceptor ;
272- return StartClient ::new ;
276+ return acceptor (() -> acceptor );
273277 }
274278
275279 public ClientTransportAcceptor acceptor (Supplier <Function <RSocket , RSocket >> acceptor ) {
276- this .acceptor = acceptor ;
277- return StartClient ::new ;
280+ return acceptor (
281+ (SocketAcceptor )
282+ (setup , sendingSocket ) -> Mono .just (acceptor .get ().apply (sendingSocket )));
278283 }
279284
285+ @ Deprecated
280286 public ClientTransportAcceptor acceptor (
281287 BiFunction <ConnectionSetupPayload , RSocket , RSocket > biAcceptor ) {
282- this .biAcceptor = biAcceptor ;
288+ return acceptor (
289+ (SocketAcceptor )
290+ (setup , sendingSocket ) -> Mono .just (biAcceptor .apply (setup , sendingSocket )));
291+ }
292+
293+ public ClientTransportAcceptor acceptor (SocketAcceptor acceptor ) {
294+ this .acceptor = acceptor ;
283295 return StartClient ::new ;
284296 }
285297
@@ -346,6 +358,8 @@ public Mono<RSocket> start() {
346358 rSocketRequester = new MultiSubscriberRSocket (rSocketRequester );
347359 }
348360
361+ RSocket wrappedRSocketRequester = plugins .applyRequester (rSocketRequester );
362+
349363 ByteBuf setupFrame =
350364 SetupFrameFlyweight .encode (
351365 allocator ,
@@ -357,34 +371,38 @@ public Mono<RSocket> start() {
357371 dataMimeType ,
358372 setupPayload );
359373
360- RSocket wrappedRSocketRequester = plugins .applyRequester (rSocketRequester );
361-
362- RSocket rSocketHandler ;
363- if (biAcceptor != null ) {
364- ConnectionSetupPayload setup = ConnectionSetupPayload .create (setupFrame );
365- rSocketHandler = biAcceptor .apply (setup , wrappedRSocketRequester );
366- } else {
367- rSocketHandler = acceptor .get ().apply (wrappedRSocketRequester );
368- }
369-
370- RSocket wrappedRSocketHandler = plugins .applyResponder (rSocketHandler );
371-
372- ResponderLeaseHandler responderLeaseHandler =
373- isLeaseEnabled
374- ? new ResponderLeaseHandler .Impl <>(
375- CLIENT_TAG , allocator , leases .sender (), errorConsumer , leases .stats ())
376- : ResponderLeaseHandler .None ;
377-
378- RSocket rSocketResponder =
379- new RSocketResponder (
380- allocator ,
381- multiplexer .asServerConnection (),
382- wrappedRSocketHandler ,
383- payloadDecoder ,
384- errorConsumer ,
385- responderLeaseHandler );
374+ ConnectionSetupPayload setup = ConnectionSetupPayload .create (setupFrame );
375+
376+ return plugins
377+ .applySocketAcceptorInterceptor (acceptor )
378+ .accept (setup , wrappedRSocketRequester )
379+ .flatMap (
380+ rSocketHandler -> {
381+ RSocket wrappedRSocketHandler = plugins .applyResponder (rSocketHandler );
382+
383+ ResponderLeaseHandler responderLeaseHandler =
384+ isLeaseEnabled
385+ ? new ResponderLeaseHandler .Impl <>(
386+ CLIENT_TAG ,
387+ allocator ,
388+ leases .sender (),
389+ errorConsumer ,
390+ leases .stats ())
391+ : ResponderLeaseHandler .None ;
392+
393+ RSocket rSocketResponder =
394+ new RSocketResponder (
395+ allocator ,
396+ multiplexer .asServerConnection (),
397+ wrappedRSocketHandler ,
398+ payloadDecoder ,
399+ errorConsumer ,
400+ responderLeaseHandler );
386401
387- return wrappedConnection .sendOne (setupFrame ).thenReturn (wrappedRSocketRequester );
402+ return wrappedConnection
403+ .sendOne (setupFrame )
404+ .thenReturn (wrappedRSocketRequester );
405+ });
388406 });
389407 }
390408
@@ -476,6 +494,11 @@ public ServerRSocketFactory addResponderPlugin(RSocketInterceptor interceptor) {
476494 return this ;
477495 }
478496
497+ public ServerRSocketFactory addSocketAcceptorPlugin (SocketAcceptorInterceptor interceptor ) {
498+ plugins .addSocketAcceptorPlugin (interceptor );
499+ return this ;
500+ }
501+
479502 public ServerTransportAcceptor acceptor (SocketAcceptor acceptor ) {
480503 this .acceptor = acceptor ;
481504 return new ServerStart <>();
@@ -644,7 +667,8 @@ private Mono<Void> acceptSetup(
644667 }
645668 RSocket wrappedRSocketRequester = plugins .applyRequester (rSocketRequester );
646669
647- return acceptor
670+ return plugins
671+ .applySocketAcceptorInterceptor (acceptor )
648672 .accept (setupPayload , wrappedRSocketRequester )
649673 .onErrorResume (
650674 err -> sendError (multiplexer , rejectedSetupError (err )).then (Mono .error (err )))
0 commit comments