'use strict';

const {Services} = ChromeUtils.import("resource://gre/modules/Services.jsm");
// Bug 788960 and later bug 1329245 have taught us that attempting to connect to
// a port that is not listening or is no longer listening fails to consistently
// result in the error (or any) event we expect on Darwin/OSX/"OS X".
const isOSX = (Services.appinfo.OS === "Darwin");
const testConnectingToNonListeningPort = !isOSX;

const SERVER_BACKLOG = -1;

const SOCKET_EVENTS = ['open', 'data', 'drain', 'error', 'close'];

function concatUint8Arrays(a, b) {
  let newArr = new Uint8Array(a.length + b.length);
  newArr.set(a, 0);
  newArr.set(b, a.length);
  return newArr;
}

function assertUint8ArraysEqual(a, b, comparingWhat) {
  if (a.length !== b.length) {
    ok(false, comparingWhat + ' arrays do not have the same length; ' +
       a.length + ' versus ' + b.length);
    return;
  }
  for (let i = 0; i < a.length; i++) {
    if (a[i] !== b[i]) {
      ok(false, comparingWhat + ' arrays differ at index ' + i +
         a[i] + ' versus ' + b[i]);
      return;
    }
  }
  ok(true, comparingWhat + ' arrays were equivalent.');
}

/**
 * Helper method to add event listeners to a socket and provide two Promise-returning
 * helpers (see below for docs on them).  This *must* be called during the turn of
 * the event loop where TCPSocket's constructor is called or the onconnect method is being
 * invoked.
 */
function listenForEventsOnSocket(socket, socketType) {
  let wantDataLength = null;
  let wantDataAndClose = false;
  let pendingResolve = null;
  let receivedEvents = [];
  let receivedData = null;
  let handleGenericEvent = function(event) {
    dump('(' + socketType + ' event: ' + event.type + ')\n');
    if (pendingResolve && wantDataLength === null) {
      pendingResolve(event);
      pendingResolve = null;
    } else {
      receivedEvents.push(event);
    }
  };

  socket.onopen = handleGenericEvent;
  socket.ondrain = handleGenericEvent;
  socket.onerror = handleGenericEvent;
  socket.onclose = function(event) {
    if (!wantDataAndClose) {
      handleGenericEvent(event);
    } else if (pendingResolve) {
      dump('(' + socketType + ' event: close)\n');
      pendingResolve(receivedData);
      pendingResolve = null;
      wantDataAndClose = false;
    }
  }
  socket.ondata = function(event) {
    dump('(' + socketType + ' event: ' + event.type + ' length: ' +
         event.data.byteLength + ')\n');
    ok(socketCompartmentInstanceOfArrayBuffer(event.data),
       'payload is ArrayBuffer');
    var arr = new Uint8Array(event.data);
    if (receivedData === null) {
      receivedData = arr;
    } else {
      receivedData = concatUint8Arrays(receivedData, arr);
    }
    if (wantDataLength !== null &&
        receivedData.length >= wantDataLength) {
      pendingResolve(receivedData);
      pendingResolve = null;
      receivedData = null;
      wantDataLength = null;
    }
  };


  return {
    /**
     * Return a Promise that will be resolved with the next (non-data) event
     * received by the socket.  If there are queued events, the Promise will
     * be immediately resolved (but you won't see that until a future turn of
     * the event loop).
     */
    waitForEvent: function() {
      if (pendingResolve) {
        throw new Error('only one wait allowed at a time.');
      }

      if (receivedEvents.length) {
        return Promise.resolve(receivedEvents.shift());
      }

      dump('(' + socketType + ' waiting for event)\n');
      return new Promise(function(resolve, reject) {
        pendingResolve = resolve;
      });
    },
    /**
     * Return a Promise that will be resolved with a Uint8Array of at least the
     * given length.  We buffer / accumulate received data until we have enough
     * data.  Data is buffered even before you call this method, so be sure to
     * explicitly wait for any and all data sent by the other side.
     */
    waitForDataWithAtLeastLength: function(length) {
      if (pendingResolve) {
        throw new Error('only one wait allowed at a time.');
      }
      if (receivedData && receivedData.length >= length) {
        let promise = Promise.resolve(receivedData);
        receivedData = null;
        return promise;
      }
      dump('(' + socketType + ' waiting for ' + length + ' bytes)\n');
      return new Promise(function(resolve, reject) {
        pendingResolve = resolve;
        wantDataLength = length;
      });
    },
    waitForAnyDataAndClose: function() {
      if (pendingResolve) {
        throw new Error('only one wait allowed at a time.');
      }

      return new Promise(function(resolve, reject) {
        pendingResolve = resolve;
        // we may receive no data before getting close, in which case we want to
        // return an empty array
        receivedData = new Uint8Array();
        wantDataAndClose = true;
      });
    }
  };
}

/**
 * Return a promise that is resolved when the server receives a connection.  The
 * promise is resolved with { socket, queue } where `queue` is the result of
 * calling listenForEventsOnSocket(socket).  This must be done because we need
 * to add the event listener during the connection.
 */
function waitForConnection(listeningServer) {
  return new Promise(function(resolve, reject) {
    // Because of the event model of sockets, we can't use the
    // listenForEventsOnSocket mechanism; we need to hook up listeners during
    // the connect event.
    listeningServer.onconnect = function(event) {
      // Clobber the listener to get upset if it receives any more connections
      // after this.
      listeningServer.onconnect = function() {
        ok(false, 'Received a connection when not expecting one.');
      };
      ok(true, 'Listening server accepted socket');
      resolve({
        socket: event.socket,
        queue: listenForEventsOnSocket(event.socket, 'server')
      });
    };
  });
}

function defer() {
  var deferred = {};
  deferred.promise = new Promise(function(resolve, reject) {
    deferred.resolve = resolve;
    deferred.reject = reject;
  });
  return deferred;
}


async function test_basics() {
  // See bug 903830; in e10s mode we never get to find out the localPort if we
  // let it pick a free port by choosing 0.  This is the same port the xpcshell
  // test was using.
  let serverPort = 8085;

  // - Start up a listening socket.
  let listeningServer = createServer(serverPort,
                                     { binaryType: 'arraybuffer' },
                                     SERVER_BACKLOG);

  let connectedPromise = waitForConnection(listeningServer);

  // -- Open a connection to the server
  let clientSocket = createSocket('127.0.0.1', serverPort,
                                  { binaryType: 'arraybuffer' });
  let clientQueue = listenForEventsOnSocket(clientSocket, 'client');

  // (the client connects)
  is((await clientQueue.waitForEvent()).type, 'open', 'got open event');
  is(clientSocket.readyState, 'open', 'client readyState is open');

  // (the server connected)
  let { socket: serverSocket, queue: serverQueue } = await connectedPromise;
  is(serverSocket.readyState, 'open', 'server readyState is open');

  // -- Simple send / receive
  // - Send data from client to server
  // (But not so much we cross the drain threshold.)
  let smallUint8Array = new Uint8Array(256);
  for (let i = 0; i < smallUint8Array.length; i++) {
    smallUint8Array[i] = i;
  }
  is(clientSocket.send(smallUint8Array.buffer, 0, smallUint8Array.length), true,
     'Client sending less than 64k, buffer should not be full.');

  let serverReceived = await serverQueue.waitForDataWithAtLeastLength(256);
  assertUint8ArraysEqual(serverReceived, smallUint8Array,
                         'Server received/client sent');

  // - Send data from server to client
  // (But not so much we cross the drain threshold.)
  is(serverSocket.send(smallUint8Array.buffer, 0, smallUint8Array.length), true,
     'Server sending less than 64k, buffer should not be full.');

  let clientReceived = await clientQueue.waitForDataWithAtLeastLength(256);
  assertUint8ArraysEqual(clientReceived, smallUint8Array,
                         'Client received/server sent');

  // -- Perform sending multiple times with different buffer slices
  // - Send data from client to server
  // (But not so much we cross the drain threshold.)
  is(clientSocket.send(smallUint8Array.buffer, 0, 7),
     true, 'Client sending less than 64k, buffer should not be full.');
  is(clientSocket.send(smallUint8Array.buffer, 7, smallUint8Array.length - 7),
     true, 'Client sending less than 64k, buffer should not be full.');

  serverReceived = await serverQueue.waitForDataWithAtLeastLength(256);
  assertUint8ArraysEqual(serverReceived, smallUint8Array,
                         'Server received/client sent');

  // - Send data from server to client
  // (But not so much we cross the drain threshold.)
  is(serverSocket.send(smallUint8Array.buffer, 0, 7),
     true, 'Server sending less than 64k, buffer should not be full.');
  is(serverSocket.send(smallUint8Array.buffer, 7, smallUint8Array.length - 7),
     true, 'Server sending less than 64k, buffer should not be full.');

  clientReceived = await clientQueue.waitForDataWithAtLeastLength(256);
  assertUint8ArraysEqual(clientReceived, smallUint8Array,
                         'Client received/server sent');


  // -- Send "big" data in both directions
  // (Enough to cross the buffering/drain threshold; 64KiB)
  let bigUint8Array = new Uint8Array(65536 + 3);
  for (let i = 0; i < bigUint8Array.length; i++) {
    bigUint8Array[i] = i % 256;
  }
  // Do this twice so we have confidence that the 'drain' event machinery
  // doesn't break after the first use.
  for (let iSend = 0; iSend < 2; iSend++) {
    // - Send "big" data from the client to the server
    is(clientSocket.send(bigUint8Array.buffer, 0, bigUint8Array.length), false,
       'Client sending more than 64k should result in the buffer being full.');
    is((await clientQueue.waitForEvent()).type, 'drain',
       'The drain event should fire after a large send that indicated full.');

    serverReceived = await serverQueue.waitForDataWithAtLeastLength(
      bigUint8Array.length);
    assertUint8ArraysEqual(serverReceived, bigUint8Array,
                           'server received/client sent');

    // - Send "big" data from the server to the client
    is(serverSocket.send(bigUint8Array.buffer, 0, bigUint8Array.length), false,
       'Server sending more than 64k should result in the buffer being full.');
    is((await serverQueue.waitForEvent()).type, 'drain',
       'The drain event should fire after a large send that indicated full.');

    clientReceived = await clientQueue.waitForDataWithAtLeastLength(
      bigUint8Array.length);
    assertUint8ArraysEqual(clientReceived, bigUint8Array,
                           'client received/server sent');
  }

  // -- Server closes the connection
  serverSocket.close();
  is(serverSocket.readyState, 'closing',
     'readyState should be closing immediately after calling close');

  is((await clientQueue.waitForEvent()).type, 'close',
     'The client should get a close event when the server closes.');
  is(clientSocket.readyState, 'closed',
     'client readyState should be closed after close event');
  is((await serverQueue.waitForEvent()).type, 'close',
     'The server should get a close event when it closes itself.');
  is(serverSocket.readyState, 'closed',
     'server readyState should be closed after close event');

  // -- Re-establish connection
  connectedPromise = waitForConnection(listeningServer);
  clientSocket = createSocket('127.0.0.1', serverPort,
                              { binaryType: 'arraybuffer' });
  clientQueue = listenForEventsOnSocket(clientSocket, 'client');
  is((await clientQueue.waitForEvent()).type, 'open', 'got open event');

  let connectedResult = await connectedPromise;
  // destructuring assignment is not yet ES6 compliant, must manually unpack
  serverSocket = connectedResult.socket;
  serverQueue = connectedResult.queue;

  // -- Client closes the connection
  clientSocket.close();
  is(clientSocket.readyState, 'closing',
     'client readyState should be losing immediately after calling close');

  is((await clientQueue.waitForEvent()).type, 'close',
     'The client should get a close event when it closes itself.');
  is(clientSocket.readyState, 'closed',
     'client readyState should be closed after the close event is received');
  is((await serverQueue.waitForEvent()).type, 'close',
     'The server should get a close event when the client closes.');
  is(serverSocket.readyState, 'closed',
     'server readyState should be closed after the close event is received');


  // -- Re-establish connection
  connectedPromise = waitForConnection(listeningServer);
  clientSocket = createSocket('127.0.0.1', serverPort,
                              { binaryType: 'arraybuffer' });
  clientQueue = listenForEventsOnSocket(clientSocket, 'client');
  is((await clientQueue.waitForEvent()).type, 'open', 'got open event');

  connectedResult = await connectedPromise;
  // destructuring assignment is not yet ES6 compliant, must manually unpack
  serverSocket = connectedResult.socket;
  serverQueue = connectedResult.queue;

  // -- Call close after enqueueing a lot of data, make sure it goes through.
  // We'll have the client send and close.
  is(clientSocket.send(bigUint8Array.buffer, 0, bigUint8Array.length), false,
     'Client sending more than 64k should result in the buffer being full.');
  clientSocket.close();
  // The drain will still fire
  is((await clientQueue.waitForEvent()).type, 'drain',
     'The drain event should fire after a large send that returned true.');
  // Then we'll get a close
  is((await clientQueue.waitForEvent()).type, 'close',
     'The close event should fire after the drain event.');

  // The server will get its data
  serverReceived = await serverQueue.waitForDataWithAtLeastLength(
    bigUint8Array.length);
  assertUint8ArraysEqual(serverReceived, bigUint8Array,
                         'server received/client sent');
  // And a close.
  is((await serverQueue.waitForEvent()).type, 'close',
     'The drain event should fire after a large send that returned true.');


  // -- Re-establish connection
  connectedPromise = waitForConnection(listeningServer);
  clientSocket = createSocket('127.0.0.1', serverPort,
                              { binaryType: 'string' });
  clientQueue = listenForEventsOnSocket(clientSocket, 'client');
  is((await clientQueue.waitForEvent()).type, 'open', 'got open event');

  connectedResult = await connectedPromise;
  // destructuring assignment is not yet ES6 compliant, must manually unpack
  serverSocket = connectedResult.socket;
  serverQueue = connectedResult.queue;

  // -- Attempt to send non-string data.
  // Restore the original behavior by replacing toString with
  // Object.prototype.toString. (bug 1121938)
  bigUint8Array.toString = Object.prototype.toString;
  is(clientSocket.send(bigUint8Array), true,
     'Client sending a large non-string should only send a small string.');
  clientSocket.close();
  // The server will get its data
  serverReceived = await serverQueue.waitForDataWithAtLeastLength(
    bigUint8Array.toString().length);
  // Then we'll get a close
  is((await clientQueue.waitForEvent()).type, 'close',
     'The close event should fire after the drain event.');

  // -- Re-establish connection (Test for Close Immediately)
  connectedPromise = waitForConnection(listeningServer);
  clientSocket = createSocket('127.0.0.1', serverPort,
                               { binaryType: 'arraybuffer' });
  clientQueue = listenForEventsOnSocket(clientSocket, 'client');
  is((await clientQueue.waitForEvent()).type, 'open', 'got open event');

  connectedResult = await connectedPromise;
  // destructuring assignment is not yet ES6 compliant, must manually unpack
  serverSocket = connectedResult.socket;
  serverQueue = connectedResult.queue;

  // -- Attempt to send two non-string data.
  is(clientSocket.send(bigUint8Array.buffer, 0, bigUint8Array.length), false,
     'Server sending more than 64k should result in the buffer being full.');
  is(clientSocket.send(bigUint8Array.buffer, 0, bigUint8Array.length), false,
     'Server sending more than 64k should result in the buffer being full.');
  clientSocket.closeImmediately();

  serverReceived = await serverQueue.waitForAnyDataAndClose();

  is(serverReceived.length < (2 * bigUint8Array.length), true, 'Received array length less than sent array length');

  // -- Close the listening server (and try to connect)
  // We want to verify that the server actually closes / stops listening when
  // we tell it to.
  listeningServer.close();

  // (We don't run this check on OS X where it's flakey; see definition up top.)
  if (testConnectingToNonListeningPort) {
    // - try and connect, get an error
    clientSocket = createSocket('127.0.0.1', serverPort,
                                { binaryType: 'arraybuffer' });
    clientQueue = listenForEventsOnSocket(clientSocket, 'client');
    is((await clientQueue.waitForEvent()).type, 'error', 'fail to connect');
    is(clientSocket.readyState, 'closed',
       'client readyState should be closed after the failure to connect');
  }
}

add_task(test_basics);
