mirror of
https://github.com/signalapp/Signal-Server
synced 2026-04-02 17:32:49 +01:00
270 lines
11 KiB
Java
270 lines
11 KiB
Java
/*
|
|
* Copyright 2013-2020 Signal Messenger, LLC
|
|
* SPDX-License-Identifier: AGPL-3.0-only
|
|
*/
|
|
package org.whispersystems.websocket;
|
|
|
|
import com.google.common.annotations.VisibleForTesting;
|
|
import com.google.common.net.HttpHeaders;
|
|
import com.google.protobuf.UninitializedMessageException;
|
|
import org.eclipse.jetty.websocket.api.MessageTooLargeException;
|
|
import org.eclipse.jetty.websocket.api.RemoteEndpoint;
|
|
import org.eclipse.jetty.websocket.api.Session;
|
|
import org.eclipse.jetty.websocket.api.WebSocketListener;
|
|
import org.glassfish.jersey.internal.MapPropertiesDelegate;
|
|
import org.glassfish.jersey.server.ApplicationHandler;
|
|
import org.glassfish.jersey.server.ContainerRequest;
|
|
import org.glassfish.jersey.server.ContainerResponse;
|
|
import org.slf4j.Logger;
|
|
import org.slf4j.LoggerFactory;
|
|
import org.whispersystems.websocket.logging.WebsocketRequestLog;
|
|
import org.whispersystems.websocket.messages.InvalidMessageException;
|
|
import org.whispersystems.websocket.messages.WebSocketMessage;
|
|
import org.whispersystems.websocket.messages.WebSocketMessageFactory;
|
|
import org.whispersystems.websocket.messages.WebSocketRequestMessage;
|
|
import org.whispersystems.websocket.messages.WebSocketResponseMessage;
|
|
import org.whispersystems.websocket.session.ContextPrincipal;
|
|
import org.whispersystems.websocket.session.WebSocketSessionContext;
|
|
import org.whispersystems.websocket.setup.WebSocketConnectListener;
|
|
|
|
import javax.ws.rs.core.MultivaluedMap;
|
|
import javax.ws.rs.core.Response;
|
|
import java.io.ByteArrayInputStream;
|
|
import java.io.ByteArrayOutputStream;
|
|
import java.io.IOException;
|
|
import java.net.URI;
|
|
import java.nio.ByteBuffer;
|
|
import java.security.Principal;
|
|
import java.util.HashMap;
|
|
import java.util.LinkedList;
|
|
import java.util.List;
|
|
import java.util.Map;
|
|
import java.util.Optional;
|
|
import java.util.Set;
|
|
import java.util.concurrent.CompletableFuture;
|
|
import java.util.concurrent.ConcurrentHashMap;
|
|
|
|
|
|
@SuppressWarnings("OptionalUsedAsFieldOrParameterType")
|
|
public class WebSocketResourceProvider<T extends Principal> implements WebSocketListener {
|
|
|
|
/**
|
|
* A static exception instance passed to outstanding requests (via {@code completeExceptionally} in
|
|
* {@link #onWebSocketClose(int, String)}
|
|
*/
|
|
public static final IOException CONNECTION_CLOSED_EXCEPTION = new IOException("Connection closed!");
|
|
private static final Logger logger = LoggerFactory.getLogger(WebSocketResourceProvider.class);
|
|
|
|
private final Map<Long, CompletableFuture<WebSocketResponseMessage>> requestMap = new ConcurrentHashMap<>();
|
|
|
|
private final T authenticated;
|
|
private final WebSocketMessageFactory messageFactory;
|
|
private final Optional<WebSocketConnectListener> connectListener;
|
|
private final ApplicationHandler jerseyHandler;
|
|
private final WebsocketRequestLog requestLog;
|
|
private final long idleTimeoutMillis;
|
|
private final String remoteAddress;
|
|
|
|
private Session session;
|
|
private RemoteEndpoint remoteEndpoint;
|
|
private WebSocketSessionContext context;
|
|
|
|
private static final Set<String> EXCLUDED_UPGRADE_REQUEST_HEADERS = Set.of("connection", "upgrade");
|
|
|
|
public WebSocketResourceProvider(String remoteAddress,
|
|
ApplicationHandler jerseyHandler,
|
|
WebsocketRequestLog requestLog,
|
|
T authenticated,
|
|
WebSocketMessageFactory messageFactory,
|
|
Optional<WebSocketConnectListener> connectListener,
|
|
long idleTimeoutMillis)
|
|
{
|
|
this.remoteAddress = remoteAddress;
|
|
this.jerseyHandler = jerseyHandler;
|
|
this.requestLog = requestLog;
|
|
this.authenticated = authenticated;
|
|
this.messageFactory = messageFactory;
|
|
this.connectListener = connectListener;
|
|
this.idleTimeoutMillis = idleTimeoutMillis;
|
|
}
|
|
|
|
@Override
|
|
public void onWebSocketConnect(Session session) {
|
|
this.session = session;
|
|
this.remoteEndpoint = session.getRemote();
|
|
this.context = new WebSocketSessionContext(new WebSocketClient(session, remoteEndpoint, messageFactory, requestMap));
|
|
this.context.setAuthenticated(authenticated);
|
|
this.session.setIdleTimeout(idleTimeoutMillis);
|
|
|
|
connectListener.ifPresent(listener -> listener.onWebSocketConnect(this.context));
|
|
}
|
|
|
|
@Override
|
|
public void onWebSocketError(Throwable cause) {
|
|
logger.debug("onWebSocketError", cause);
|
|
|
|
final int closeCode;
|
|
final String message;
|
|
if (cause instanceof MessageTooLargeException) {
|
|
closeCode = 1009;
|
|
message = "Frame too large";
|
|
} else {
|
|
closeCode = 1011;
|
|
message = "Server error";
|
|
}
|
|
|
|
close(session, closeCode, message);
|
|
}
|
|
|
|
@Override
|
|
public void onWebSocketBinary(byte[] payload, int offset, int length) {
|
|
try {
|
|
WebSocketMessage webSocketMessage = messageFactory.parseMessage(payload, offset, length);
|
|
|
|
switch (webSocketMessage.getType()) {
|
|
case REQUEST_MESSAGE:
|
|
handleRequest(webSocketMessage.getRequestMessage());
|
|
break;
|
|
case RESPONSE_MESSAGE:
|
|
handleResponse(webSocketMessage.getResponseMessage());
|
|
break;
|
|
default:
|
|
close(session, 1018, "Badly formatted");
|
|
break;
|
|
}
|
|
} catch (UninitializedMessageException | InvalidMessageException e) {
|
|
logger.debug("Parsing", e);
|
|
close(session, 1018, "Badly formatted");
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public void onWebSocketClose(int statusCode, String reason) {
|
|
if (context != null) {
|
|
context.notifyClosed(statusCode, reason);
|
|
|
|
for (long requestId : requestMap.keySet()) {
|
|
CompletableFuture<WebSocketResponseMessage> outstandingRequest = requestMap.remove(requestId);
|
|
|
|
if (outstandingRequest != null) {
|
|
outstandingRequest.completeExceptionally(CONNECTION_CLOSED_EXCEPTION);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public void onWebSocketText(String message) {
|
|
logger.debug("onWebSocketText!");
|
|
}
|
|
|
|
private void handleRequest(WebSocketRequestMessage requestMessage) {
|
|
ContainerRequest containerRequest = new ContainerRequest(null, URI.create(requestMessage.getPath()), requestMessage.getVerb(), new WebSocketSecurityContext(new ContextPrincipal(context)), new MapPropertiesDelegate(new HashMap<>()), jerseyHandler.getConfiguration());
|
|
containerRequest.headers(getCombinedHeaders(session.getUpgradeRequest().getHeaders(), requestMessage.getHeaders()));
|
|
|
|
if (requestMessage.getBody().isPresent()) {
|
|
containerRequest.setEntityStream(new ByteArrayInputStream(requestMessage.getBody().get()));
|
|
}
|
|
|
|
ByteArrayOutputStream responseBody = new ByteArrayOutputStream();
|
|
CompletableFuture<ContainerResponse> responseFuture = (CompletableFuture<ContainerResponse>) jerseyHandler.apply(containerRequest, responseBody);
|
|
|
|
responseFuture.thenAccept(response -> {
|
|
sendResponse(requestMessage, response, responseBody);
|
|
requestLog.log(remoteAddress, containerRequest, response);
|
|
}).exceptionally(exception -> {
|
|
logger.warn("Websocket Error: " + requestMessage.getVerb() + " " + requestMessage.getPath() + "\n" + requestMessage.getBody(), exception);
|
|
sendErrorResponse(requestMessage, Response.status(500).build());
|
|
requestLog.log(remoteAddress, containerRequest, new ContainerResponse(containerRequest, Response.status(500).build()));
|
|
return null;
|
|
});
|
|
}
|
|
|
|
@VisibleForTesting
|
|
static Map<String, List<String>> getCombinedHeaders(final Map<String, List<String>> upgradeRequestHeaders, final Map<String, String> requestMessageHeaders) {
|
|
final Map<String, List<String>> combinedHeaders = new HashMap<>();
|
|
|
|
upgradeRequestHeaders.entrySet().stream()
|
|
.filter(entry -> shouldIncludeUpgradeRequestHeader(entry.getKey()))
|
|
.forEach(entry -> combinedHeaders.put(entry.getKey(), entry.getValue()));
|
|
|
|
requestMessageHeaders.entrySet().stream()
|
|
.filter(entry -> shouldIncludeRequestMessageHeader(entry.getKey()))
|
|
.forEach(entry -> combinedHeaders.put(entry.getKey(), List.of(entry.getValue())));
|
|
|
|
return combinedHeaders;
|
|
}
|
|
|
|
@VisibleForTesting
|
|
static boolean shouldIncludeUpgradeRequestHeader(final String header) {
|
|
return !EXCLUDED_UPGRADE_REQUEST_HEADERS.contains(header.toLowerCase()) && !header.toLowerCase().contains("websocket-");
|
|
}
|
|
|
|
@VisibleForTesting
|
|
static boolean shouldIncludeRequestMessageHeader(final String header) {
|
|
return !HttpHeaders.X_FORWARDED_FOR.equalsIgnoreCase(header.trim());
|
|
}
|
|
|
|
private void handleResponse(WebSocketResponseMessage responseMessage) {
|
|
CompletableFuture<WebSocketResponseMessage> future = requestMap.remove(responseMessage.getRequestId());
|
|
|
|
if (future != null) {
|
|
future.complete(responseMessage);
|
|
}
|
|
}
|
|
|
|
private void close(Session session, int status, String message) {
|
|
session.close(status, message);
|
|
}
|
|
|
|
private void sendResponse(WebSocketRequestMessage requestMessage, ContainerResponse response, ByteArrayOutputStream responseBody) {
|
|
if (requestMessage.hasRequestId()) {
|
|
byte[] body = responseBody.toByteArray();
|
|
|
|
if (body.length <= 0) {
|
|
body = null;
|
|
}
|
|
|
|
byte[] responseBytes = messageFactory.createResponse(requestMessage.getRequestId(),
|
|
response.getStatus(),
|
|
response.getStatusInfo().getReasonPhrase(),
|
|
getHeaderList(response.getStringHeaders()),
|
|
Optional.ofNullable(body))
|
|
.toByteArray();
|
|
|
|
remoteEndpoint.sendBytesByFuture(ByteBuffer.wrap(responseBytes));
|
|
}
|
|
}
|
|
|
|
private void sendErrorResponse(WebSocketRequestMessage requestMessage, Response error) {
|
|
if (requestMessage.hasRequestId()) {
|
|
WebSocketMessage response = messageFactory.createResponse(requestMessage.getRequestId(),
|
|
error.getStatus(),
|
|
"Error response",
|
|
getHeaderList(error.getStringHeaders()),
|
|
Optional.empty());
|
|
|
|
remoteEndpoint.sendBytesByFuture(ByteBuffer.wrap(response.toByteArray()));
|
|
}
|
|
}
|
|
|
|
|
|
@VisibleForTesting
|
|
WebSocketSessionContext getContext() {
|
|
return context;
|
|
}
|
|
|
|
@VisibleForTesting
|
|
static List<String> getHeaderList(final MultivaluedMap<String, String> headerMap) {
|
|
final List<String> headers = new LinkedList<>();
|
|
|
|
if (headerMap != null) {
|
|
for (String key : headerMap.keySet()) {
|
|
headers.add(key + ":" + headerMap.getFirst(key));
|
|
}
|
|
}
|
|
|
|
return headers;
|
|
}
|
|
}
|