!7

// TODO: we are iterating over all channels and all sockest

// 100 times per second - this seems excessive

static int updateInterval = 10;

static long timeout = 30000+10000; // idle time + latency

static int port = 6000;

static int maxLineLength = 10*1024*1024; // change #1020500 too if you change this

import java.nio.*;

import java.nio.ByteBuffer;

import java.nio.channels.*;

static ConnectedInstances cI;

static ReliableSingleThread rst = new(r update);

static int theNumber;

static Map<SocketChannel, State> dataMapper = synchroMap();

static ServerSocketChannel serverChannel;

static Selector selector;

static Map<S, Channel> channels = synchroMap();

static Lock channelsLock = lock();

static java.util.Timer mainTimer;

sbool debug;

// state of a client connection

srecord noeq State(SocketChannel socket) {

  S socketID = aGlobalID(), computerID;

  int countSeen = -1;

  new LineBuffer lineBuffer;

  long lastMessage = sysNow();

  Map<S, Channel> subbedChannels = synchroLinkedHashMap();

  Lock lock = lock();

  // TODO: clear eventually to avoid stalling attacks

  L<byte[]> toSend = synchroLinkedList();

}

// a virtual "channel" that clients can "subscribe" to

// (what that means differs between channels)

// Each channel is identified by a unique name

// A is a custom data type for each subscriber

sclass Channel<A> {

  S name;

  Set<SocketChannel> sockets = synchroSet();

  Map<SocketChannel, A> dataMapper = synchroMap();

  Lock lock = lock();

  event subCountChanged;

  // add a subscriber

  void addSocket(SocketChannel socket) {

    lock lock;

    sockets.add(socket);

    subCountChanged();

    initialMessage(socket);

  }

  // remove a subscriber

  void removeSocket(SocketChannel socket) {

    lock lock;

    sockets.remove(socket);

    subCountChanged();

  }

  int subCount() { ret l(sockets); }

  void sendToSubscribers(S line) {  

    for (SocketChannel socket : cloneList(sockets))

      sendLineToSocket(socket, line);

  }

  void update() {}

  void initialMessage(SocketChannel socket) {}

  void handleLine(S line, SocketChannel socket) {}

}

// channel listing the currently logged in computer IDs

// TODO: optimize the protocol to send diffs only

sclass ComputerIDsChannel extends Channel<O> {

  L<S> value;

  void initialMessage(SocketChannel socket) {

    sendLineToSocket(socket, stateMsg());

  }

  void update() {

    if (set_trueIfChanged(this, value := asList(cI.set())))

      sendToSubscribers(stateMsg());

  }

  S stateMsg() {

    ret "ComputerIDs = " + struct(value);

  }

}

// channel listing the current count of logged in computers

sclass ComputerCountChannel extends Channel<O> {

  int value;

  void initialMessage(SocketChannel socket) {

    sendLineToSocket(socket, stateMsg());

  }

  void update() {

    if (set_trueIfChanged(this, value := cI.value()))

      sendToSubscribers(stateMsg());

  }

  S stateMsg() {

    ret "ComputerCount = " + value;

  }

}

// a public chat for broadcast messages between subscribers

sclass PublicChatChannel extends Channel<O> {

  L<S> msgs = synchroLinkedList();

  void handleLine(S line, SocketChannel socket) {

    postMessage(unquote(line));

  }

  void postMessage(S msg) {

    msgs.add(msg);

  }

  void update {

    S msg;

    while ((msg = syncPopFirst(msgs)) != null) {

      S quoted = quote(msg);

      //print("Chat send: " + shorten(quoted, 20));

      quoted = name + ":" + quoted;

      for (SocketChannel s : sockets)

        sendLineToSocket(s, quoted);

    }

  }

}

// a chat for DMs (messages from one computer to another)

// (is it private though? probably not as we don't check if

// the computer IDs sent to the server are correct)

sclass PrivateChatChannel extends Channel<O> {

  S prefix = "";

  new MultiMap<S, SocketChannel> clients; // key: computer ID

  void initialMessage(SocketChannel socket) {

    clients.put(computerID(socket), socket);

  }

  void removeSocket(SocketChannel socket) {

    super.removeSocket(socket);

    clients.remove(computerID(socket), socket);

  }

  void handleLine(S line, SocketChannel socket) {

    pcall {

      int i = indexOf(line, ' ');

      S computerID = takeFirst(line, i);

      while licensed {

        SocketChannel socket2 = clients.getFirst(computerID);

        if (socket2 == null) break;

        try {

          S sender = computerID(socket);

          if (l(line) >= 10000)

            print("Forwarding line from " + sender + " to " + computerID + ": " + l(line));

          sendLineToSocket(socket2, prefix + "privately from " + sender + ": " + substring(line, i+1));

          break;

        } catch e {

          closeSocket(socket2);

          clients.remove(computerID, socket2);

        }

      }

    }

  }

}

// channel distributing snippet updates (edits, compiles etc)

sclass SnippetUpdatesChannel extends Channel<O> {

  ManualTailFile tail;

  *() {

    tail = manualTailFile_newOnly(programFile(#1019175, "changes.log"),

      vfAppendToLineBuffer(unquotingLineBuffer(voidfunc(S line) { sendLine(line) })));

  }

  void update {

    tail.check();

  }

  void sendLine(S line) {

    print("SnippetUpdatesChannel send: " + line);

    line = assertNoLineBreaks(line);

    print("Sockets: " + l(sockets));

    //S quoted = quote(line);

    for (SocketChannel s : sockets)

      sendLineToSocket(s, "snippetUpdates:" + line);

  }

}

// channel broadcast its current subscriber count

sclass PresenceCountingChannel extends Channel<Void> {

  *() {

    onSubCountChanged(-> sendToSubscribers(stateMsg()));

  }

  void initialMessage(SocketChannel socket) {

    sendLineToSocket(socket, stateMsg());

  }

  S stateMsg() {

    ret name + ":" + subCount();

  }

}

static State state(SocketChannel socket) {

  ret dataMapper.get(socket);

}

sS computerID(SocketChannel socket) {

  State state = state(socket);

  ret state == null ? null : state.computerID;

}

p {

  cI = new ConnectedInstances;

  cI.setAfterglowInterval(1000);

  cI.connected2.onChange(rst);

  cI.start();

  openServerSocket();

  // restart socket every hour (and once now for testing - XX)

  /*doEvery(10.0, 60*60.0, r {

    pcall { cleanMeUp_socket(); }

    openServerSocket();

  });*/

}

svoid openServerSocket ctex {

  print("Opening server socket");

  selector = Selector.open();

  serverChannel = ServerSocketChannel.open();

  serverChannel.configureBlocking(false);

  InetSocketAddress sockAddr = new InetSocketAddress(port);

  serverChannel.socket().bind(sockAddr);

  print("Server socket bound");

  serverChannel.register(selector, SelectionKey.OP_ACCEPT);

  thread "Server Accept Loop" { acceptLoop(); }

  mainTimer = doEvery(updateInterval, r {

    for (SocketChannel socket : keysList(dataMapper)) {

      State state = state(socket);

      if (state != null && sysTime() >= state.lastMessage+timeout) pcall {

        closeSocket(socket);

      }

    }

    for (Channel channel : values(channels)) pcall {

      lock channel.lock;

      channel.update();

    }

  });

}

svoid cleanMeUp_socket() ctex {

  if (serverChannel != null)

    serverChannel.close();

  for (SocketChannel socket : keys(getAndClearMap(dataMapper))) pcall {

    socket.close();

  }

}

svoid acceptLoop ctex {

  while licensed {

    // wait for events

    selector.select();

    Iterator<SelectionKey> keys = selector.selectedKeys().iterator();

    while (keys.hasNext()) {

      SelectionKey key = keys.next();

      keys.remove();

      debug "have key";

      if (!key.isValid())

        continue with debug "key not valid";

      pcall {

        if (key.isAcceptable()) {

          debug "Key is acceptable";

          acceptOnSocket(key);

        } else if (key.isReadable()) {

          debug "Key is readable";

          readFromSocket(key);

        } else if (key.isWritable()) {

          debug "Key is writable";

          writeToSocket(key);

        } else

          debug "Unknown key type";

      }

    }

  }

}

//accept a connection made to this channel's socket

static void acceptOnSocket(SelectionKey key) throws IOException {

  ServerSocketChannel serverChannel = cast key.channel();

  SocketChannel channel = serverChannel.accept();

  channel.configureBlocking(false);

  Socket actualSocket = channel.socket();

  SocketAddress remoteAddr = actualSocket.getRemoteSocketAddress();

  print("Connected to: " + remoteAddr);

  // register channel with selector for further IO

  dataMapper.put(channel, new State(channel));

  channel.register(selector, SelectionKey.OP_READ);

}

svoid onSocketClose(SocketChannel socket) {

  Socket actualSocket = socket.socket();

  SocketAddress remoteAddr = actualSocket.getRemoteSocketAddress();

  print("Connection closed by client: " + remoteAddr);

  State state = state(socket);

  if (state != null && state.computerID != null)

    cI.lostConnection(state.computerID);

  dataMapper.remove(socket);

  removeFromAllChannels(socket);

}

svoid closeSocket(SocketChannel socket) {

  pcall { socket.close(); }

  pcall { onSocketClose(socket); }

}

static void readFromSocket(SelectionKey key) throws IOException {

  final SocketChannel socket = cast key.channel();

  ByteBuffer buffer = ByteBuffer.allocate(1024);

  int numRead = socket.read(buffer);

  if (numRead == -1) {

    closeSocket(socket);

    key.cancel();

    ret;

  }

  S s = new String(buffer.array(), 0, numRead);

  //print("Got data from client: " + s);

  final State state = state(socket);

  if (state != null) {

    state.lastMessage = sysNow();

    state.lineBuffer.append(s, voidfunc(S line) {

      //if (nempty(line)) print("Got line: " + line);

      pcall { handleLine(line, socket, state); }

    });

    if (state.lineBuffer.size() > maxLineLength) {

      print("Bad client (line too long");

      closeSocket(socket);

    }

  }

}

static void writeToSocket(SelectionKey key) throws IOException {

  final SocketChannel socket = cast key.channel();

  socket.register(selector, SelectionKey.OP_READ);

  State state = state(socket);

  lock state.lock;

  if (empty(state.toSend)) ret;

  byte[] data = popFirst(state.toSend);

  print("Sending postponed data (" + n2(l(data), "byte") + ") to " + state.computerID);

  sendBytesToSocket(socket, data);

}

svoid sendBytesToSocket(SocketChannel socket, byte[] bytes) throws IOException {

  int length = l(bytes);

  State state = state(socket);

  lock state.lock;

  if (nempty(state.toSend)) {

    print("Postponing send to " + state.computerID + " because data already queued");

    ret with state.toSend.add(bytes);

  }

  ByteBuffer buf = ByteBuffer.wrap(bytes);

  int n = socket.write(buf);

  if (n < length) {

    byte[] remaining = subByteArray(bytes, n);

    state.toSend.add(remaining);

    print("Postponing sending " + n2(l(remaining), "byte") + " to " + state.computerID);

    socket.register(selector, SelectionKey.OP_READ | SelectionKey.OP_WRITE);

  }

}

svoid handleLine(S line, SocketChannel socket, State state) {

  new Matches m;

  if (empty(line)) ret with sendLineToSocket(socket, ""); // idle

  if (jmatch("computerID=*", line, m)) {

    if (state.computerID != null || !possibleComputerID($1)) ret;

    print("Got computer ID: " + $1);

    state.computerID = $1;

    cI.gotConnection(state.computerID);

    ret;

  }

  if (jmatch("sub *", line, m)) {

    Channel channel = getChannelForSocket(socket, $1);

    if (channel == null)

      sendLineToSocket(socket, "Channel " + $1 + " not found");

    else

      sendLineToSocket(socket, "Subbed " + $1);

    ret;

  }

  L<S> tok = javaTok(line);

  if (eq(get(tok, 3), ":")) {

    Channel channel = state.subbedChannels.get(get(tok, 1));

    if (channel != null)

      channel.handleLine(joinSubList(tok, 5), socket);

  }

  // ignore unknown line

}

html { ret "Use port " + port; }

svoid sendLineToSocket(SocketChannel socket, S s) ctex {

  byte[] bytes = toUtf8(s + "\n");

  sendBytesToSocket(socket, bytes);

}

svoid removeFromAllChannels(SocketChannel socket) {

  lock channelsLock;

  for (S name, Channel c : cloneMap(channels)) { pcall {

    c.removeSocket(socket);

    if (empty(c.sockets)) {

      print("Deleting channel: " + name);

      channels.remove(name);

    }

  }}

}

static Channel getChannelForSocket(SocketChannel socket, S name) {

  lock channelsLock;

  Channel c = channels.get(name);

  if (c == null) {

    c = makeChannel(name);

    c.name = name;

    if (c == null) null;

    print("New channel: " + name);

    channels.put(name, c);

    lock c.lock;

    c.update();

  }

  c.addSocket(socket);

  state(socket).subbedChannels.put(name, c);

  ret c;

}

static Channel makeChannel(S name) {

  if (eq(name, "computerIDs")) ret new ComputerIDsChannel;

  if (eq(name, "computerCount")) ret new ComputerCountChannel;

  if (eq(name, "chat")) ret new PublicChatChannel;

  if (eq(name, "machineChat")) ret new PublicChatChannel;

  if (eq(name, "privateChat")) ret new PrivateChatChannel;

  if (eq(name, "privateMachineChat")) ret nu(PrivateChatChannel, prefix := name + ":");

  if (eq(name, "snippetUpdates")) ret new SnippetUpdatesChannel;

  if (eq(name, "gazelleChat")) ret new PublicChatChannel;

  if (eq(name, "gazellePresence")) ret new PresenceCountingChannel;

  ret null with print("Unknown channel: " + name);

}

svoid update {

  print("OS Instances update");

  if (empty(dataMapper)) ret;

  int value = cI.value();

  //if (value == theNumber) ret;

  theNumber = value;

  print("OS Instances update done");

}