55import io .rsocket .frame .SetupFrameFlyweight ;
66import io .rsocket .frame .VersionFlyweight ;
77import io .rsocket .internal .ClientServerInputMultiplexer ;
8+ import io .rsocket .plugins .DuplexConnectionInterceptor ;
9+ import io .rsocket .plugins .PluginRegistry ;
10+ import io .rsocket .plugins .Plugins ;
11+ import io .rsocket .plugins .RSocketInterceptor ;
812import io .rsocket .transport .ClientTransport ;
913import io .rsocket .transport .ServerTransport ;
1014import io .rsocket .util .PayloadImpl ;
@@ -105,6 +109,7 @@ class ClientRSocketFactory
105109 private Supplier <io .rsocket .transport .ClientTransport > transportClient ;
106110 private Consumer <Throwable > errorConsumer = Throwable ::printStackTrace ;
107111 private int mtu = 0 ;
112+ private PluginRegistry plugins = new PluginRegistry (Plugins .defaultPlugins ());
108113 private int flags = SetupFrameFlyweight .FLAGS_STRICT_INTERPRETATION ;
109114
110115 private Payload setupPayload = PayloadImpl .EMPTY ;
@@ -116,6 +121,21 @@ class ClientRSocketFactory
116121 private String dataMineType = "application/binary" ;
117122 private String metadataMimeType = "application/binary" ;
118123
124+ public ClientRSocketFactory addConnectionPlugin (DuplexConnectionInterceptor interceptor ) {
125+ plugins .addConnectionPlugin (interceptor );
126+ return this ;
127+ }
128+
129+ public ClientRSocketFactory addClientPlugin (RSocketInterceptor interceptor ) {
130+ plugins .addClientPlugin (interceptor );
131+ return this ;
132+ }
133+
134+ public ClientRSocketFactory addServerPlugin (RSocketInterceptor interceptor ) {
135+ plugins .addServerPlugin (interceptor );
136+ return this ;
137+ }
138+
119139 @ Override
120140 public ClientRSocketFactory keepAlive () {
121141 tickPeriod = Duration .ofSeconds (20 );
@@ -225,15 +245,13 @@ public Mono<RSocket> start() {
225245 metadataMimeType ,
226246 setupPayload );
227247
228- ClientServerInputMultiplexer multiplexer ;
229248 if (mtu > 0 ) {
230- multiplexer =
231- new ClientServerInputMultiplexer (
232- new FragmentationDuplexConnection (connection , mtu ));
233- } else {
234- multiplexer = new ClientServerInputMultiplexer (connection );
249+ connection = new FragmentationDuplexConnection (connection , mtu );
235250 }
236251
252+ ClientServerInputMultiplexer multiplexer =
253+ new ClientServerInputMultiplexer (connection , plugins );
254+
237255 RSocketClient rSocketClient =
238256 new RSocketClient (
239257 multiplexer .asClientConnection (),
@@ -243,23 +261,25 @@ public Mono<RSocket> start() {
243261 ackTimeout ,
244262 missedAcks );
245263
246- return Plugins .CLIENT_REACTIVE_SOCKET_INTERCEPTOR
247- .apply (rSocketClient )
248- .then (
249- wrappedClientRSocket -> {
250- RSocket unwrappedServerSocket =
251- acceptor .get ().apply (wrappedClientRSocket );
252- return Plugins .SERVER_REACTIVE_SOCKET_INTERCEPTOR
253- .apply (unwrappedServerSocket )
254- .doOnNext (
255- rSocket ->
256- new RSocketServer (
257- multiplexer .asServerConnection (),
258- rSocket ,
259- errorConsumer ))
260- .then (connection .sendOne (setupFrame ))
261- .then (Mono .just (wrappedClientRSocket ));
262- });
264+ Mono <RSocket > wrappedRSocketClient =
265+ Mono .just (rSocketClient ).map (plugins ::applyClient );
266+
267+ DuplexConnection finalConnection = connection ;
268+ return wrappedRSocketClient .then (
269+ wrappedClientRSocket -> {
270+ RSocket unwrappedServerSocket = acceptor .get ().apply (wrappedClientRSocket );
271+
272+ Mono <RSocket > wrappedRSocketServer =
273+ Mono .just (unwrappedServerSocket ).map (plugins ::applyServer );
274+
275+ return wrappedRSocketServer
276+ .doOnNext (
277+ rSocket ->
278+ new RSocketServer (
279+ multiplexer .asServerConnection (), rSocket , errorConsumer ))
280+ .then (finalConnection .sendOne (setupFrame ))
281+ .then (Mono .just (wrappedClientRSocket ));
282+ });
263283 });
264284 }
265285 }
@@ -274,9 +294,25 @@ class ServerRSocketFactory
274294 private Supplier <io .rsocket .transport .ServerTransport > transportServer ;
275295 private Consumer <Throwable > errorConsumer = Throwable ::printStackTrace ;
276296 private int mtu = 0 ;
297+ private PluginRegistry plugins = new PluginRegistry (Plugins .defaultPlugins ());
277298
278299 private ServerRSocketFactory () {}
279300
301+ public ServerRSocketFactory addConnectionPlugin (DuplexConnectionInterceptor interceptor ) {
302+ plugins .addConnectionPlugin (interceptor );
303+ return this ;
304+ }
305+
306+ public ServerRSocketFactory addClientPlugin (RSocketInterceptor interceptor ) {
307+ plugins .addClientPlugin (interceptor );
308+ return this ;
309+ }
310+
311+ public ServerRSocketFactory addServerPlugin (RSocketInterceptor interceptor ) {
312+ plugins .addServerPlugin (interceptor );
313+ return this ;
314+ }
315+
280316 @ Override
281317 public Transport <io .rsocket .transport .ServerTransport , Closeable > acceptor (
282318 Supplier <SocketAcceptor > acceptor ) {
@@ -312,15 +348,13 @@ public Mono<Closeable> start() {
312348 .get ()
313349 .start (
314350 connection -> {
315- ClientServerInputMultiplexer multiplexer ;
316351 if (mtu > 0 ) {
317- multiplexer =
318- new ClientServerInputMultiplexer (
319- new FragmentationDuplexConnection (connection , mtu ));
320- } else {
321- multiplexer = new ClientServerInputMultiplexer (connection );
352+ connection = new FragmentationDuplexConnection (connection , mtu );
322353 }
323354
355+ ClientServerInputMultiplexer multiplexer =
356+ new ClientServerInputMultiplexer (connection , plugins );
357+
324358 return multiplexer
325359 .asStreamZeroConnection ()
326360 .receive ()
@@ -348,16 +382,10 @@ private Mono<? extends Void> processSetupFrame(
348382 new RSocketClient (
349383 multiplexer .asServerConnection (), errorConsumer , StreamIdSupplier .serverSupplier ());
350384
351- Mono <RSocket > wrappedRSocketClient =
352- Plugins .CLIENT_REACTIVE_SOCKET_INTERCEPTOR .apply (rSocketClient );
385+ Mono <RSocket > wrappedRSocketClient = Mono .just (rSocketClient ).map (plugins ::applyClient );
353386
354387 return wrappedRSocketClient
355- .then (
356- sender ->
357- acceptor
358- .get ()
359- .accept (setupPayload , sender )
360- .then (Plugins .SERVER_REACTIVE_SOCKET_INTERCEPTOR ::apply ))
388+ .then (sender -> acceptor .get ().accept (setupPayload , sender ).map (plugins ::applyServer ))
361389 .map (
362390 handler ->
363391 new RSocketServer (multiplexer .asClientConnection (), handler , errorConsumer ))
0 commit comments