module S1AP_Emulation {

/* S1AP Emulation, runs on top of S1AP_CodecPort. It multiplexes/demultiplexes
 * the individual subscribers by their UE association (MME_UE_S1AP_ID/
 * ENB_UE_S1AP_ID identifiers), so there can be separate TTCN-3 components
 * handling each of them.
 *
 * The S1AP_Emulation.main() function processes S1AP primitives from the S1AP
 * socket via the S1AP_CodecPort, and dispatches them to the per-subscriber
 * components.
 *
 * For each new subscruber, the S1apOps.create_cb() is called. It can create
 * or resolve a TTCN-3 component, and returns a component reference to which
 * that subscriber traffic is routed/dispatched.
 *
 * If a pre-existing component wants to register to handle a future inbound UE
 * association, it can do so by registering an "expect" with the expected
 * MME_UE_S1AP_ID/ENB_UE_S1AP_ID identifiers. It is also possible to register
 * an expect for a specific procedureCode, in case the expected message is non
 * UE related (unit-data).
 *
 * Inbound non-UE related S1AP messages (such as RESET, SETUP, OVERLOAD) are
 * dispatched to the S1apOps.unitdata_cb() callback, which is registered with
 * an argument to the main() function below.
 *
 * (C) 2019 by Harald Welte <laforge@gnumonks.org>
 * All rights reserved.
 *
 * Released under the terms of GNU General Public License, Version 2 or
 * (at your option) any later version.
 *
 * SPDX-License-Identifier: GPL-2.0-or-later
 */

import from S1AP_CodecPort all;
import from S1AP_CodecPort_CtrlFunct all;
import from S1AP_Types all;
import from S1AP_Constants all;
import from S1AP_PDU_Contents all;
import from S1AP_PDU_Descriptions all;
import from S1AP_IEs all;
import from S1AP_Templates all;
import from SCTP_Templates all;

import from NAS_EPS_Types all;
import from NAS_Templates all;

import from LTE_CryptoFunctions all;

import from General_Types all;
import from Osmocom_Types all;
import from IPL4asp_Types all;
import from DNS_Helpers all;


type component S1AP_ConnHdlr {
	port S1AP_Conn_PT S1AP;
	/* procedure based port to register for incoming connections */
	port S1APEM_PROC_PT S1AP_PROC;
}

/* port between individual per-connection components and this dispatcher */
type port S1AP_Conn_PT message {
	inout S1AP_PDU, PDU_NAS_EPS, S1APEM_Config;
} with { extension "internal" };

type record NAS_Keys {
	octetstring k_nas_int,
	octetstring k_nas_enc
};
type record ResetNAScounts {
/* empty */
};
type union S1APEM_Config {
	NAS_Keys set_nas_keys,
	ResetNAScounts reset_nas_counts,
	NAS_ALG_INT set_nas_alg_int
};

type enumerated S1APEM_EventUpDown {
	S1APEM_EVENT_DOWN,
	S1APEM_EVENT_UP
}

/* an event indicating us whether or not a connection is physically up or down,
 * and whether we have received an ID_ACK */
type union S1APEM_Event {
	S1APEM_EventUpDown	up_down
}

/* global test port e.g. for non-imsi/conn specific messages */
type port S1AP_PT message {
	inout S1AP_PDU, S1APEM_Event;
} with { extension "internal" };


/* represents a single S1AP Association */
type record AssociationData {
	S1AP_ConnHdlr	comp_ref,			/* component handling this UE connection */
	uint24_t	enb_ue_s1ap_id optional,	/* eNB side S1AP ID */
	uint32_t	mme_ue_s1ap_id optional,	/* MME side S1AP ID */
	EUTRAN_CGI	cgi optional,
	TAI		tai optional,
	NAS_UE_State	nus
};

type component S1AP_Emulation_CT {
	/* Port facing to the UDP SUT */
	port S1AP_CODEC_PT S1AP;
	/* All S1AP_ConnHdlr S1AP ports connect here
	 * S1AP_Emulation_CT.main needs to figure out what messages
	 * to send where with CLIENT.send() to vc_conn */
	port S1AP_Conn_PT S1AP_CLIENT;
	/* currently tracked connections */
	var AssociationData S1apAssociationTable[16];
	/* pending expected S1AP Association (UE oriented) */
	var ExpectData S1apExpectTable[8];
	/* pending expected S1AP PDU */
	var ExpectDataProc S1apExpectTableProc[8];
	/* procedure based port to register for incoming connections */
	port S1APEM_PROC_PT S1AP_PROC;
	/* test port for unit data messages */
	port S1AP_PT S1AP_UNIT;

	var S1AP_conn_parameters g_pars;
	var charstring g_s1ap_id;
	var integer g_s1ap_conn_id := -1;
}

type function S1APCreateCallback(S1AP_PDU msg, template (omit) MME_UE_S1AP_ID mme_id,
				 template (omit) ENB_UE_S1AP_ID enb_id, charstring id)
runs on S1AP_Emulation_CT return S1AP_ConnHdlr;

type function S1APUnitdataCallback(S1AP_PDU msg)
runs on S1AP_Emulation_CT return template S1AP_PDU;

type record S1APOps {
	S1APCreateCallback create_cb,
	S1APUnitdataCallback unitdata_cb
}

type record S1AP_conn_parameters {
	HostName remote_ip,
	PortNumber remote_sctp_port,
	HostName local_ip,
	PortNumber local_sctp_port,
	NAS_Role role
}

function tr_S1AP_RecvFrom_R(template S1AP_PDU msg)
runs on S1AP_Emulation_CT return template S1AP_RecvFrom {
	var template S1AP_RecvFrom mrf := {
		connId := g_s1ap_conn_id,
		remName := ?,
		remPort := ?,
		locName := ?,
		locPort := ?,
		msg := msg
	}
	return mrf;
}

private function f_s1ap_ids_known(template (omit) MME_UE_S1AP_ID mme_id,
				  template (omit) ENB_UE_S1AP_ID enb_id)
runs on S1AP_Emulation_CT return boolean {
	var integer i;
	log("f_s1ap_ids_known(",mme_id,", ",enb_id,")");
	for (i := 0; i < sizeof(S1apAssociationTable); i := i+1) {
		log("tbl[",i,"]: mme=", S1apAssociationTable[i].mme_ue_s1ap_id,
		    ", enb=", S1apAssociationTable[i].enb_ue_s1ap_id);
		/* skip empty records */
		if (S1apAssociationTable[i].mme_ue_s1ap_id == omit and
		    S1apAssociationTable[i].enb_ue_s1ap_id == omit) {
			log("skipping empty ", i);
			continue;
		}
		if (S1apAssociationTable[i].mme_ue_s1ap_id == omit) {
			log("entry ", i, " has no MME ID yet (enb=", S1apAssociationTable[i].enb_ue_s1ap_id);
			/* Table doesn't yet know the MME side ID, let's look-up only
			 * based on the eNB side ID */
			if (match(S1apAssociationTable[i].enb_ue_s1ap_id, enb_id)) {
				/* update table with MME side ID */
				S1apAssociationTable[i].mme_ue_s1ap_id := valueof(mme_id);
				return true;
			}
		} else if (match(S1apAssociationTable[i].enb_ue_s1ap_id, enb_id) and
			match(S1apAssociationTable[i].mme_ue_s1ap_id, mme_id)) {
			return true;
		}
	}
	return false;
}

private function f_comp_known(S1AP_ConnHdlr client)
runs on S1AP_Emulation_CT return boolean {
	var integer i;
	for (i := 0; i < sizeof(S1apAssociationTable); i := i+1) {
		if (S1apAssociationTable[i].comp_ref == client) {
			return true;
		}
	}
	return false;
}

private function f_assoc_id_by_s1ap_ids(template (omit) MME_UE_S1AP_ID mme_id,
				    template (omit) ENB_UE_S1AP_ID enb_id)
runs on S1AP_Emulation_CT return integer {
	var integer i;
	for (i := 0; i < sizeof(S1apAssociationTable); i := i+1) {
		if (istemplatekind(enb_id, "omit") or
		    match(S1apAssociationTable[i].enb_ue_s1ap_id, enb_id)) {
			if (istemplatekind(mme_id, "omit")) {
				return i;
			} else {
				if (match(S1apAssociationTable[i].mme_ue_s1ap_id, mme_id)) {
					return i;
				}
			}
		}
	}
	setverdict(fail, "S1AP Association Table not found by ENB-ID=", enb_id, " MME-ID=", mme_id);
	mtc.stop;
}

private function f_assoc_id_by_comp(S1AP_ConnHdlr client)
runs on S1AP_Emulation_CT return integer {
	var integer i;
	for (i := 0; i < sizeof(S1apAssociationTable); i := i+1) {
		if (S1apAssociationTable[i].comp_ref == client) {
			return i;
		}
	}
	setverdict(fail, "S1AP Association Table not found by component ", client);
	mtc.stop;
}

private function f_assoc_by_comp(S1AP_ConnHdlr client)
runs on S1AP_Emulation_CT return AssociationData {
	var integer i := f_assoc_id_by_comp(client);
	return S1apAssociationTable[i];
}

private function f_s1ap_id_table_add(S1AP_ConnHdlr comp_ref,
				     template (omit) MME_UE_S1AP_ID mme_id, ENB_UE_S1AP_ID enb_id)
runs on S1AP_Emulation_CT return integer {
	var integer i;
	for (i := 0; i < sizeof(S1apAssociationTable); i := i+1) {
		if (not isvalue(S1apAssociationTable[i].enb_ue_s1ap_id)) {
			S1apAssociationTable[i].enb_ue_s1ap_id := enb_id;
			if (istemplatekind(mme_id, "omit")) {
				S1apAssociationTable[i].mme_ue_s1ap_id := omit;
			} else {
				S1apAssociationTable[i].mme_ue_s1ap_id := valueof(mme_id);
			}
			S1apAssociationTable[i].comp_ref := comp_ref;
			return i;
		}
	}
	testcase.stop("S1AP Association Table full!");
	return -1;
}

private function f_s1ap_id_table_del(S1AP_ConnHdlr comp_ref, ENB_UE_S1AP_ID enb_id)
runs on S1AP_Emulation_CT {
	var integer i;
	for (i := 0; i < sizeof(S1apAssociationTable); i := i+1) {
		if (S1apAssociationTable[i].comp_ref == comp_ref and
		    S1apAssociationTable[i].enb_ue_s1ap_id == enb_id) {
			S1apAssociationTable[i].enb_ue_s1ap_id := omit;
			S1apAssociationTable[i].mme_ue_s1ap_id := omit;
			S1apAssociationTable[i].comp_ref := null;
			return;
		}
	}
	setverdict(fail, "S1AP Association Table: Couldn't find to-be-deleted entry!");
	mtc.stop;
}


private function f_s1ap_id_table_init()
runs on S1AP_Emulation_CT {
	for (var integer i := 0; i < sizeof(S1apAssociationTable); i := i+1) {
		S1apAssociationTable[i].mme_ue_s1ap_id := omit;
		S1apAssociationTable[i].enb_ue_s1ap_id := omit;
		S1apAssociationTable[i].cgi := omit;
		S1apAssociationTable[i].tai := omit;
		S1apAssociationTable[i].nus := valueof(t_NAS_UE_State(g_pars.role));
		S1apAssociationTable[i].comp_ref := null;
	}
}

private function f_s1ap_xceive(template (value) S1AP_PDU tx,
				template S1AP_PDU rx_t := ?)
runs on S1AP_Emulation_CT  return S1AP_PDU {
	timer T := 10.0;
	var S1AP_RecvFrom mrf;

	S1AP.send(t_S1AP_Send(g_s1ap_conn_id, tx));
	alt {
	[] S1AP.receive(tr_S1AP_RecvFrom_R(rx_t)) -> value mrf { }
	[] S1AP.receive(tr_SctpAssocChange) { repeat; }
	[] S1AP.receive(tr_SctpPeerAddrChange)  { repeat; }
	[] T.timeout {
		setverdict(fail, "Timeout waiting for ", rx_t);
		mtc.stop;
		}
	}
	return mrf.msg;
}

/*
private function f_nas_try_decaps(PDU_NAS_EPS nas) return PDU_NAS_EPS
{
	var PDU_NAS_EPS_SecurityProtectedNASMessage secp_nas;
	if (not match(nas, tr_NAS_EMM_SecurityProtected)) {
		return nas;
	}
	secp_nas := nas.ePS_messages.ePS_MobilityManagement.pDU_NAS_EPS_SecurityProtectedNASMessage;
	select (secp_nas.securityHeaderType) {
	case ('0011'B) {
		var octetstring knas_int := '530ce32318f26264eab26bc116870b86'O;
		var octetstring data_with_seq := int2oct(secp_nas.sequenceNumber, 1) & secp_nas.nAS_Message;
		var OCT4 exp_mac := f_snow_3g_f9(knas_int, secp_nas.sequenceNumber, 0,
						 is_downlink:=true, data:=data_with_seq);
		if (exp_mac != secp_nas.messageAuthenticationCode) {
			setverdict(fail, "Received NAS MAC ", secp_nas.messageAuthenticationCode,
				   " doesn't match expected MAC ", exp_mac, ": ", nas);
			mtc.stop;
		}
		return dec_PDU_NAS_EPS(secp_nas.nAS_Message);
		}
	case else  {
		setverdict(fail, "Implement SecHdrType for ", secp_nas);
		mtc.stop;
		}
	}
}
*/

function handle_S1AP_UeContextReleaseCmd(template (present) S1AP_PDU rel_cmd) runs on S1AP_Emulation_CT {
	if (ispresent(rel_cmd.initiatingMessage.value_.uEContextReleaseCommand.protocolIEs[0].value_.uE_S1AP_IDs.uE_S1AP_ID_pair)) {
		var template MME_UE_S1AP_ID mme_ue_id;
		var template ENB_UE_S1AP_ID enb_ue_id;
		var integer assoc_id;
		var S1AP_ConnHdlr vc_conn

		mme_ue_id := rel_cmd.initiatingMessage.value_.uEContextReleaseCommand.protocolIEs[0].value_.uE_S1AP_IDs.uE_S1AP_ID_pair.mME_UE_S1AP_ID;
		enb_ue_id := rel_cmd.initiatingMessage.value_.uEContextReleaseCommand.protocolIEs[0].value_.uE_S1AP_IDs.uE_S1AP_ID_pair.eNB_UE_S1AP_ID;

		assoc_id := f_assoc_id_by_s1ap_ids(mme_ue_id, enb_ue_id);
		vc_conn := S1apAssociationTable[assoc_id].comp_ref;

		f_s1ap_id_table_del(vc_conn, valueof(enb_ue_id));
	} else {
		/* TODO: The UE CONTEXT RELEASE COMMAND (see also: 3GPP TS 36.413, section 9.1.4.6), may identify the
		 * context by either an uE_S1AP_ID_pair (MME_UE_S1AP_ID and ENB_UE_S1AP_ID) or an MME_UE_S1AP_ID alone.
		 * The latter case is not implemented here yet. */
		setverdict(fail, "complete implementation of UeContextReleaseCmd handling");
		mtc.stop;
	}
}

private function SendToS1apExpectTableProc(S1AP_PDU msg) runs on S1AP_Emulation_CT {
	var integer procedureCode;
	var S1AP_ConnHdlr vc_conn;

	if (ispresent(msg.initiatingMessage.procedureCode)) {
		procedureCode := msg.initiatingMessage.procedureCode;
	} else if (ispresent(msg.unsuccessfulOutcome.procedureCode)) {
		procedureCode := msg.unsuccessfulOutcome.procedureCode;
	} else if (ispresent(msg.successfulOutcome.procedureCode)) {
		procedureCode := msg.successfulOutcome.procedureCode;
	} else {
		return;
	}

	for (var integer i := 0; i < sizeof(S1apExpectTableProc); i := i+1) {
		if (S1apExpectTableProc[i].procedureCode == procedureCode) {
			vc_conn := S1apExpectTableProc[i].vc_conn;
			if (vc_conn != null) {
				S1AP_CLIENT.send(msg) to vc_conn;
			}
		}
	}

	return;
}

function main(S1APOps ops, S1AP_conn_parameters p, charstring id) runs on S1AP_Emulation_CT {
	var Result res;
	g_pars := p;
	g_s1ap_id := id;
	f_s1ap_id_table_init();
	f_expect_table_init();

	map(self:S1AP, system:S1AP_CODEC_PT);
	if (p.remote_sctp_port == -1) {
		res := S1AP_CodecPort_CtrlFunct.f_IPL4_listen(S1AP, p.local_ip, p.local_sctp_port,
							      { sctp := valueof(ts_SctpTuple(18)) });
	} else {
		res := S1AP_CodecPort_CtrlFunct.f_IPL4_connect(S1AP, p.remote_ip, p.remote_sctp_port,
							       p.local_ip, p.local_sctp_port, -1,
							       { sctp := valueof(ts_SctpTuple(18)) });
	}
	if (not ispresent(res.connId)) {
		setverdict(fail, "Could not connect S1AP socket, check your configuration");
		mtc.stop;
	}
	g_s1ap_conn_id := res.connId;

	/* notify user about SCTP establishment */
	if (p.remote_sctp_port != -1) {
		S1AP_UNIT.send(S1APEM_Event:{up_down:=S1APEM_EVENT_UP})
	}

	while (true) {
		var S1AP_ConnHdlr vc_conn;
		var PDU_NAS_EPS nas;
		var MME_UE_S1AP_ID mme_id;
		var ENB_UE_S1AP_ID enb_id;
		var integer procedureCode;
		var S1AP_RecvFrom mrf;
		var S1AP_PDU msg;
		var S1APEM_Config s1cfg;
		var charstring vlr_name, mme_name;
		var integer ai;
		var octetstring kasme;

		alt {
		/* Configuration primitive from client */
		[] S1AP_CLIENT.receive(S1APEM_Config:{set_nas_keys:=?}) -> value s1cfg sender vc_conn {
			var integer assoc_id := f_assoc_id_by_comp(vc_conn);
			S1apAssociationTable[assoc_id].nus.k_nas_int := s1cfg.set_nas_keys.k_nas_int;
			S1apAssociationTable[assoc_id].nus.k_nas_enc := s1cfg.set_nas_keys.k_nas_enc;
			}
		/* Configuration primitive from client */
		[] S1AP_CLIENT.receive(S1APEM_Config:{reset_nas_counts:=?}) -> value s1cfg sender vc_conn {
			var integer assoc_id := f_assoc_id_by_comp(vc_conn);
			S1apAssociationTable[assoc_id].nus.rx_count := 0;
			S1apAssociationTable[assoc_id].nus.tx_count := 0;
			}
		/* Configuration primitive from client */
		[] S1AP_CLIENT.receive(S1APEM_Config:{set_nas_alg_int:=?}) -> value s1cfg sender vc_conn {
			var integer assoc_id := f_assoc_id_by_comp(vc_conn);
			S1apAssociationTable[assoc_id].nus.alg_int := s1cfg.set_nas_alg_int;
			}
		/* S1AP from client: InitialUE */
		[] S1AP_CLIENT.receive(tr_S1AP_InitialUE) -> value msg sender vc_conn {
			/* create a table entry about this connection */
			ai := f_s1ap_id_table_add(vc_conn, omit, valueof(f_S1AP_get_ENB_UE_S1AP_ID(msg)));
			/* Store CGI + TAI so we can use it for generating UlNasTransport from NAS */
			S1apAssociationTable[ai].tai := msg.initiatingMessage.value_.InitialUEMessage.protocolIEs[2].value_.TAI;
			S1apAssociationTable[ai].cgi := msg.initiatingMessage.value_.InitialUEMessage.protocolIEs[3].value_.EUTRAN_CGI;
			/* Pass message through */
			S1AP.send(t_S1AP_Send(g_s1ap_conn_id, msg));
			}
		/* NAS from client: Wrap in S1AP Uplink NAS Transport */
		[] S1AP_CLIENT.receive(PDU_NAS_EPS:?) -> value nas sender vc_conn {
			var integer assoc_id := f_assoc_id_by_comp(vc_conn);
			var AssociationData ad := S1apAssociationTable[assoc_id];
			nas := f_nas_encaps(S1apAssociationTable[assoc_id].nus, nas, new_ctx := false);
			var octetstring nas_enc := enc_PDU_NAS_EPS(nas);
			S1AP.send(t_S1AP_Send(g_s1ap_conn_id,
					      ts_S1AP_UlNasTransport(ad.mme_ue_s1ap_id,
								     ad.enb_ue_s1ap_id,
								     nas_enc, ad.cgi, ad.tai)));
			}
		/* S1AP from client: pass on transparently */
		[] S1AP_CLIENT.receive(S1AP_PDU:?) -> value msg sender vc_conn {
			/* Pass message through */
			/* FIXME: validate S1AP_IDs ? */
			S1AP.send(t_S1AP_Send(g_s1ap_conn_id, msg));
			}

		/* non-UE related S1AP: pass through unmodified/unverified */
		[] S1AP_UNIT.receive(S1AP_PDU:?) -> value msg sender vc_conn {
			/* Pass message through */
			S1AP.send(t_S1AP_Send(g_s1ap_conn_id, msg));
			}

		/* S1AP received from peer (MME) */
		[] S1AP.receive(tr_S1AP_RecvFrom_R(?)) -> value mrf {
			if (match(mrf.msg, tr_S1AP_nonUErelated)) {
				/* non-UE-related S1AP message */
				SendToS1apExpectTableProc(mrf.msg);
				var template S1AP_PDU resp := ops.unitdata_cb.apply(mrf.msg);
				if (isvalue(resp)) {
					S1AP.send(t_S1AP_Send(g_s1ap_conn_id, valueof(resp)));
				}
			} else {
				/* Ue-related S1AP message */
				/* obtain MME + ENB UE S1AP ID */
				var template (omit) MME_UE_S1AP_ID mme_ue_id := f_S1AP_get_MME_UE_S1AP_ID(mrf.msg);
				var template (omit) ENB_UE_S1AP_ID enb_ue_id := f_S1AP_get_ENB_UE_S1AP_ID(mrf.msg);
				/* check if those IDs are known in our table */
				if (f_s1ap_ids_known(mme_ue_id, enb_ue_id)) {
					/* if yes, dispatch to the ConnHdlr for this Ue-Connection */
					var template (omit) octetstring nas_enc;
					var integer assoc_id := f_assoc_id_by_s1ap_ids(mme_ue_id, enb_ue_id);
					vc_conn := S1apAssociationTable[assoc_id].comp_ref;
					nas_enc := f_S1AP_get_NAS_PDU(mrf.msg);
					if (isvalue(nas_enc)) {
						nas := dec_PDU_NAS_EPS(valueof(nas_enc));
						if (match(nas, tr_NAS_EMM_SecurityProtected)) {
							nas := f_nas_try_decaps(S1apAssociationTable[assoc_id].nus, nas);
						}
						/* DL/UlNasTransport are not interesting, don't send them */
						if (not match(mrf.msg, (tr_S1AP_DlNasTransport, tr_S1AP_UlNasTransport))) {
							/* send raw S1AP */
							S1AP_CLIENT.send(mrf.msg) to vc_conn;
						}
						/* send decoded NAS */
						S1AP_CLIENT.send(nas) to vc_conn;
					} else {
						/* send raw S1AP */
						S1AP_CLIENT.send(mrf.msg) to vc_conn;
					}
				} else {
					/* if not, call create_cb so it can create new ConnHdlr */
					vc_conn := ops.create_cb.apply(mrf.msg, mme_ue_id, enb_ue_id, id);
					f_s1ap_id_table_add(vc_conn, mme_ue_id, valueof(enb_ue_id));
					S1AP_CLIENT.send(mrf.msg) to vc_conn;
				}
				if (match(mrf.msg, tr_S1AP_UeContextReleaseCmd)) {
					handle_S1AP_UeContextReleaseCmd(mrf.msg);
				}
			}
			}
		[] S1AP.receive(tr_SctpAssocChange) { }
		[] S1AP.receive(tr_SctpPeerAddrChange)  { }
		[] S1AP_PROC.getcall(S1APEM_register:{?,?,?}) -> param(mme_id, enb_id, vc_conn) {
			f_create_expect(mme_id, enb_id, vc_conn);
			S1AP_PROC.reply(S1APEM_register:{mme_id, enb_id, vc_conn}) to vc_conn;
			}
		[] S1AP_PROC.getcall(S1APEM_register_proc:{?,?}) -> param(procedureCode, vc_conn) {
			f_create_expect_proc(procedureCode, vc_conn);
			S1AP_PROC.reply(S1APEM_register_proc:{procedureCode, vc_conn}) to vc_conn;
			}
		[] S1AP_PROC.getcall(S1APEM_derive_nas_token:{?, ?, -}) -> param(kasme, vc_conn) {
			var integer assoc_id := f_assoc_id_by_comp(vc_conn);
			var OCT32 nas_token := f_kdf_nas_token(kasme, S1apAssociationTable[assoc_id].nus.tx_count)
			S1apAssociationTable[assoc_id].nus.tx_count := S1apAssociationTable[assoc_id].nus.tx_count + 1;
			S1AP_PROC.reply(S1APEM_derive_nas_token:{kasme, vc_conn, nas_token}) to vc_conn;
			}
		}
	}
}

/* "Expect" Handling */

type record ExpectData {
	MME_UE_S1AP_ID mme_id optional,
	ENB_UE_S1AP_ID enb_id optional,
	S1AP_ConnHdlr vc_conn
}

/* represents a single S1AP PDU that we expect. When a matching PDU is seen, it is forwarded to the registered
 * component */
type record ExpectDataProc {
	integer procedureCode optional,
	S1AP_ConnHdlr vc_conn				/* component handling this UE connection */
};

signature S1APEM_register(in MME_UE_S1AP_ID mme_id, in ENB_UE_S1AP_ID enb_id, in S1AP_ConnHdlr hdlr);
signature S1APEM_register_proc(in integer procedureCode, in S1AP_ConnHdlr hdlr);
signature S1APEM_derive_nas_token(in octetstring kasme, in S1AP_ConnHdlr hdlr, out OCT32 nas_token);

type port S1APEM_PROC_PT procedure {
	inout S1APEM_register;
	inout S1APEM_register_proc;
	inout S1APEM_derive_nas_token;
} with { extension "internal" };

/* Function that can be used as create_cb and will use the expect table */
function ExpectedCreateCallback(S1AP_PDU msg,
				template (omit) MME_UE_S1AP_ID mme_id,
				template (omit) ENB_UE_S1AP_ID enb_id, charstring id)
runs on S1AP_Emulation_CT return S1AP_ConnHdlr {
	var S1AP_ConnHdlr ret := null;
	var integer i;

	for (i := 0; i < sizeof(S1apExpectTable); i := i+1) {
		if (not ispresent(S1apExpectTable[i].mme_id) and not ispresent(S1apExpectTable[i].enb_id)) {
			continue;
		}

		if (valueof(mme_id) == S1apExpectTable[i].mme_id and valueof(enb_id) == S1apExpectTable[i].enb_id) {
			ret := S1apExpectTable[i].vc_conn;
			/* Release this entry */
			S1apExpectTable[i].mme_id := omit;
			S1apExpectTable[i].enb_id := omit;
			S1apExpectTable[i].vc_conn := null;
			log("Found Expect[", i, "] for ", msg, " handled at ", ret);
			return ret;
		}
	}
	setverdict(fail, "Couldn't find Expect for ", msg);
	mtc.stop;
}

private function f_create_expect(template (omit) MME_UE_S1AP_ID mme_id,
				 template (omit) ENB_UE_S1AP_ID enb_id,
				 S1AP_ConnHdlr hdlr)
runs on S1AP_Emulation_CT {
	var integer i;

	/* Check an entry like this is not already presnt */
	for (i := 0; i < sizeof(S1apExpectTable); i := i+1) {
		if (not ispresent(S1apExpectTable[i].mme_id) and not ispresent(S1apExpectTable[i].enb_id)) {
			continue;
		}
		if (valueof(mme_id) == S1apExpectTable[i].mme_id and valueof(enb_id) == S1apExpectTable[i].enb_id) {
			setverdict(fail, "UE MME id / UE ENB id pair already present", mme_id, enb_id);
			mtc.stop;
		}
	}
	for (i := 0; i < sizeof(S1apExpectTable); i := i+1) {
		if (not ispresent(S1apExpectTable[i].mme_id) and not ispresent(S1apExpectTable[i].enb_id)) {
			S1apExpectTable[i].mme_id := valueof(mme_id);
			S1apExpectTable[i].enb_id := valueof(enb_id);
			S1apExpectTable[i].vc_conn := hdlr;
			log("Created Expect[", i, "] for UE MME id:", mme_id, ", UE ENB id:", enb_id, " to be handled at ", hdlr);
			return;
		}
	}
	testcase.stop("No space left in S1apExpectTable")
}

/* client/conn_hdlr side function to use procedure port to create expect in emulation */
function f_create_s1ap_expect(template (omit) MME_UE_S1AP_ID mme_id,
			      template (omit) ENB_UE_S1AP_ID enb_id) runs on S1AP_ConnHdlr {
	S1AP_PROC.call(S1APEM_register:{mme_id, enb_id, self}) {
		[] S1AP_PROC.getreply(S1APEM_register:{?,?,?}) {};
	}
}

private function f_create_expect_proc(integer procedureCode, S1AP_ConnHdlr hdlr) runs on S1AP_Emulation_CT {
	var integer i;

	/* Check an entry like this is not already presnt */
	for (i := 0; i < sizeof(S1apExpectTableProc); i := i+1) {
		if (S1apExpectTableProc[i].vc_conn == null) {
			continue;
		}
		if (S1apExpectTableProc[i].procedureCode == procedureCode) {
			setverdict(fail, "procedureCode ", procedureCode, " already present");
			mtc.stop;
		}
	}
	for (i := 0; i < sizeof(S1apExpectTableProc); i := i+1) {
		if (S1apExpectTableProc[i].vc_conn == null) {
			S1apExpectTableProc[i].procedureCode := procedureCode;
			S1apExpectTableProc[i].vc_conn := hdlr;
			log("Created Expect[", i, "] for PDU:", procedureCode, " to be handled at ", hdlr);
			return;
		}
	}
	testcase.stop("No space left in S1apExpectTableProc")
}

/* client/conn_hdlr side function to use procedure port to create expect (PDU) in emulation */
function f_create_s1ap_expect_proc(integer procedureCode, S1AP_ConnHdlr hdlr) runs on S1AP_ConnHdlr
{
	S1AP_PROC.call(S1APEM_register_proc:{procedureCode, self}) {
		[] S1AP_PROC.getreply(S1APEM_register_proc:{?,?}) {};
	}

	log(procedureCode);
}

/* Derive NAS Token (and post-increment ul_count): */
function f_s1apem_derive_nas_token(in octetstring kasme) runs on S1AP_ConnHdlr return OCT32
{
	var OCT32 nas_token;
	S1AP_PROC.call(S1APEM_derive_nas_token:{kasme, self, -}) {
		[] S1AP_PROC.getreply(S1APEM_derive_nas_token:{kasme, self, ?}) -> param(nas_token) {
			return nas_token;
		};
	}
}

private function f_expect_table_init()
runs on S1AP_Emulation_CT {
	var integer i;
	for (i := 0; i < sizeof(S1apExpectTable); i := i + 1) {
		S1apExpectTable[i].mme_id := omit;
		S1apExpectTable[i].enb_id := omit;
		S1apExpectTable[i].vc_conn := null;
	}

	for (i := 0; i < sizeof(S1apExpectTableProc); i := i + 1) {
		S1apExpectTableProc[i].procedureCode := omit;
		S1apExpectTableProc[i].vc_conn := null;
	}
}

function DummyUnitdataCallback(S1AP_PDU msg)
runs on S1AP_Emulation_CT return template S1AP_PDU {
	log("Ignoring S1AP ", msg);
	return omit;
}


function f_S1AP_get_ENB_UE_S1AP_ID(S1AP_PDU s1ap) return template (omit) ENB_UE_S1AP_ID
{
	if (ischosen(s1ap.initiatingMessage)) {
		var InitiatingMessage im := s1ap.initiatingMessage;
		select (s1ap) {
		case (tr_S1AP_InitialUE) {
			return im.value_.InitialUEMessage.protocolIEs[0].value_.ENB_UE_S1AP_ID;
			}
		case (tr_S1AP_DlNasTransport) {
			return im.value_.DownlinkNASTransport.protocolIEs[1].value_.ENB_UE_S1AP_ID;
			}
		case (tr_S1AP_UlNasTransport) {
			return im.value_.UplinkNASTransport.protocolIEs[1].value_.ENB_UE_S1AP_ID;
			}
		case (tr_S1AP_IntialCtxSetupReq) {
			return im.value_.initialContextSetupRequest.protocolIEs[1].value_.ENB_UE_S1AP_ID;
			}
		case (tr_S1AP_UeContextReleaseReq) {
			return im.value_.UEContextReleaseRequest.protocolIEs[1].value_.ENB_UE_S1AP_ID;
			}
		case (tr_S1AP_UeContextReleaseCmd) {
			if (ispresent(im.value_.uEContextReleaseCommand.protocolIEs[0].value_.uE_S1AP_IDs.uE_S1AP_ID_pair)) {
				return im.value_.uEContextReleaseCommand.protocolIEs[0].value_.uE_S1AP_IDs.uE_S1AP_ID_pair.eNB_UE_S1AP_ID;
			} else {
				return omit;
			}
			}
		case (tr_S1AP_ConnEstInd) {
			return im.value_.ConnectionEstablishmentIndication.protocolIEs[1].value_.ENB_UE_S1AP_ID;
			}
		/* TODO */
		}
	} else if (ischosen(s1ap.successfulOutcome)) {
		var SuccessfulOutcome so := s1ap.successfulOutcome;
		select (s1ap) {
		case (tr_S1AP_InitialCtxSetupResp) {
			return so.value_.initialContextSetupResponse.protocolIEs[1].value_.ENB_UE_S1AP_ID;
			}
		case (tr_S1AP_UeContextReleaseCompl) {
			return so.value_.UEContextReleaseComplete.protocolIEs[1].value_.ENB_UE_S1AP_ID;
			}
		/* TODO */
		}
	} else if (ischosen(s1ap.unsuccessfulOutcome)) {
		var UnsuccessfulOutcome uo := s1ap.unsuccessfulOutcome;
		select (s1ap) {
		case (tr_S1AP_InitialCtxSetupFail) {
			return uo.value_.initialContextSetupFailure.protocolIEs[1].value_.ENB_UE_S1AP_ID;
			}
		/* TODO */
		}
	}
	return omit;
}

function f_S1AP_get_MME_UE_S1AP_ID(S1AP_PDU s1ap) return template (omit) MME_UE_S1AP_ID
{
	if (ischosen(s1ap.initiatingMessage)) {
		var InitiatingMessage im := s1ap.initiatingMessage;
		select (s1ap) {
		case (tr_S1AP_DlNasTransport) {
			return im.value_.DownlinkNASTransport.protocolIEs[0].value_.MME_UE_S1AP_ID;
			}
		case (tr_S1AP_UlNasTransport) {
			return im.value_.UplinkNASTransport.protocolIEs[0].value_.MME_UE_S1AP_ID;
			}
		case (tr_S1AP_IntialCtxSetupReq) {
			return im.value_.initialContextSetupRequest.protocolIEs[0].value_.MME_UE_S1AP_ID;
			}
		case (tr_S1AP_UeContextReleaseReq) {
			return im.value_.UEContextReleaseRequest.protocolIEs[0].value_.MME_UE_S1AP_ID;
			}
		case (tr_S1AP_UeContextReleaseCmd) {
			if (ispresent(im.value_.uEContextReleaseCommand.protocolIEs[0].value_.uE_S1AP_IDs.uE_S1AP_ID_pair)) {
				return im.value_.uEContextReleaseCommand.protocolIEs[0].value_.uE_S1AP_IDs.uE_S1AP_ID_pair.mME_UE_S1AP_ID;
			} else {
				return im.value_.uEContextReleaseCommand.protocolIEs[0].value_.uE_S1AP_IDs.mME_UE_S1AP_ID;
			}
			}
		case (tr_S1AP_ConnEstInd) {
			return im.value_.ConnectionEstablishmentIndication.protocolIEs[0].value_.MME_UE_S1AP_ID;
			}
		/* TODO */
		}
	} else if (ischosen(s1ap.successfulOutcome)) {
		var SuccessfulOutcome so := s1ap.successfulOutcome;
		select (s1ap) {
		case (tr_S1AP_InitialCtxSetupResp) {
			return so.value_.initialContextSetupResponse.protocolIEs[0].value_.MME_UE_S1AP_ID;
			}
		case (tr_S1AP_UeContextReleaseCompl) {
			return so.value_.UEContextReleaseComplete.protocolIEs[0].value_.MME_UE_S1AP_ID;
			}
		/* TODO */
		}
	} else if (ischosen(s1ap.unsuccessfulOutcome)) {
		var UnsuccessfulOutcome uo := s1ap.unsuccessfulOutcome;
		select (s1ap) {
		case (tr_S1AP_InitialCtxSetupFail) {
			return uo.value_.initialContextSetupFailure.protocolIEs[0].value_.MME_UE_S1AP_ID;
			}
		/* TODO */
		}
	}
	return omit;
}

function f_S1AP_get_NAS_PDU(S1AP_PDU s1ap) return template (omit) NAS_PDU
{
	var integer i, j;

	if (ischosen(s1ap.initiatingMessage)) {
		var InitiatingMessage im := s1ap.initiatingMessage;
		select (s1ap) {
		case (tr_S1AP_DlNasTransport) {
			var DownlinkNASTransport msg := im.value_.DownlinkNASTransport;
			for (i := 0; i < lengthof(msg.protocolIEs); i := i+1) {
				if (msg.protocolIEs[i].id == id_NAS_PDU) {
					return msg.protocolIEs[i].value_.NAS_PDU;
				}
			}
			}
		case (tr_S1AP_UlNasTransport) {
			var UplinkNASTransport msg := im.value_.UplinkNASTransport;
			for (i := 0; i < lengthof(msg.protocolIEs); i := i+1) {
				if (msg.protocolIEs[i].id == id_NAS_PDU) {
					return msg.protocolIEs[i].value_.NAS_PDU;
				}
			}
			}
		case (tr_S1AP_IntialCtxSetupReq) {
			var InitialContextSetupRequest msg := im.value_.initialContextSetupRequest;
			for (i := 0; i < lengthof(msg.protocolIEs); i := i+1) {
				if (msg.protocolIEs[i].id == id_E_RABToBeSetupListCtxtSUReq) {
					var E_RABToBeSetupListCtxtSUReq rab_req := msg.protocolIEs[i].value_.E_RABToBeSetupListCtxtSUReq;
					for (j := 0; j < lengthof(rab_req); j := j+1) {
						var E_RABToBeSetupItemCtxtSUReq it := rab_req[j].value_.E_RABToBeSetupItemCtxtSUReq;
						return it.nAS_PDU;
					}
				}
			}
			return omit;
		}
		}
	}
	return omit;
}


}