tankz/src/common/haxe/ru/m/connect/neko/NekoWSConnection.hx

package ru.m.connect.neko;

import haxe.crypto.BaseCode;
import haxe.crypto.Sha1;
import haxe.io.Bytes;
import haxe.io.BytesBuffer;
import protohx.Message;
import sys.net.Socket;

class NekoWSConnection<O:Message, I:Message> extends NekoConnection<O, I> {
  private var opened:Bool;

  public function new(socket:Socket, i:Class<I>) {
    super(socket, i);
    opened = false;
  }

  override private function sendPacket(packet:O):Void {
    var data = PacketUtil.toBytes(packet);
    writeData(data, socket);
  }

  override public function pushData(bytes:Bytes):Void {
    if (!opened) {
      var str:String = bytes.getString(0, bytes.length);
      if (StringTools.startsWith(str, "GET")) {
        var r = ~/Sec-WebSocket-Key:\s*([A-z0-9=+\/]+)/;
        r.match(str);
        opened = true;
        sendServerHandShake(socket, r.matched(1));
      }
    } else {
      var data = parseData(bytes);
      if (data != null) {
        var packet:I = PacketUtil.fromBytes(data, queue.packetClass);
        receive(packet);
      }
    }
  }

  private function encodeBase64(content:String):String {
    var suffix = switch (content.length % 3) {
      case 2: "=";
      case 1: "==";
      default: "";
    };
    return BaseCode.encode(content, "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/") + suffix;
  }

  private function hex2data(hex:String):String {
    var data = "";
    for (i in 0...Std.int(hex.length / 2)) {
      data += String.fromCharCode(Std.parseInt("0x" + hex.substr(i * 2, 2)));
    }
    return data;
  }

  private function sendServerHandShake(socket:sys.net.Socket, inpKey:String) {
    var outKey = encodeBase64(hex2data(Sha1.encode(StringTools.trim(inpKey) + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")));

    var s = "HTTP/1.1 101 Switching Protocols\r\n" + "Upgrade: websocket\r\n" + "Connection: Upgrade\r\n"
      + "Sec-WebSocket-Accept: " + outKey + "\r\n" + "\r\n";

    socket.output.writeString(s);
  }

  private function writeData(data:Bytes, socket:sys.net.Socket, isServer = true):Void {
    socket.output.writeByte(130);

    var len = 0;
    if (data.length < 126)
      len = data.length;
    else if (data.length < 65536)
      len = 126;
    else
      len = 127;

    socket.output.writeByte(len | (!isServer ? 0x80 : 0x00));

    if (data.length >= 126) {
      if (data.length < 65536) {
        socket.output.writeByte((data.length >> 8) & 0xFF);
        socket.output.writeByte(data.length & 0xFF);
      } else {
        socket.output.writeByte((data.length >> 24) & 0xFF);
        socket.output.writeByte((data.length >> 16) & 0xFF);
        socket.output.writeByte((data.length >> 8) & 0xFF);
        socket.output.writeByte(data.length & 0xFF);
      }
    }

    if (isServer) {
      socket.output.writeBytes(data, 0, data.length);
    } else {
      var mask = [Std.random(256), Std.random(256), Std.random(256), Std.random(256)];
      socket.output.writeByte(mask[0]);
      socket.output.writeByte(mask[1]);
      socket.output.writeByte(mask[2]);
      socket.output.writeByte(mask[3]);
      var maskedData = new BytesBuffer();
      for (i in 0...data.length) {
        maskedData.addByte(data.get(i) ^ mask[i % 4]);
      }
      socket.output.writeBytes(maskedData.getBytes(), 0, maskedData.length);
    }
  }

  private function parseData(bytes:Bytes):Bytes {
    var p = 0;
    var opcode = bytes.get(p++);

    if (opcode == 0x00) {
      var data = new BytesBuffer();
      var b:Int;
      while ((b = bytes.get(p++)) != 0xFF) {
        data.addByte(b);
      }
      return data.getBytes();
    }

    // 130 = binary data
    if (opcode == 130) {
      var len = bytes.get(p++);

      // mask
      if (len & 0x80 != 0) {
        len &= 0x7F;

        if (len == 126) {
          var b2 = bytes.get(p++);
          var b3 = bytes.get(p++);
          len = (b2 << 8) + b3;
        } else if (len == 127) {
          var b2 = bytes.get(p++);
          var b3 = bytes.get(p++);
          var b4 = bytes.get(p++);
          var b5 = bytes.get(p++);
          len = (b2 << 24) + (b3 << 16) + (b4 << 8) + b5;
        }

        // Lib.println("len = " + len);

        // direct array init not work corectly!
        var mask = [];
        mask.push(bytes.get(p++));
        mask.push(bytes.get(p++));
        mask.push(bytes.get(p++));
        mask.push(bytes.get(p++));

        // Lib.println("mask = " + mask);

        var data = new BytesBuffer();
        for (i in 0...len) {
          data.addByte(bytes.get(p++) ^ mask[i % 4]);
        }

        // Lib.println("readed = " + data.toString());
        return data.getBytes();
      } else {
        throw "Expected masked data.";
      }
    }

    if (opcode == 136) {
      // socket.close();
      opened = false;
      return null;
    } else {
      throw "Unsupported websocket opcode: " + opcode;
    }
    return null;
  }
}