io.rsocket.core.ServerSetup Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of rsocket-core Show documentation
Core functionality for the RSocket library
The newest version!
/*
 * Copyright 2015-2020 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package io.rsocket.core;

import static io.rsocket.keepalive.KeepAliveHandler.*;

import io.netty.buffer.ByteBuf;
import io.rsocket.DuplexConnection;
import io.rsocket.RSocketErrorException;
import io.rsocket.exceptions.RejectedResumeException;
import io.rsocket.exceptions.UnsupportedSetupException;
import io.rsocket.frame.ResumeFrameCodec;
import io.rsocket.frame.SetupFrameCodec;
import io.rsocket.keepalive.KeepAliveHandler;
import io.rsocket.resume.*;
import java.nio.channels.ClosedChannelException;
import java.time.Duration;
import java.util.function.BiFunction;
import java.util.function.Function;
import reactor.core.publisher.Mono;
import reactor.util.function.Tuple2;

abstract class ServerSetup {

  final Duration timeout;

  protected ServerSetup(Duration timeout) {
    this.timeout = timeout;
  }

  Mono> init(DuplexConnection connection) {
    return Mono.>create(
            sink -> sink.onRequest(__ -> new SetupHandlingDuplexConnection(connection, sink)))
        .timeout(this.timeout)
        .or(connection.onClose().then(Mono.error(ClosedChannelException::new)));
  }

  abstract Mono acceptRSocketSetup(
      ByteBuf frame,
      DuplexConnection clientServerConnection,
      BiFunction> then);

  abstract Mono acceptRSocketResume(ByteBuf frame, DuplexConnection connection);

  void dispose() {}

  void sendError(DuplexConnection duplexConnection, RSocketErrorException exception) {
    duplexConnection.sendErrorAndClose(exception);
    duplexConnection.receive().subscribe();
  }

  static class DefaultServerSetup extends ServerSetup {

    DefaultServerSetup(Duration timeout) {
      super(timeout);
    }

    @Override
    public Mono acceptRSocketSetup(
        ByteBuf frame,
        DuplexConnection duplexConnection,
        BiFunction> then) {

      if (SetupFrameCodec.resumeEnabled(frame)) {
        sendError(duplexConnection, new UnsupportedSetupException("resume not supported"));
        return duplexConnection.onClose();
      } else {
        return then.apply(new DefaultKeepAliveHandler(), duplexConnection);
      }
    }

    @Override
    public Mono acceptRSocketResume(ByteBuf frame, DuplexConnection duplexConnection) {
      sendError(duplexConnection, new RejectedResumeException("resume not supported"));
      return duplexConnection.onClose();
    }
  }

  static class ResumableServerSetup extends ServerSetup {
    private final SessionManager sessionManager;
    private final Duration resumeSessionDuration;
    private final Duration resumeStreamTimeout;
    private final Function resumeStoreFactory;
    private final boolean cleanupStoreOnKeepAlive;

    ResumableServerSetup(
        Duration timeout,
        SessionManager sessionManager,
        Duration resumeSessionDuration,
        Duration resumeStreamTimeout,
        Function resumeStoreFactory,
        boolean cleanupStoreOnKeepAlive) {
      super(timeout);
      this.sessionManager = sessionManager;
      this.resumeSessionDuration = resumeSessionDuration;
      this.resumeStreamTimeout = resumeStreamTimeout;
      this.resumeStoreFactory = resumeStoreFactory;
      this.cleanupStoreOnKeepAlive = cleanupStoreOnKeepAlive;
    }

    @Override
    public Mono acceptRSocketSetup(
        ByteBuf frame,
        DuplexConnection duplexConnection,
        BiFunction> then) {

      if (SetupFrameCodec.resumeEnabled(frame)) {
        ByteBuf resumeToken = SetupFrameCodec.resumeToken(frame);

        final ResumableFramesStore resumableFramesStore = resumeStoreFactory.apply(resumeToken);
        final ResumableDuplexConnection resumableDuplexConnection =
            new ResumableDuplexConnection(
                "server", resumeToken, duplexConnection, resumableFramesStore);
        final ServerRSocketSession serverRSocketSession =
            new ServerRSocketSession(
                resumeToken,
                resumableDuplexConnection,
                duplexConnection,
                resumableFramesStore,
                resumeSessionDuration,
                cleanupStoreOnKeepAlive);

        sessionManager.save(serverRSocketSession, resumeToken);

        return then.apply(
            new ResumableKeepAliveHandler(
                resumableDuplexConnection, serverRSocketSession, serverRSocketSession),
            resumableDuplexConnection);
      } else {
        return then.apply(new DefaultKeepAliveHandler(), duplexConnection);
      }
    }

    @Override
    public Mono acceptRSocketResume(ByteBuf frame, DuplexConnection duplexConnection) {
      ServerRSocketSession session = sessionManager.get(ResumeFrameCodec.token(frame));
      if (session != null) {
        session.resumeWith(frame, duplexConnection);
        return duplexConnection.onClose();
      } else {
        sendError(duplexConnection, new RejectedResumeException("unknown resume token"));
        return duplexConnection.onClose();
      }
    }

    @Override
    public void dispose() {
      sessionManager.dispose();
    }
  }
}