Package org.java_websocket.server

Source Code of org.java_websocket.server.WebSocketServer$WebSocketWorker

package org.java_websocket.server;

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.ByteChannel;
import java.nio.channels.CancelledKeyException;
import java.nio.channels.ClosedByInterruptException;
import java.nio.channels.SelectableChannel;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;
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;

import org.java_websocket.SocketChannelIOHelper;
import org.java_websocket.WebSocket;
import org.java_websocket.WebSocketAdapter;
import org.java_websocket.WebSocketFactory;
import org.java_websocket.WebSocketImpl;
import org.java_websocket.WrappedByteChannel;
import org.java_websocket.drafts.Draft;
import org.java_websocket.exceptions.InvalidDataException;
import org.java_websocket.framing.CloseFrame;
import org.java_websocket.framing.Framedata;
import org.java_websocket.handshake.ClientHandshake;
import org.java_websocket.handshake.Handshakedata;
import org.java_websocket.handshake.ServerHandshakeBuilder;

/**
* <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.
*
*/
public 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.
   *
   * 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.
   *
   * 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.
   *
   * The default implementation allows access from all remote domains, but
   * only on the port that this WebSocketServer is listening on.
   *
   * 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(org.java_websocket.framing.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;
  }
}
TOP

Related Classes of org.java_websocket.server.WebSocketServer$WebSocketWorker

TOP
Copyright © 2018 www.massapi.com. All rights reserved.
All source code are property of their respective owners. Java is a trademark of Sun Microsystems, Inc and owned by ORACLE Inc. Contact coftware#gmail.com.