From 3c5105bb7485a4a3532708bd36371b5d08b62b7c Mon Sep 17 00:00:00 2001
From: Valentin <valentin@devling.xyz>
Date: Mon, 4 Dec 2023 10:42:01 +0000
Subject: [PATCH] Generate servingNetworkName from vPLMN instead of hPLMN, Fix
 #32

---
 .../control_test_engine/gnb/context/context.go |  2 +-
 .../control_test_engine/gnb/context/msg.go     |  2 ++
 .../gnb/nas/service/service.go                 |  2 ++
 .../control_test_engine/ue/context/context.go  | 18 +++++++++++-------
 .../ue/nas/service/service.go                  |  2 ++
 internal/control_test_engine/ue/ue.go          |  5 ++++-
 6 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/internal/control_test_engine/gnb/context/context.go b/internal/control_test_engine/gnb/context/context.go
index 1881a637..2dcde594 100644
--- a/internal/control_test_engine/gnb/context/context.go
+++ b/internal/control_test_engine/gnb/context/context.go
@@ -386,7 +386,7 @@ func (gnb *GNBContext) GetSliceInBytes() ([]byte, []byte) {
 	return sstBytes, nil
 }
 
-func (gnb *GNBContext) getMccAndMnc() (string, string) {
+func (gnb *GNBContext) GetMccAndMnc() (string, string) {
 	return gnb.controlInfo.mcc, gnb.controlInfo.mnc
 }
 
diff --git a/internal/control_test_engine/gnb/context/msg.go b/internal/control_test_engine/gnb/context/msg.go
index 12063d4e..d00d8113 100644
--- a/internal/control_test_engine/gnb/context/msg.go
+++ b/internal/control_test_engine/gnb/context/msg.go
@@ -15,4 +15,6 @@ type UEMessage struct {
 	ConnectionClosed bool
 	AmfId int64
 	Msin string
+	Mcc string
+	Mnc string
 }
diff --git a/internal/control_test_engine/gnb/nas/service/service.go b/internal/control_test_engine/gnb/nas/service/service.go
index 7cf0db4e..6ab3bb0f 100644
--- a/internal/control_test_engine/gnb/nas/service/service.go
+++ b/internal/control_test_engine/gnb/nas/service/service.go
@@ -29,6 +29,8 @@ func gnbListen(gnb *context.GNBContext) {
 		// select AMF and get sctp association
 		// make a tun interface
 		ue := gnb.NewGnBUe(message.GNBTx, message.GNBRx, message.Msin)
+		mcc, mnc := gnb.GetMccAndMnc()
+		message.GNBTx <- context.UEMessage{Mcc: mcc, Mnc: mnc}
 		ue.SetPduSessions(message.GNBPduSessions)
 
 		if ue == nil {
diff --git a/internal/control_test_engine/ue/context/context.go b/internal/control_test_engine/ue/context/context.go
index fbe77ba8..0315484a 100644
--- a/internal/control_test_engine/ue/context/context.go
+++ b/internal/control_test_engine/ue/context/context.go
@@ -62,6 +62,8 @@ type Amf struct {
 	amfSetId    uint16
 	amfPointer  uint8
 	amfUeId     int64
+	mcc         string
+	mnc         string
 }
 
 type UEPDUSession struct {
@@ -165,9 +167,6 @@ func (ue *UEContext) NewRanUeContext(msin string,
 		}
 	}
 
-	// added snn.
-	ue.UeSecurity.Snn = ue.deriveSNN()
-
 	ue.scenarioChan = scenarioChan
 
 	// added initial state for MM(NULL)
@@ -377,12 +376,11 @@ func (pduSession *UEPDUSession) GetStateSM() int {
 func (ue *UEContext) deriveSNN() string {
 	// 5G:mnc093.mcc208.3gppnetwork.org
 	var resu string
-	if len(ue.UeSecurity.mnc) == 2 {
-		resu = "5G:mnc0" + ue.UeSecurity.mnc + ".mcc" + ue.UeSecurity.mcc + ".3gppnetwork.org"
+	if len(ue.amfInfo.mnc) == 2 {
+		resu = "5G:mnc0" + ue.amfInfo.mnc + ".mcc" + ue.amfInfo.mcc + ".3gppnetwork.org"
 	} else {
-		resu = "5G:mnc" + ue.UeSecurity.mnc + ".mcc" + ue.UeSecurity.mcc + ".3gppnetwork.org"
+		resu = "5G:mnc" + ue.amfInfo.mnc + ".mcc" + ue.amfInfo.mcc + ".3gppnetwork.org"
 	}
-
 	return resu
 }
 
@@ -506,6 +504,12 @@ func (ue *UEContext) GetAmfUeId() int64 {
 	return ue.amfInfo.amfUeId
 }
 
+func (ue *UEContext) SetAmfMccAndMnc(mcc string, mnc string) {
+	ue.amfInfo.mcc = mcc
+	ue.amfInfo.mnc = mnc
+	ue.UeSecurity.Snn = ue.deriveSNN()
+}
+
 func (ue *UEContext) Get5gGuti() [4]uint8 {
 	return ue.UeSecurity.Guti
 }
diff --git a/internal/control_test_engine/ue/nas/service/service.go b/internal/control_test_engine/ue/nas/service/service.go
index 57d84f87..2ecb9333 100644
--- a/internal/control_test_engine/ue/nas/service/service.go
+++ b/internal/control_test_engine/ue/nas/service/service.go
@@ -16,4 +16,6 @@ func InitConn(ue *context.UEContext, gnb *gnbContext.GNBContext) {
 
 	// Send channels to gNB
 	inboundChannel <- gnbContext.UEMessage{GNBTx: ue.GetGnbTx(), GNBRx: ue.GetGnbRx(), Msin: ue.GetMsin()}
+	msg := <-ue.GetGnbTx()
+	ue.SetAmfMccAndMnc(msg.Mcc, msg.Mnc)
 }
diff --git a/internal/control_test_engine/ue/ue.go b/internal/control_test_engine/ue/ue.go
index e4b4076c..f0d9127d 100644
--- a/internal/control_test_engine/ue/ue.go
+++ b/internal/control_test_engine/ue/ue.go
@@ -87,8 +87,11 @@ func gnbMsgHandler(msg context2.UEMessage, ue *context.UEContext) {
 		// handling NAS message.
 		ue.SetAmfUeId(msg.AmfId)
 		state.DispatchState(ue, msg.Nas)
-	} else {
+	} else if msg.GNBPduSessions[0] != nil {
+		// Setup PDU Session
 		serviceGtp.SetupGtpInterface(ue, msg)
+	} else {
+		log.Error("[UE] Received unknown message from gNodeB", msg)
 	}
 }