LNXSDK/leenkx/Sources/leenkx/network/krom/KromSocket.hx

package leenkx.network.krom;

import haxe.io.Bytes;
import haxe.io.BytesInput;
import haxe.io.BytesOutput;
import haxe.io.Error;

#if (haxe_ver < 4)
typedef JsBuffer = js.html.ArrayBuffer;
#else
typedef JsBuffer = js.lib.ArrayBuffer;
#end

@:native("krom_socket_create") extern function krom_socket_create():Int;
@:native("krom_socket_close") extern function krom_socket_close(id:Int):Void;
@:native("krom_socket_bind") extern function krom_socket_bind(id:Int, addr:String, port:Int):Bool;
@:native("krom_socket_listen") extern function krom_socket_listen(id:Int, backlog:Int):Bool;
@:native("krom_socket_accept") extern function krom_socket_accept(id:Int):Int;
@:native("krom_socket_connect") extern function krom_socket_connect(id:Int, host:String, port:Int):Bool;
@:native("krom_socket_send") extern function krom_socket_send(id:Int, data:JsBuffer):Int;
@:native("krom_socket_recv") extern function krom_socket_recv(id:Int, maxLen:Int):Dynamic;
@:native("krom_socket_set_blocking") extern function krom_socket_set_blocking(id:Int, blocking:Bool):Void;
@:native("krom_socket_is_connected") extern function krom_socket_is_connected(id:Int):Bool;

class KromSocket {
	private var _socketId:Int = -1;
	private var _host:KromHost = null;
	private var _port:Int = 0;
	private var _blocking:Bool = true;
	
	public var input:KromSocketInput = null;
	public var output:KromSocketOutput = null;
	
	private static var _connections:Array<KromSocket> = [];
	private var _newConnections:Array<KromSocket> = [];
	
	public function new() {
		_socketId = krom_socket_create();
		if (_socketId >= 0) {
			input = new KromSocketInput(this);
			output = new KromSocketOutput(this);
		}
	}
	
	private function setSocketId(id:Int):Void {
		_socketId = id;
		if (_socketId >= 0) {
			input = new KromSocketInput(this);
			output = new KromSocketOutput(this);
		}
	}
	
	public function getSocketId():Int {
		return _socketId;
	}
	
	public function bind(host:KromHost, port:Int):Void {
		_host = host;
		_port = port;
		if (_socketId >= 0) {
			if (!krom_socket_bind(_socketId, host.host, port)) {
				trace("Failed to bind to " + host.host + ":" + port);
			}
		}
	}
	
	public function listen(connections:Int):Void {
		if (_host == null) {
			throw "You must bind the Socket to an address!";
		}
		if (_socketId >= 0) {
			if (!krom_socket_listen(_socketId, connections)) {
				trace("Failed to listen");
			}
		}
	}
	
	public function accept():KromSocket {
		if (_socketId < 0) {
			throw "Blocking";
		}
		
		var clientId = krom_socket_accept(_socketId);
		if (clientId < 0) {
			throw "Blocking";
		}
		
		var clientSocket = new KromSocket();
		clientSocket.setSocketId(clientId);
		_connections.push(clientSocket);
		return clientSocket;
	}
	
	public function connect(host:KromHost, port:Int):Void {
		if (_socketId >= 0) {
			if (!krom_socket_connect(_socketId, host.host, port)) {
				throw "Connection failed to " + host.host + ":" + port;
			}
		}
	}
	
	public function setBlocking(blocking:Bool):Void {
		_blocking = blocking;
		if (_socketId >= 0) {
			krom_socket_set_blocking(_socketId, blocking);
		}
	}
	
	public function setTimeout(timeout:Int):Void {
		// TODO: Re-investigate timeout handling
	}
	
	public function peer():{host:KromHost, port:Int} {
		return {host: _host != null ? _host : new KromHost("0.0.0.0"), port: _port};
	}
	
	public function host():{host:KromHost, port:Int} {
		return {host: _host != null ? _host : new KromHost("0.0.0.0"), port: _port};
	}
	
	public function close():Void {
		if (_socketId >= 0) {
			krom_socket_close(_socketId);
			_socketId = -1;
		}
		_connections.remove(this);
	}
	
	public function sendRaw(data:Bytes):Int {
		var arrayBuffer:JsBuffer = data.getData();
		var exactBuffer:JsBuffer = untyped arrayBuffer.slice(0, data.length);
		return krom_socket_send(_socketId, exactBuffer);
	}
	
	// throws "ConnectionClosed" if peer disconnected
	public function recvRaw(maxLength:Int):Bytes {
		if (_socketId < 0) return null;
		var result:Dynamic = krom_socket_recv(_socketId, maxLength);
		// V8 returns ArrayBuffer data, null would-block, or -2 connection closed
		if (result == -2) {
			_socketId = -1;
			throw "ConnectionClosed";
		}
		if (result == null) return null;
		return Bytes.ofData(result);
	}
	
	public function isConnected():Bool {
		if (_socketId < 0) return false;
		return krom_socket_is_connected(_socketId);
	}
	
	public static function select(read:Array<KromSocket>, write:Array<KromSocket>, others:Array<KromSocket>, ?timeout:Float):{read:Array<KromSocket>, write:Array<KromSocket>, others:Array<KromSocket>} {
		var readable:Array<KromSocket> = [];
		
		if (read != null) {
			for (sock in read) {
				if (sock._socketId >= 0) {
					readable.push(sock);
				}
			}
		}
		
		if (readable.length == 0) {
			return null;
		}
		
		return {
			read: readable,
			write: write,
			others: others
		};
	}
}

class KromSocketInput {
	private var _socket:KromSocket;
	public var hasData:Bool = false;
	private var _buffer:Bytes = null;
	private var _bufferPos:Int = 0;
	
	public function new(socket:KromSocket) {
		_socket = socket;
	}
	
	public function readBytes(buf:Bytes, pos:Int, len:Int):Int {
		var data = _socket.recvRaw(len);
		if (data == null || data.length == 0) {
			hasData = false;
			throw Error.Blocked;
		}
		hasData = true;
		var toRead = data.length < len ? data.length : len;
		buf.blit(pos, data, 0, toRead);
		return toRead;
	}
	
	public function read(nbytes:Int):Bytes {
		var buf = Bytes.alloc(nbytes);
		var read = readBytes(buf, 0, nbytes);
		if (read < nbytes) {
			return buf.sub(0, read);
		}
		return buf;
	}
}

class KromSocketOutput {
	private var _socket:KromSocket;
	private var _buffer:BytesOutput;
	
	public function new(socket:KromSocket) {
		_socket = socket;
		_buffer = new BytesOutput();
	}
	
	public function write(data:Bytes):Void {
		_buffer.write(data);
	}
	
	public function writeBytes(buf:Bytes, pos:Int, len:Int):Int {
		_buffer.writeBytes(buf, pos, len);
		return len;
	}
	
	public function flush():Void {
		var data = _buffer.getBytes();
		if (data.length > 0) {
			_socket.sendRaw(data);
			_buffer = new BytesOutput();
		}
	}
}

class KromHost {
	public var host:String;
	
	public function new(hostname:String) {
		host = hostname;
	}
	
	public static function resolve(hostname:String):String {
		return hostname;
	}
	
	public function toString():String {
		return host;
	}
}