package com.github.ambry.network;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLServerSocket;
import javax.net.ssl.SSLSocket;


/**
 * A simple server that takes size delimited byte arrays and just echos them back to the sender.
 */
public class EchoServer extends Thread {
  public final int port;
  private final ServerSocket serverSocket;
  private final List<Thread> threads;
  private final List<Socket> sockets;
  private final List<Exception> exceptions;
  private final AtomicBoolean renegotiate = new AtomicBoolean();

  public EchoServer(int port)
      throws Exception {
    this(null, port);
  }

  public EchoServer(SSLFactory sslFactory, int port)
      throws Exception {
    this.port = port;
    if (sslFactory == null) {
      this.serverSocket = new ServerSocket(port);
    } else {
      SSLContext sslContext = sslFactory.createSSLContext();
      this.serverSocket = sslContext.getServerSocketFactory().createServerSocket(port);

      // enable mutual authentication
      ((SSLServerSocket) this.serverSocket).setNeedClientAuth(true);
    }
    this.threads = Collections.synchronizedList(new ArrayList<Thread>());
    this.sockets = Collections.synchronizedList(new ArrayList<Socket>());
    this.exceptions = Collections.synchronizedList(new ArrayList<Exception>());
  }

  /**
   * Test client to handle renegotiation from server
   * It only affects the next connection rather than all connections
   * After setting renegotiate to true, the next connection to the server will be renegotiated
   */
  public void renegotiate() {
    renegotiate.set(true);
  }

  @Override
  public void run() {
    try {
      while (true) {
        final Socket socket = serverSocket.accept();
        sockets.add(socket);
        Thread thread = new Thread() {
          @Override
          public void run() {
            try {
              DataInputStream input = new DataInputStream(socket.getInputStream());
              DataOutputStream output = new DataOutputStream(socket.getOutputStream());
              while (socket.isConnected() && !socket.isClosed()) {
                long size = input.readLong();
                if (renegotiate.compareAndSet(true, false)) {
                  ((SSLSocket) socket).startHandshake();
                }
                byte[] bytes = new byte[(int) size - 8];
                input.readFully(bytes);
                output.writeLong(size);
                output.write(bytes);
                output.flush();
              }
            } catch (IOException e) {
              // ignore, caused by client closed connection
            } finally {
              try {
                socket.close();
              } catch (IOException e) {
                exceptions.add(e);
              }
            }
          }
        };
        thread.start();
        threads.add(thread);
      }
    } catch (IOException e) {
      exceptions.add(e);
    }
  }

  public void closeConnections()
      throws IOException {
    for (Socket socket : sockets) {
      socket.close();
    }
  }

  public void close()
      throws IOException, InterruptedException {
    this.serverSocket.close();
    closeConnections();
    for (Thread t : threads) {
      t.join();
    }
    join();
  }

  public int getExceptionCount() {
    return exceptions.size();
  }
}

@@ -512,16 +512,62 @@
 sender.%0A
+ * Support both plaintext and SSL connections%0A
  */%0Apubl
@@ -857,16 +857,90 @@
 ean();%0A%0A
+  /**%0A   * Create an EchoServer that supports plaintext connections%0A   */%0A
   public
@@ -1005,32 +1005,100 @@
 ll, port);%0A  %7D%0A%0A
+  /**%0A   * Create an EchoServer that supports SSL connections%0A   */%0A
   public EchoSer

add more comments