001/*
002 * Copyright 2002-2018 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.springframework.web.reactive.socket.adapter;
018
019import java.nio.ByteBuffer;
020import java.nio.charset.StandardCharsets;
021
022import io.undertow.websockets.WebSocketConnectionCallback;
023import io.undertow.websockets.core.AbstractReceiveListener;
024import io.undertow.websockets.core.BufferedBinaryMessage;
025import io.undertow.websockets.core.BufferedTextMessage;
026import io.undertow.websockets.core.CloseMessage;
027import io.undertow.websockets.core.WebSocketChannel;
028
029import org.springframework.core.io.buffer.DataBuffer;
030import org.springframework.util.Assert;
031import org.springframework.web.reactive.socket.CloseStatus;
032import org.springframework.web.reactive.socket.WebSocketHandler;
033import org.springframework.web.reactive.socket.WebSocketMessage;
034import org.springframework.web.reactive.socket.WebSocketMessage.Type;
035
036/**
037 * Undertow {@link WebSocketConnectionCallback} implementation that adapts and
038 * delegates to a Spring {@link WebSocketHandler}.
039 *
040 * @author Violeta Georgieva
041 * @author Rossen Stoyanchev
042 * @since 5.0
043 */
044public class UndertowWebSocketHandlerAdapter extends AbstractReceiveListener {
045
046        private final UndertowWebSocketSession session;
047
048
049        public UndertowWebSocketHandlerAdapter(UndertowWebSocketSession session) {
050                Assert.notNull(session, "UndertowWebSocketSession is required");
051                this.session = session;
052        }
053
054
055        @Override
056        protected void onFullTextMessage(WebSocketChannel channel, BufferedTextMessage message) {
057                this.session.handleMessage(Type.TEXT, toMessage(Type.TEXT, message.getData()));
058        }
059
060        @Override
061        protected void onFullBinaryMessage(WebSocketChannel channel, BufferedBinaryMessage message) {
062                this.session.handleMessage(Type.BINARY, toMessage(Type.BINARY, message.getData().getResource()));
063                message.getData().free();
064        }
065
066        @Override
067        protected void onFullPongMessage(WebSocketChannel channel, BufferedBinaryMessage message) {
068                this.session.handleMessage(Type.PONG, toMessage(Type.PONG, message.getData().getResource()));
069                message.getData().free();
070        }
071
072        @Override
073        protected void onFullCloseMessage(WebSocketChannel channel, BufferedBinaryMessage message) {
074                CloseMessage closeMessage = new CloseMessage(message.getData().getResource());
075                this.session.handleClose(new CloseStatus(closeMessage.getCode(), closeMessage.getReason()));
076                message.getData().free();
077        }
078
079        @Override
080        protected void onError(WebSocketChannel channel, Throwable error) {
081                this.session.handleError(error);
082        }
083
084        private <T> WebSocketMessage toMessage(Type type, T message) {
085                if (Type.TEXT.equals(type)) {
086                        byte[] bytes = ((String) message).getBytes(StandardCharsets.UTF_8);
087                        return new WebSocketMessage(Type.TEXT, this.session.bufferFactory().wrap(bytes));
088                }
089                else if (Type.BINARY.equals(type)) {
090                        DataBuffer buffer = this.session.bufferFactory().allocateBuffer().write((ByteBuffer[]) message);
091                        return new WebSocketMessage(Type.BINARY, buffer);
092                }
093                else if (Type.PONG.equals(type)) {
094                        DataBuffer buffer = this.session.bufferFactory().allocateBuffer().write((ByteBuffer[]) message);
095                        return new WebSocketMessage(Type.PONG, buffer);
096                }
097                else {
098                        throw new IllegalArgumentException("Unexpected message type: " + message);
099                }
100        }
101
102}