package io.undertow.websockets.jsr.annotated;
import io.undertow.servlet.api.InstanceHandle;
import io.undertow.websockets.core.AbstractReceiveListener;
import io.undertow.websockets.core.BufferedBinaryMessage;
import io.undertow.websockets.core.BufferedTextMessage;
import io.undertow.websockets.core.CloseMessage;
import io.undertow.websockets.core.StreamSourceFrameChannel;
import io.undertow.websockets.core.WebSocketCallback;
import io.undertow.websockets.core.WebSocketChannel;
import io.undertow.websockets.core.WebSockets;
import io.undertow.websockets.jsr.DefaultPongMessage;
import io.undertow.websockets.jsr.OrderedExecutor;
import io.undertow.websockets.jsr.UndertowSession;
import org.xnio.Buffers;
import org.xnio.Pooled;
import javax.websocket.CloseReason;
import javax.websocket.DecodeException;
import javax.websocket.Endpoint;
import javax.websocket.EndpointConfig;
import javax.websocket.PongMessage;
import javax.websocket.SendHandler;
import javax.websocket.SendResult;
import javax.websocket.Session;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.Reader;
import java.io.StringReader;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Executor;
/**
* @author Stuart Douglas
*/
public class AnnotatedEndpoint extends Endpoint {
private final InstanceHandle<?> instance;
private Executor executor;
private final BoundMethod webSocketOpen;
private final BoundMethod webSocketClose;
private final BoundMethod webSocketError;
private final BoundMethod textMessage;
private final BoundMethod binaryMessage;
private final BoundMethod pongMessage;
AnnotatedEndpoint(final InstanceHandle<?> instance, final BoundMethod webSocketOpen, final BoundMethod webSocketClose, final BoundMethod webSocketError, final BoundMethod textMessage, final BoundMethod binaryMessage, final BoundMethod pongMessage) {
this.instance = instance;
this.webSocketOpen = webSocketOpen;
this.webSocketClose = webSocketClose;
this.webSocketError = webSocketError;
this.textMessage = textMessage;
this.binaryMessage = binaryMessage;
this.pongMessage = pongMessage;
}
@Override
public void onOpen(final Session session, final EndpointConfig endpointConfiguration) {
this.executor = new OrderedExecutor(((UndertowSession)session).getWebSocketChannel().getWorker());
UndertowSession s = (UndertowSession) session;
boolean partialText = textMessage == null || (textMessage.hasParameterType(boolean.class) && !textMessage.getMessageType().equals(boolean.class));
boolean partialBinary = binaryMessage == null || (binaryMessage.hasParameterType(boolean.class) && !binaryMessage.getMessageType().equals(boolean.class));
s.setReceiveListener(new AnnotatedEndpointFrameHandler((UndertowSession) session, partialText, partialBinary));
if (webSocketOpen != null) {
final Map<Class<?>, Object> params = new HashMap<Class<?>, Object>();
params.put(Session.class, session);
params.put(EndpointConfig.class, endpointConfiguration);
params.put(Map.class, session.getPathParameters());
invokeMethod(params, webSocketOpen, s);
}
}
private void invokeMethod(final Map<Class<?>, Object> params, final BoundMethod method, final UndertowSession session) {
session.getContainer().invokeEndpointMethod(executor, new Runnable() {
@Override
public void run() {
try {
method.invoke(instance.getInstance(), params);
} catch (DecodeException e) {
onError(session, e);
}
}
});
}
@Override
public void onClose(final Session session, final CloseReason closeReason) {
if (webSocketClose != null) {
final Map<Class<?>, Object> params = new HashMap<Class<?>, Object>();
params.put(Session.class, session);
params.put(Map.class, session.getPathParameters());
params.put(CloseReason.class, closeReason);
invokeMethod(params, webSocketClose, (UndertowSession) session);
}
}
@Override
public void onError(final Session session, final Throwable thr) {
try {
if (webSocketError != null) {
final Map<Class<?>, Object> params = new HashMap<Class<?>, Object>();
params.put(Session.class, session);
params.put(Throwable.class, thr);
params.put(Map.class, session.getPathParameters());
((UndertowSession) session).getContainer().invokeEndpointMethod(executor, new Runnable() {
@Override
public void run() {
try {
webSocketError.invoke(instance.getInstance(), params);
} catch (DecodeException e) {
throw new RuntimeException(e); //not much we can do here
}
}
});
}
} finally {
((UndertowSession) session).forceClose();
}
}
private class AnnotatedEndpointFrameHandler extends AbstractReceiveListener {
//because fragmented messages can be split on code points we may need
//to buffer data between frames
BufferedTextMessage bufferedTextMessage;
private final UndertowSession session;
private final boolean partialText;
private final boolean partialBinary;
private final SendHandler errorReportingSendHandler = new SendHandler() {
@Override
public void onResult(final SendResult result) {
if (!result.isOK()) {
AnnotatedEndpoint.this.onError(null, result.getException());
}
}
};
public AnnotatedEndpointFrameHandler(final UndertowSession session, boolean partialText, boolean partialBinary) {
this.session = session;
this.partialText = partialText;
this.partialBinary = partialBinary;
}
@Override
protected long getMaxTextBufferSize() {
if (textMessage != null) {
return textMessage.getMaxMessageSize();
}
//TODO: what do we do when there is no handler?
return 1;
}
@Override
protected long getMaxPongBufferSize() {
if (pongMessage != null) {
return pongMessage.getMaxMessageSize();
}
return -1;
}
@Override
protected long getMaxBinaryBufferSize() {
if (binaryMessage != null) {
return binaryMessage.getMaxMessageSize();
}
return 1;
}
@Override
protected void onFullCloseMessage(final WebSocketChannel channel, BufferedBinaryMessage message) throws IOException {
Pooled<ByteBuffer[]> data = message.getData();
final ByteBuffer buffer = WebSockets.mergeBuffers(data.getResource());
final CloseMessage cm = new CloseMessage(buffer);
data.free();
try {
if (webSocketClose != null) {
try {
final Map<Class<?>, Object> params = new HashMap<Class<?>, Object>();
params.put(Session.class, session);
params.put(Map.class, session.getPathParameters());
params.put(CloseReason.class, new CloseReason(CloseReason.CloseCodes.getCloseCode(cm.getReason()), cm.getString()));
invokeMethod(params, webSocketClose, session);
} catch (Exception e) {
AnnotatedEndpoint.this.onError(session, e);
}
}
} finally {
//execute this in the executor to preserve ordering, otherwise the socket
//may be closed while invocations are active
executor.execute(new Runnable() {
@Override
public void run() {
WebSockets.sendClose(buffer.duplicate(), channel, null);
}
});
}
}
@Override
protected void onFullPongMessage(WebSocketChannel channel, BufferedBinaryMessage bufferedBinaryMessage) throws IOException {
if (pongMessage == null) {
return;
}
Pooled<ByteBuffer[]> pooled = bufferedBinaryMessage.getData();
try {
PongMessage message = DefaultPongMessage.create(WebSockets.mergeBuffers(pooled.getResource()));
final Map<Class<?>, Object> params = new HashMap<Class<?>, Object>();
params.put(Session.class, session);
params.put(Map.class, session.getPathParameters());
params.put(PongMessage.class, message);
session.getContainer().invokeEndpointMethod(executor, new Runnable() {
@Override
public void run() {
final Object result;
try {
result = pongMessage.invoke(instance.getInstance(), params);
} catch (Exception e) {
AnnotatedEndpoint.this.onError(session, e);
return;
}
sendResult(result);
}
});
} finally {
pooled.free();
}
}
@Override
protected void onError(WebSocketChannel channel, Throwable error) {
AnnotatedEndpoint.this.onError(session, error);
}
@Override
protected void onText(final WebSocketChannel webSocketChannel, final StreamSourceFrameChannel messageChannel) throws IOException {
if (!partialText) {
super.onText(webSocketChannel, messageChannel);
} else {
if (bufferedTextMessage == null) {
bufferedTextMessage = new BufferedTextMessage(false);
}
bufferedTextMessage.read(messageChannel, new WebSocketCallback<BufferedTextMessage>() {
@Override
public void complete(WebSocketChannel channel, BufferedTextMessage context) {
try {
handleTextMessage(context, context.isComplete());
} finally {
if (messageChannel.isFinalFragment()) {
bufferedTextMessage = null;
}
}
}
@Override
public void onError(WebSocketChannel channel, BufferedTextMessage context, Throwable throwable) {
AnnotatedEndpoint.this.onError(session, throwable);
bufferedTextMessage = null;
}
});
}
}
@Override
protected void onFullTextMessage(WebSocketChannel channel, BufferedTextMessage message) throws IOException {
handleTextMessage(message, true);
}
private void handleTextMessage(BufferedTextMessage message, boolean finalFragment) {
final String data = message.getData();
Object messageObject;
if (textMessage.isDecoderRequired()) {
try {
messageObject = session.getEncoding().decodeText(textMessage.getMessageType(), data);
} catch (DecodeException e) {
AnnotatedEndpoint.this.onError(session, e);
return;
}
} else if (textMessage.getMessageType().equals(Reader.class)) {
messageObject = new StringReader(data);
} else {
messageObject = data;
}
final Map<Class<?>, Object> params = new HashMap<Class<?>, Object>();
params.put(Session.class, session);
params.put(Map.class, session.getPathParameters());
params.put(textMessage.getMessageType(), messageObject);
params.put(boolean.class, finalFragment);
session.getContainer().invokeEndpointMethod(executor, new Runnable() {
@Override
public void run() {
final Object result;
try {
result = textMessage.invoke(instance.getInstance(), params);
} catch (Exception e) {
AnnotatedEndpoint.this.onError(session, e);
return;
}
sendResult(result);
}
});
}
private void sendResult(final Object result) {
if (result != null) {
if (result instanceof String) {
session.getAsyncRemote().sendText((String) result, errorReportingSendHandler);
} else if (result instanceof byte[]) {
session.getAsyncRemote().sendBinary(ByteBuffer.wrap((byte[]) result), errorReportingSendHandler);
} else if (result instanceof ByteBuffer) {
session.getAsyncRemote().sendBinary((ByteBuffer) result, errorReportingSendHandler);
} else {
session.getAsyncRemote().sendObject(result, errorReportingSendHandler);
}
}
}
@Override
protected void onBinary(WebSocketChannel webSocketChannel, final StreamSourceFrameChannel messageChannel) throws IOException {
if (!partialBinary) {
super.onText(webSocketChannel, messageChannel);
} else {
BufferedBinaryMessage buffered = new BufferedBinaryMessage(session.getMaxBinaryMessageBufferSize(), false);
buffered.read(messageChannel, new WebSocketCallback<BufferedBinaryMessage>() {
@Override
public void complete(WebSocketChannel channel, BufferedBinaryMessage context) {
handleBinaryMessage(context, context.isComplete());
}
@Override
public void onError(WebSocketChannel channel, BufferedBinaryMessage context, Throwable throwable) {
AnnotatedEndpoint.this.onError(session, throwable);
}
});
}
}
@Override
protected void onFullBinaryMessage(WebSocketChannel channel, BufferedBinaryMessage message) throws IOException {
handleBinaryMessage(message, true);
}
protected byte[] toArray(ByteBuffer... payload) {
if (payload.length == 1) {
ByteBuffer buf = payload[0];
if (buf.hasArray() && buf.arrayOffset() == 0 && buf.position() == 0) {
return buf.array();
}
}
int size = (int) Buffers.remaining(payload);
byte[] data = new byte[size];
int pos = 0;
for (ByteBuffer buf : payload) {
int toWrite = buf.remaining();
buf.get(data, pos, toWrite);
pos += toWrite;
}
return data;
}
private void handleBinaryMessage(BufferedBinaryMessage message, boolean finalFragment) {
final Pooled<ByteBuffer[]> pooled = message.getData();
try {
final Map<Class<?>, Object> params = new HashMap<Class<?>, Object>();
params.put(Session.class, session);
params.put(Map.class, session.getPathParameters());
if (binaryMessage.isDecoderRequired()) {
try {
params.put(binaryMessage.getMessageType(), session.getEncoding().decodeBinary(binaryMessage.getMessageType(), toArray(pooled.getResource())));
} catch (Exception e) {
AnnotatedEndpoint.this.onError(session, e);
return;
}
} else if (binaryMessage.getMessageType() == ByteBuffer.class) {
params.put(ByteBuffer.class, WebSockets.mergeBuffers(pooled.getResource()));
} else if (binaryMessage.getMessageType() == byte[].class) {
params.put(byte[].class, toArray(pooled.getResource()));
} else if (binaryMessage.getMessageType() == InputStream.class) {
params.put(InputStream.class, new ByteArrayInputStream(toArray(pooled.getResource())));
} else {
try {
params.put(binaryMessage.getMessageType(), session.getEncoding().decodeBinary(binaryMessage.getMessageType(), toArray(pooled.getResource())));
} catch (DecodeException e) {
AnnotatedEndpoint.this.onError(session, e);
return;
}
//decoders
throw new RuntimeException("decoders are not implemented yet");
}
params.put(boolean.class, finalFragment);
session.getContainer().invokeEndpointMethod(executor, new Runnable() {
@Override
public void run() {
final Object result;
try {
result = binaryMessage.invoke(instance.getInstance(), params);
} catch (Exception e) {
AnnotatedEndpoint.this.onError(session, e);
return;
}
sendResult(result);
}
});
} finally {
pooled.free();
}
}
}
}