//..............................................................................
//
// TCP packet decoder
//

import "io_IpDecoder.jnc"
import "log_Writer.jnc"
import "std_HashTable.jnc"

threadlocal log.Writer* g_logWriter;

//..............................................................................

class TcpDecoder: io.IpDecoder {
	enum Defaults {
		ConnectionTimeout = 5 * 60000, // 5 minutes of inactivity
		CleanupPeriod     = 10,        // remove timed out connections every 10 new connections
	}

protected:
	class TcpPeer {
		io.SocketAddress m_address;
		TcpSequencer m_sequencer;
		uint32_t m_finSeqNumber;
	}

	// use a simplified TCP state machine for *both* peers (not the RFC 793 one)

	enum ConnectionState {
		Connecting, // an initial SYN from the client has been seen
		Connected,  // a SYN | ACK from the server has been seen
		PickedUp,   // an existing connection has been picked up
		Closing,    // a single FIN has been seen (from client or server)
		Closed,     // two FINs from both sides has been seen
		Reset,      // a RST has been seen
	}

	bitflag enum ConnectionSide {
		Client,
		Server,
		Both = Client | Server
	}

	struct ConnectionKey {
		io.SocketAddress m_srcAddress;
		io.SocketAddress m_dstAddress;
	}

	class Connection {
		std.ListEntry* m_listEntry;
		ConnectionState m_state;
		ConnectionSide m_syn;
		ConnectionSide m_fin;
		uint64_t m_timestamp;
		TcpPeer m_clientPeer;
		TcpPeer m_serverPeer;
	}

protected:
	io.IpDefragmenter m_ipDefragmenter;
	std.HashTable m_connectionMap;
	std.List m_connectionList;
	Connection* m_currentConnection;
	size_t m_connectionSeqId;
	size_t m_outOfSeqLimit = TcpSequencer.Defaults.OutOfSeqLimit;

public:
	size_t bindable readonly m_connectionCount;

	// the following settings are freely adjustible

	io.SocketAddress m_filterAddress;
	size_t m_cleanupPeriod = Defaults.CleanupPeriod;
	uint_t m_connectionTimeout = Defaults.ConnectionTimeout;

public:
	void setup(
		size_t ipFramentLimit,
		uint_t ipFragmentTimeout = io.IpDefragmenter.Defaults.FragmentTimeout,
		size_t outOfSeqLimit = TcpSequencer.Defaults.OutOfSeqLimit
	);

	void reset();

protected:
	override bool decodeTransportHdr(
		uint64_t timestamp,
		io.SocketAddress* srcAddress,
		io.SocketAddress* dstAddress,
		io.IpProtocol protocol,
		void const* p,
		size_t size
	);

	Connection* createConnection(
		uint64_t timestamp,
		io.SocketAddress const* srcAddress,
		io.SocketAddress const* dstAddress,
		io.TcpHdr const* tcpHdr
	);

	void setCurrentConnection(
		uint64_t timestamp,
		Connection* connection
	);

	void finalizeConnection(
		uint64_t timestamp,
		Connection* connection,
		uint32_t seqNumber,
		uint32_t ackNumber,
		bool isClient
	);

	void removeTimedOutConnections(uint64_t timestamp);
}

// . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

void TcpDecoder.setup(
	size_t ipFramentLimit,
	uint_t ipFragmentTimeout,
	size_t outOfSeqLimit
) {
	m_ipDefragmenter.m_fragmentLimit = ipFramentLimit;
	m_ipDefragmenter.m_fragmentTimeout = ipFragmentTimeout;

	std.ListEntry* entry = m_connectionList.m_head;
	for (; entry; entry = entry.m_next) {
		Connection* connection = entry.m_data;
		connection.m_clientPeer.m_sequencer.m_outOfSeqLimit = outOfSeqLimit;
		connection.m_serverPeer.m_sequencer.m_outOfSeqLimit = outOfSeqLimit;
	}

	m_outOfSeqLimit = outOfSeqLimit;
}

void TcpDecoder.reset() {
	m_ipDefragmenter.reset();
	m_connectionList.clear();
	m_connectionMap.clear();
	m_currentConnection = null;
	m_connectionSeqId = 0;
}

bool TcpDecoder.decodeTransportHdr(
	uint64_t timestamp,
	io.SocketAddress* srcAddress,
	io.SocketAddress* dstAddress,
	io.IpProtocol protocol,
	void const* p,
	size_t size
) {
	if (protocol != io.IpProtocol.Tcp || size < sizeof(io.TcpHdr))
		return false;

	io.TcpHdr const* tcpHdr = (io.TcpHdr const*) p;
	size_t tcpHdrSize = tcpHdr.m_dataOffset * 4;
	if (tcpHdrSize < sizeof(io.TcpHdr) || tcpHdrSize > size)
		return false;

	srcAddress.m_port = tcpHdr.m_srcPort;
	dstAddress.m_port = tcpHdr.m_dstPort;

	p += tcpHdrSize;
	size -= tcpHdrSize;

	Connection* connection;
	ConnectionKey key = { *srcAddress, *dstAddress };
	std.MapEntry* entry = m_connectionMap.visit(key);
	if (entry.m_value) {
		connection = entry.m_value;
		if (connection.m_state >= ConnectionState.Closed)
			return true; // discard
	} else {
		if (!size && !(tcpHdr.m_flags & (io.TcpFlags.Syn | io.TcpFlags.Fin)))
			return false; // don't pick up empty segments

		if (m_filterAddress.m_family &&
			!srcAddress.isMatch(m_filterAddress) &&
			!dstAddress.isMatch(m_filterAddress))
			return false;

		connection = createConnection(timestamp, srcAddress, dstAddress, tcpHdr);
		entry.m_value = connection;
		key = { *dstAddress, *srcAddress };
		m_connectionMap[key] = connection;

		if (!(++m_connectionSeqId % m_cleanupPeriod))
			removeTimedOutConnections(timestamp);
	}

	bool isClient = connection.m_clientPeer.m_address.isEqual(srcAddress);
	TcpPeer* peer = isClient ? &connection.m_clientPeer : &connection.m_serverPeer;

	// a RST kills a connection instantly

	if (tcpHdr.m_flags & io.TcpFlags.Rst) {
		finalizeConnection(timestamp, connection, tcpHdr.m_seqNumber, tcpHdr.m_ackNumber, isClient);
		connection.m_state = ConnectionState.Reset;

		g_logWriter.write(
			timestamp,
			TcpFlowLogRecordCode.Reset,
			dstAddress,
			sizeof(io.SocketAddress),
			srcAddress,
			sizeof(io.SocketAddress)
		);

		connection.m_timestamp = timestamp;
		m_connectionList.moveToTail(connection.m_listEntry);
		return true;
	}

	// a SYN from server

	if ((tcpHdr.m_flags & io.TcpFlags.Syn) &&
		connection.m_state == ConnectionState.Connecting &&
		!isClient
	) {
		connection.m_state = ConnectionState.Connected;
		connection.m_syn |= ConnectionSide.Server;
		connection.m_serverPeer.m_sequencer.m_seqNumber = tcpHdr.m_seqNumber + 1; // fin pseudo-byte

		if (connection != m_currentConnection) {
			setCurrentConnection(timestamp, connection);
			g_logWriter.write(
				timestamp,
				SocketLogRecordCode.ConnectingFrom,
				srcAddress,
				sizeof(io.SocketAddress),
				dstAddress,
				sizeof(io.SocketAddress)
			);
		}

		g_logWriter.write(timestamp, SocketLogRecordCode.ConnectCompleted);
	}

	// if connected, sequence the payload

	if (size &&
		connection.m_state >= ConnectionState.Connected &&
		connection.m_state < ConnectionState.Closed
	) {
		setCurrentConnection(timestamp, connection);
		peer.m_sequencer.write(timestamp, tcpHdr.m_seqNumber, p, size);
	}

	// state transitions on FIN

	if (tcpHdr.m_flags & io.TcpFlags.Fin) {
		if (isClient) {
			connection.m_fin |= ConnectionSide.Client;
			connection.m_clientPeer.m_finSeqNumber = tcpHdr.m_seqNumber;
		} else {
			connection.m_fin |= ConnectionSide.Server;
			connection.m_serverPeer.m_finSeqNumber = tcpHdr.m_seqNumber;
		}

		if (connection.m_fin != (ConnectionSide.Client | ConnectionSide.Server))
			connection.m_state = ConnectionState.Closing;
		else {
			// account for the FIN pseudo-byte
			connection.m_clientPeer.m_sequencer.m_seqNumber = connection.m_clientPeer.m_finSeqNumber + 1;
			connection.m_serverPeer.m_sequencer.m_seqNumber = connection.m_serverPeer.m_finSeqNumber + 1;
			finalizeConnection(timestamp, connection, tcpHdr.m_seqNumber + 1, tcpHdr.m_ackNumber, isClient);
			connection.m_state = ConnectionState.Closed;

			g_logWriter.write(
				timestamp,
				TcpFlowLogRecordCode.Disconnected,
				srcAddress,
				sizeof(io.SocketAddress),
				dstAddress,
				sizeof(io.SocketAddress)
			);
		}
	}

	connection.m_timestamp = timestamp;
	m_connectionList.moveToTail(connection.m_listEntry);
	return true;
}

TcpDecoder.Connection* TcpDecoder.createConnection(
	uint64_t timestamp,
	io.SocketAddress const* srcAddress,
	io.SocketAddress const* dstAddress,
	io.TcpHdr const* tcpHdr
) {
	Connection* connection = new Connection;
	connection.m_timestamp = timestamp;
	connection.m_clientPeer.m_sequencer.m_outOfSeqLimit = m_outOfSeqLimit;
	connection.m_clientPeer.m_sequencer.m_logRecordCode = log.StdRecordCode.Tx;
	connection.m_serverPeer.m_sequencer.m_outOfSeqLimit = m_outOfSeqLimit;
	connection.m_serverPeer.m_sequencer.m_logRecordCode = log.StdRecordCode.Rx;

	if (!(tcpHdr.m_flags & io.TcpFlags.Syn)) { // pick up existing connection
		connection.m_state = ConnectionState.PickedUp;
		connection.m_syn = 0; // never seen any
		connection.m_clientPeer.m_address = *srcAddress;
		connection.m_clientPeer.m_sequencer.m_seqNumber = tcpHdr.m_seqNumber;
		connection.m_serverPeer.m_address = *dstAddress;
		connection.m_serverPeer.m_sequencer.m_seqNumber = tcpHdr.m_ackNumber;

		g_logWriter.write(
			timestamp,
			TcpFlowLogRecordCode.ConnectionPickedUp,
			dstAddress,
			sizeof(io.SocketAddress),
			srcAddress,
			sizeof(io.SocketAddress)
		);

	} else if (tcpHdr.m_flags & io.TcpFlags.Ack) { // syn-ack from server
		connection.m_state = ConnectionState.Connected;
		connection.m_syn = ConnectionSide.Client | ConnectionSide.Server;
		connection.m_clientPeer.m_address = *dstAddress;
		connection.m_clientPeer.m_sequencer.m_seqNumber = tcpHdr.m_ackNumber;
		connection.m_serverPeer.m_address = *srcAddress;
		connection.m_serverPeer.m_sequencer.m_seqNumber = tcpHdr.m_seqNumber + 1; // syn pseudo-byte

		g_logWriter.write(
			timestamp,
			SocketLogRecordCode.ConnectingFrom,
			srcAddress,
			sizeof(io.SocketAddress),
			dstAddress,
			sizeof(io.SocketAddress)
		);

		g_logWriter.write(timestamp, SocketLogRecordCode.ConnectCompleted);

	} else { // initial syn from client
		connection.m_state = ConnectionState.Connecting;
		connection.m_syn = ConnectionSide.Client;
		connection.m_clientPeer.m_address = *srcAddress;
		connection.m_clientPeer.m_sequencer.m_seqNumber = tcpHdr.m_seqNumber + 1; // syn pseudo-byte
		connection.m_serverPeer.m_address = *dstAddress;
		connection.m_serverPeer.m_sequencer.m_seqNumber = 0;

		g_logWriter.write(
			timestamp,
			SocketLogRecordCode.ConnectingFrom,
			dstAddress,
			sizeof(io.SocketAddress),
			srcAddress,
			sizeof(io.SocketAddress)
		);
	}

	connection.m_listEntry = m_connectionList.add(connection);
	m_connectionCount = m_connectionList.m_count;
	m_currentConnection = connection;
	return connection;
}

void TcpDecoder.setCurrentConnection(
	uint64_t timestamp,
	TcpDecoder.Connection* connection
) {
	if (connection == m_currentConnection)
		return;

	m_currentConnection = connection;

	g_logWriter.write(
		timestamp,
		TcpFlowLogRecordCode.ConnectionChanged,
		connection.m_serverPeer.m_address,
		sizeof(io.SocketAddress),
		connection.m_clientPeer.m_address,
		sizeof(io.SocketAddress)
	);
}

void TcpDecoder.finalizeConnection(
	uint64_t timestamp,
	TcpDecoder.Connection* connection,
	uint32_t seqNumber,
	uint32_t ackNumber,
	bool isClient
) {
	TcpSequencer* srcSequencer = isClient ?
		&connection.m_clientPeer.m_sequencer :
		&connection.m_serverPeer.m_sequencer;

	TcpSequencer* dstSequencer = isClient ?
		&connection.m_serverPeer.m_sequencer :
		&connection.m_clientPeer.m_sequencer;

	if (srcSequencer.m_outOfSeqList.m_count || dstSequencer.m_outOfSeqList.m_count) {
		setCurrentConnection(timestamp, connection);
		srcSequencer.flush(timestamp);
		dstSequencer.flush(timestamp);
	}

	if (srcSequencer.m_totalLostSize ||
		dstSequencer.m_totalLostSize ||
		srcSequencer.m_seqNumber != seqNumber ||
		dstSequencer.m_seqNumber != ackNumber
	) {
		setCurrentConnection(timestamp, connection);
		srcSequencer.finalize(timestamp, seqNumber);
		dstSequencer.finalize(timestamp, ackNumber);

		g_logWriter.write(
			timestamp,
			TcpFlowLogRecordCode.TotalDataLost,
			&connection.m_clientPeer.m_sequencer.m_totalLostSize,
			sizeof(uint32_t),
			&connection.m_serverPeer.m_sequencer.m_totalLostSize,
			sizeof(uint32_t)
		);
	}
}

void TcpDecoder.removeTimedOutConnections(uint64_t timestamp) {
	timestamp -= m_connectionTimeout * 10000;

	while (m_connectionList.m_count) {
		Connection* connection = m_connectionList.m_head.m_data;
		if (connection == m_currentConnection) // don't touch current connection
			continue;

		if (timestamp < connection.m_timestamp) // the rest is even more recent
			break;

		m_connectionList.remove(connection.m_listEntry);
		ConnectionKey key = { connection.m_clientPeer.m_address, connection.m_serverPeer.m_address };
		m_connectionMap.removeKey(key);
		key = { connection.m_serverPeer.m_address, connection.m_clientPeer.m_address };
		m_connectionMap.removeKey(key);
	}

	m_connectionCount = m_connectionList.m_count;
}

//..............................................................................
