package jxm.JWS;
import jxm.JWS.EX.InvalidDataException;
import jxm.JWS.FR.CloseFrame;
import jxm.JWS.FR.Framedata;
import jxm.JWS.HS.ClientHandshake;
import jxm.JWS.HS.Handshakedata;
import jxm.JWS.DF.Draft;
import jxm.JWS.HS.ServerHandshakeBuilder;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
/**
* <tt>WebSocketServer</tt> is an abstract class that only takes care of the
* HTTP handshake portion of WebSockets. It's up to a subclass to add
* functionality/purpose to the server.
*/
abstract class WebSocketServer extends WebSocketAdapter implements Runnable {
public static int DECODERS = Runtime.getRuntime().availableProcessors();
/**
* Holds the list of active WebSocket connections. "Active" means WebSocket
* handshake is complete and socket can be written to, or read from.
*/
private final Collection<WebSocket> connections;
/**
* The port number that this WebSocket server should listen on. Default is
* WebSocket.DEFAULT_PORT.
*/
private final InetSocketAddress address;
/**
* The socket channel for this WebSocket server.
*/
private ServerSocketChannel server;
/**
* The 'Selector' used to get event keys from the underlying socket.
*/
private Selector selector;
/**
* The Draft of the WebSocket protocol the Server is adhering to.
*/
private List<Draft> drafts;
private Thread selectorthread;
private volatile AtomicBoolean isclosed = new AtomicBoolean(false);
private List<WebSocketWorker> decoders;
private List<WebSocketImpl> iqueue;
private BlockingQueue<ByteBuffer> buffers;
private int queueinvokes = 0;
private AtomicInteger queuesize = new AtomicInteger(0);
private WebSocketServerFactory wsf = new DefaultWebSocketServerFactory();
/**
* Creates a WebSocketServer that will attempt to
* listen on port <var>WebSocket.DEFAULT_PORT</var>.
*
* @see #WebSocketServer(InetSocketAddress, int, List, Collection) more details here
*/
public WebSocketServer() throws UnknownHostException {
this(new InetSocketAddress(WebSocket.DEFAULT_PORT), DECODERS, null);
}
/**
* Creates a WebSocketServer that will attempt to bind/listen on the given <var>address</var>.
*
* @see #WebSocketServer(InetSocketAddress, int, List, Collection) more details here
*/
public WebSocketServer(InetSocketAddress address) {
this(address, DECODERS, null);
}
/**
* @see #WebSocketServer(InetSocketAddress, int, List, Collection) more details here
*/
public WebSocketServer(InetSocketAddress address, int decoders) {
this(address, decoders, null);
}
/**
* @see #WebSocketServer(InetSocketAddress, int, List, Collection) more details here
*/
public WebSocketServer(InetSocketAddress address, List<Draft> drafts) {
this(address, DECODERS, drafts);
}
/**
* @see #WebSocketServer(InetSocketAddress, int, List, Collection) more details here
*/
public WebSocketServer(InetSocketAddress address, int decodercount, List<Draft> drafts) {
this(address, decodercount, drafts, new HashSet<WebSocket>());
}
/**
* Creates a WebSocketServer that will attempt to bind/listen on the given <var>address</var>,
* and comply with <tt>Draft</tt> version <var>draft</var>.
*
* @param address The address (host:port) this server should listen on.
* @param decodercount The number of {@link WebSocketWorker}s that will be used to process the incoming network data. By default this will be <code>Runtime.getRuntime().availableProcessors()</code>
* @param drafts The versions of the WebSocket protocol that this server
* instance should comply to. Clients that use an other protocol version will be rejected.
* @param connectionscontainer Allows to specify a collection that will be used to store the websockets in. <br>
* If you plan to often iterate through the currently connected websockets you may want to use a collection that does not require synchronization like a {@link CopyOnWriteArraySet}. In that case make sure that you overload {@link #removeConnection(WebSocket)} and {@link #addConnection(WebSocket)}.<br>
* By default a {@link HashSet} will be used.
* @see #removeConnection(WebSocket) for more control over syncronized operation
* @see <a href="https://github.com/TooTallNate/Java-WebSocket/wiki/Drafts" > more about drafts
*/
public WebSocketServer(InetSocketAddress address, int decodercount, List<Draft> drafts, Collection<WebSocket> connectionscontainer) {
if (address == null || decodercount < 1 || connectionscontainer == null) {
throw new IllegalArgumentException("address and connectionscontainer must not be null and you need at least 1 decoder");
}
if (drafts == null)
this.drafts = Collections.emptyList();
else
this.drafts = drafts;
this.address = address;
this.connections = connectionscontainer;
iqueue = new LinkedList<WebSocketImpl>();
decoders = new ArrayList<WebSocketWorker>(decodercount);
buffers = new LinkedBlockingQueue<ByteBuffer>();
for (int i = 0; i < decodercount; i++) {
WebSocketWorker ex = new WebSocketWorker();
decoders.add(ex);
ex.start();
}
}
/**
* Starts the server selectorthread that binds to the currently set port number and
* listeners for WebSocket connection requests. Creates a fixed thread pool with the size {@link WebSocketServer#DECODERS}<br>
* May only be called once.
* <p/>
* Alternatively you can call {@link WebSocketServer#run()} directly.
*
* @throws IllegalStateException
*/
public void start() {
if (selectorthread != null)
throw new IllegalStateException(getClass().getName() + " can only be started once.");
new Thread(this).start();
;
}
/**
* Closes all connected clients sockets, then closes the underlying
* ServerSocketChannel, effectively killing the server socket selectorthread,
* freeing the port the server was bound to and stops all internal workerthreads.
* <p/>
* If this method is called before the server is started it will never start.
*
* @param timeout Specifies how many milliseconds the overall close handshaking may take altogether before the connections are closed without proper close handshaking.<br>
* @throws IOException When {@link ServerSocketChannel}.close throws an IOException
* @throws InterruptedException
*/
public void stop(int timeout) throws InterruptedException {
if (!isclosed.compareAndSet(false, true)) { // this also makes sure that no further connections will be added to this.connections
return;
}
List<WebSocket> socketsToClose = null;
// copy the connections in a list (prevent callback deadlocks)
synchronized (connections) {
socketsToClose = new ArrayList<WebSocket>(connections);
}
for (WebSocket ws : socketsToClose) {
ws.close(CloseFrame.GOING_AWAY);
}
synchronized (this) {
if (selectorthread != null) {
if (Thread.currentThread() != selectorthread) {
}
if (selectorthread != Thread.currentThread()) {
if (socketsToClose.size() > 0)
selectorthread.join(timeout);// isclosed will tell the selectorthread to go down after the last connection was closed
selectorthread.interrupt();// in case the selectorthread did not terminate in time we send the interrupt
selectorthread.join();
}
}
}
}
public void stop() throws IOException, InterruptedException {
stop(0);
}
/**
* Returns a WebSocket[] of currently connected clients.
* Its iterators will be failfast and its not judicious
* to modify it.
*
* @return The currently connected clients.
*/
public Collection<WebSocket> connections() {
return this.connections;
}
public InetSocketAddress getAddress() {
return this.address;
}
/**
* Gets the port number that this server listens on.
*
* @return The port number.
*/
public int getPort() {
int port = getAddress().getPort();
if (port == 0 && server != null) {
port = server.socket().getLocalPort();
}
return port;
}
public List<Draft> getDraft() {
return Collections.unmodifiableList(drafts);
}
// Runnable IMPLEMENTATION /////////////////////////////////////////////////
public void run() {
synchronized (this) {
if (selectorthread != null)
throw new IllegalStateException(getClass().getName() + " can only be started once.");
selectorthread = Thread.currentThread();
if (isclosed.get()) {
return;
}
}
selectorthread.setName("WebsocketSelector" + selectorthread.getId());
try {
server = ServerSocketChannel.open();
server.configureBlocking(false);
ServerSocket socket = server.socket();
socket.setReceiveBufferSize(WebSocketImpl.RCVBUF);
socket.bind(address);
selector = Selector.open();
server.register(selector, server.validOps());
} catch (IOException ex) {
handleFatal(null, ex);
return;
}
try {
while (!selectorthread.isInterrupted()) {
SelectionKey key = null;
WebSocketImpl conn = null;
try {
selector.select();
Set<SelectionKey> keys = selector.selectedKeys();
Iterator<SelectionKey> i = keys.iterator();
while (i.hasNext()) {
key = i.next();
if (!key.isValid()) {
// Object o = key.attachment();
continue;
}
if (key.isAcceptable()) {
if (!onConnect(key)) {
key.cancel();
continue;
}
SocketChannel channel = server.accept();
channel.configureBlocking(false);
WebSocketImpl w = wsf.createWebSocket(this, drafts, channel.socket());
w.key = channel.register(selector, SelectionKey.OP_READ, w);
w.channel = wsf.wrapChannel(channel, w.key);
i.remove();
allocateBuffers(w);
continue;
}
if (key.isReadable()) {
conn = (WebSocketImpl) key.attachment();
ByteBuffer buf = takeBuffer();
try {
if (SocketChannelIOHelper.read(buf, conn, conn.channel)) {
if (buf.hasRemaining()) {
conn.inQueue.put(buf);
queue(conn);
i.remove();
if (conn.channel instanceof WrappedByteChannel) {
if (((WrappedByteChannel) conn.channel).isNeedRead()) {
iqueue.add(conn);
}
}
} else
pushBuffer(buf);
} else {
pushBuffer(buf);
}
} catch (IOException e) {
pushBuffer(buf);
throw e;
}
}
if (key.isWritable()) {
conn = (WebSocketImpl) key.attachment();
if (SocketChannelIOHelper.batch(conn, conn.channel)) {
if (key.isValid())
key.interestOps(SelectionKey.OP_READ);
}
}
}
while (!iqueue.isEmpty()) {
conn = iqueue.remove(0);
WrappedByteChannel c = ((WrappedByteChannel) conn.channel);
ByteBuffer buf = takeBuffer();
try {
if (SocketChannelIOHelper.readMore(buf, conn, c))
iqueue.add(conn);
if (buf.hasRemaining()) {
conn.inQueue.put(buf);
queue(conn);
} else {
pushBuffer(buf);
}
} catch (IOException e) {
pushBuffer(buf);
throw e;
}
}
} catch (CancelledKeyException e) {
// an other thread may cancel the key
} catch (ClosedByInterruptException e) {
return; // do the same stuff as when InterruptedException is thrown
} catch (IOException ex) {
if (key != null)
key.cancel();
handleIOException(key, conn, ex);
} catch (InterruptedException e) {
return;// FIXME controlled shutdown (e.g. take care of buffermanagement)
}
}
} catch (RuntimeException e) {
// should hopefully never occur
handleFatal(null, e);
} finally {
if (decoders != null) {
for (WebSocketWorker w : decoders) {
w.interrupt();
}
}
if (server != null) {
try {
server.close();
} catch (IOException e) {
onError(null, e);
}
}
}
}
protected void allocateBuffers(WebSocket c) throws InterruptedException {
if (queuesize.get() >= 2 * decoders.size() + 1) {
return;
}
queuesize.incrementAndGet();
buffers.put(createBuffer());
}
protected void releaseBuffers(WebSocket c) throws InterruptedException {
// queuesize.decrementAndGet();
// takeBuffer();
}
public ByteBuffer createBuffer() {
return ByteBuffer.allocate(WebSocketImpl.RCVBUF);
}
private void queue(WebSocketImpl ws) throws InterruptedException {
if (ws.workerThread == null) {
ws.workerThread = decoders.get(queueinvokes % decoders.size());
queueinvokes++;
}
ws.workerThread.put(ws);
}
private ByteBuffer takeBuffer() throws InterruptedException {
return buffers.take();
}
private void pushBuffer(ByteBuffer buf) throws InterruptedException {
if (buffers.size() > queuesize.intValue())
return;
buffers.put(buf);
}
private void handleIOException(SelectionKey key, WebSocket conn, IOException ex) {
// onWebsocketError( conn, ex );// conn may be null here
if (conn != null) {
conn.closeConnection(CloseFrame.ABNORMAL_CLOSE, ex.getMessage());
} else if (key != null) {
SelectableChannel channel = key.channel();
if (channel != null && channel.isOpen()) { // this could be the case if the IOException ex is a SSLException
try {
channel.close();
} catch (IOException e) {
// there is nothing that must be done here
}
if (WebSocketImpl.DEBUG)
System.out.println("Connection closed because of" + ex);
}
}
}
private void handleFatal(WebSocket conn, Exception e) {
onError(conn, e);
try {
stop();
} catch (IOException e1) {
onError(null, e1);
} catch (InterruptedException e1) {
Thread.currentThread().interrupt();
onError(null, e1);
}
}
/**
* Gets the XML string that should be returned if a client requests a Flash
* security policy.
* <p/>
* The default implementation allows access from all remote domains, but
* only on the port that this WebSocketServer is listening on.
* <p/>
* This is specifically implemented for gitime's WebSocket client for Flash:
* http://github.com/gimite/web-socket-js
*
* @return An XML String that comforms to Flash's security policy. You MUST
* not include the null char at the end, it is appended automatically.
*/
protected String getFlashSecurityPolicy() {
return "<cross-domain-policy><allow-access-from domain=\"*\" to-ports=\"" + getPort() + "\" /></cross-domain-policy>";
}
@Override
public final void onWebsocketMessage(WebSocket conn, String message) {
onMessage(conn, message);
}
@Override
@Deprecated
public/*final*/void onWebsocketMessageFragment(WebSocket conn, Framedata frame) {// onFragment should be overloaded instead
onFragment(conn, frame);
}
@Override
public final void onWebsocketMessage(WebSocket conn, ByteBuffer blob) {
onMessage(conn, blob);
}
@Override
public final void onWebsocketOpen(WebSocket conn, Handshakedata handshake) {
if (addConnection(conn)) {
onOpen(conn, (ClientHandshake) handshake);
}
}
@Override
public final void onWebsocketClose(WebSocket conn, int code, String reason, boolean remote) {
selector.wakeup();
try {
if (removeConnection(conn)) {
onClose(conn, code, reason, remote);
}
} finally {
try {
releaseBuffers(conn);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
}
}
}
/**
* This method performs remove operations on the connection and therefore also gives control over whether the operation shall be synchronized
* <p/>
* {@link #WebSocketServer(InetSocketAddress, int, List, Collection)} allows to specify a collection which will be used to store current connections in.<br>
* Depending on the type on the connection, modifications of that collection may have to be synchronized.
*/
protected boolean removeConnection(WebSocket ws) {
boolean removed;
synchronized (connections) {
removed = this.connections.remove(ws);
assert (removed);
}
if (isclosed.get() && connections.size() == 0) {
selectorthread.interrupt();
}
return removed;
}
@Override
public ServerHandshakeBuilder onWebsocketHandshakeReceivedAsServer(WebSocket conn, Draft draft, ClientHandshake request) throws InvalidDataException {
return super.onWebsocketHandshakeReceivedAsServer(conn, draft, request);
}
/**
* @see #removeConnection(WebSocket)
*/
protected boolean addConnection(WebSocket ws) {
if (!isclosed.get()) {
synchronized (connections) {
boolean succ = this.connections.add(ws);
assert (succ);
return succ;
}
} else {
// This case will happen when a new connection gets ready while the server is already stopping.
ws.close(CloseFrame.GOING_AWAY);
return true;// for consistency sake we will make sure that both onOpen will be called
}
}
/**
* @param conn may be null if the error does not belong to a single connection
*/
@Override
public final void onWebsocketError(WebSocket conn, Exception ex) {
onError(conn, ex);
}
@Override
public final void onWriteDemand(WebSocket w) {
WebSocketImpl conn = (WebSocketImpl) w;
try {
conn.key.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
} catch (CancelledKeyException e) {
// the thread which cancels key is responsible for possible cleanup
conn.outQueue.clear();
}
selector.wakeup();
}
@Override
public void onWebsocketCloseInitiated(WebSocket conn, int code, String reason) {
onCloseInitiated(conn, code, reason);
}
@Override
public void onWebsocketClosing(WebSocket conn, int code, String reason, boolean remote) {
onClosing(conn, code, reason, remote);
}
public void onCloseInitiated(WebSocket conn, int code, String reason) {
}
public void onClosing(WebSocket conn, int code, String reason, boolean remote) {
}
public final void setWebSocketFactory(WebSocketServerFactory wsf) {
this.wsf = wsf;
}
public final WebSocketFactory getWebSocketFactory() {
return wsf;
}
/**
* Returns whether a new connection shall be accepted or not.<br>
* Therefore method is well suited to implement some kind of connection limitation.<br>
*
* @see {@link #onOpen(WebSocket, ClientHandshake)}, {@link #onWebsocketHandshakeReceivedAsServer(WebSocket, Draft, ClientHandshake)}
*/
protected boolean onConnect(SelectionKey key) {
return true;
}
private Socket getSocket(WebSocket conn) {
WebSocketImpl impl = (WebSocketImpl) conn;
return ((SocketChannel) impl.key.channel()).socket();
}
@Override
public InetSocketAddress getLocalSocketAddress(WebSocket conn) {
return (InetSocketAddress) getSocket(conn).getLocalSocketAddress();
}
@Override
public InetSocketAddress getRemoteSocketAddress(WebSocket conn) {
return (InetSocketAddress) getSocket(conn).getRemoteSocketAddress();
}
/**
* Called after an opening handshake has been performed and the given websocket is ready to be written on.
*/
public abstract void onOpen(WebSocket conn, ClientHandshake handshake);
/**
* Called after the websocket connection has been closed.
*
* @param code The codes can be looked up here: {@link CloseFrame}
* @param reason Additional information string
* @param remote Returns whether or not the closing of the connection was initiated by the remote host.
*/
public abstract void onClose(WebSocket conn, int code, String reason, boolean remote);
/**
* Callback for string messages received from the remote host
*
* @see #onMessage(WebSocket, ByteBuffer)
*/
public abstract void onMessage(WebSocket conn, String message);
/**
* Called when errors occurs. If an error causes the websocket connection to fail {@link #onClose(WebSocket, int, String, boolean)} will be called additionally.<br>
* This method will be called primarily because of IO or protocol errors.<br>
* If the given exception is an RuntimeException that probably means that you encountered a bug.<br>
*
* @param con Can be null if there error does not belong to one specific websocket. For example if the servers port could not be bound.
*/
public abstract void onError(WebSocket conn, Exception ex);
/**
* Callback for binary messages received from the remote host
*
* @see #onMessage(WebSocket, String)
*/
public void onMessage(WebSocket conn, ByteBuffer message) {
}
/**
* @see WebSocket#sendFragmentedFrame(jxm._.JWS.FR.Framedata.Opcode, ByteBuffer, boolean)
*/
public void onFragment(WebSocket conn, Framedata fragment) {
}
public class WebSocketWorker extends Thread {
private BlockingQueue<WebSocketImpl> iqueue;
public WebSocketWorker() {
iqueue = new LinkedBlockingQueue<WebSocketImpl>();
setName("WebSocketWorker-" + getId());
setUncaughtExceptionHandler(new UncaughtExceptionHandler() {
@Override
public void uncaughtException(Thread t, Throwable e) {
getDefaultUncaughtExceptionHandler().uncaughtException(t, e);
}
});
}
public void put(WebSocketImpl ws) throws InterruptedException {
iqueue.put(ws);
}
@Override
public void run() {
WebSocketImpl ws = null;
try {
while (true) {
ByteBuffer buf = null;
ws = iqueue.take();
buf = ws.inQueue.poll();
assert (buf != null);
try {
ws.decode(buf);
} finally {
pushBuffer(buf);
}
}
} catch (InterruptedException e) {
} catch (RuntimeException e) {
handleFatal(ws, e);
}
}
}
public interface WebSocketServerFactory extends WebSocketFactory {
@Override
public WebSocketImpl createWebSocket(WebSocketAdapter a, Draft d, Socket s);
public WebSocketImpl createWebSocket(WebSocketAdapter a, List<Draft> drafts, Socket s);
/**
* Allows to wrap the Socketchannel( key.channel() ) to insert a protocol layer( like ssl or proxy authentication) beyond the ws layer.
*
* @param key a SelectionKey of an open SocketChannel.
* @return The channel on which the read and write operations will be performed.<br>
*/
public ByteChannel wrapChannel(SocketChannel channel, SelectionKey key) throws IOException;
}
}