fix: possible race conditions and nil dereference bugs in AMF (#1173)

Signed-off-by: Guillaume Belanger <guillaume.belanger27@gmail.com>

Author: gruyaume

internal/amf/amf.go

@@ -151,8 +151,12 @@ func (amf *AMF) AddAmfUeToUePool(ue *AmfUe) error {
 func (amf *AMF) DeregisterAndRemoveAMFUE(ctx context.Context, ue *AmfUe) {
 	ue.Deregister(ctx)
 
-	if ue.ranUe != nil {
-		err := ue.ranUe.Remove()
+	ue.Mutex.Lock()
+	ranUe := ue.ranUe
+	ue.Mutex.Unlock()
+
+	if ranUe != nil {
+		err := ranUe.Remove()
 		if err != nil {
 			logger.AmfLog.Error("failed to remove RAN UE", zap.Error(err))
 		}

internal/amf/ngap/handle_pdu_session_resource_modify.go

@@ -57,16 +57,9 @@ func HandlePDUSessionResourceNotify(ctx context.Context, amfInstance *amf.AMF, r
 		return
 	}
 
-	var ranUe *amf.RanUe
-
-	ranUe = ran.FindUEByRanUeNgapID(rANUENGAPID.Value)
-	if ranUe == nil {
-		logger.WithTrace(ctx, ran.Log).Warn("No UE Context", zap.Int64("RanUeNgapID", rANUENGAPID.Value))
-	}
-
-	ranUe = amfInstance.FindRanUeByAmfUeNgapID(aMFUENGAPID.Value)
+	ranUe := ran.FindUEByRanUeNgapID(rANUENGAPID.Value)
 	if ranUe == nil {
-		logger.WithTrace(ctx, ran.Log).Warn("UE Context not found", zap.Int64("AmfUeNgapID", aMFUENGAPID.Value))
+		logger.WithTrace(ctx, ran.Log).Error("No UE Context", zap.Int64("RanUeNgapID", rANUENGAPID.Value), zap.Int64("AmfUeNgapID", aMFUENGAPID.Value))
 		return
 	}
 

internal/amf/ngap/handle_ue_context_release_complete.go

@@ -193,9 +193,9 @@ func HandleUEContextReleaseComplete(ctx context.Context, amfInstance *amf.AMF, r
 			logger.WithTrace(ctx, ran.Log).Error(err.Error())
 		}
 
-		// Valid Security is not exist for this UE then only delete AMfUe Context
+		// No valid security context exists for this UE, so delete the AMF UE context
 		if !amfUe.SecurityContextAvailable {
-			logger.WithTrace(ctx, ran.Log).Info("Valid Security is not exist for the UE, so deleting AmfUe Context", logger.SUPI(amfUe.Supi.String()))
+			logger.WithTrace(ctx, ran.Log).Info("No valid security context for UE, deleting AMF UE context", logger.SUPI(amfUe.Supi.String()))
 			amfInstance.DeregisterAndRemoveAMFUE(ctx, amfUe)
 		}
 	case amf.UeContextReleaseDueToNwInitiatedDeregistraion:

internal/amf/ran_ue.go

@@ -346,7 +346,7 @@ func (ranUe *RanUe) UpdateLocation(ctx context.Context, amf *AMF, userLocationIn
 
 		if ranUe.amfUe != nil {
 			ranUe.amfUe.Location = ranUe.Location
-			ranUe.amfUe.Tai = *ranUe.amfUe.Location.NrLocation.Tai
+			ranUe.amfUe.Tai = *ranUe.amfUe.Location.EutraLocation.Tai
 		}
 	case ngapType.UserLocationInformationPresentUserLocationInformationNR:
 		locationInfoNR := userLocationInformation.UserLocationInformationNR

internal/amf/status.go

@@ -25,7 +25,14 @@ func (amf *AMF) RadioNameForSubscriber(supi etsi.SUPI) string {
 	defer amf.mu.RUnlock()
 
 	ue, ok := amf.UEs[supi]
-	if !ok || ue.GetState() != Registered || ue.ranUe == nil || ue.ranUe.Radio == nil {
+	if !ok {
+		return ""
+	}
+
+	ue.Mutex.Lock()
+	defer ue.Mutex.Unlock()
+
+	if ue.state != Registered || ue.ranUe == nil || ue.ranUe.Radio == nil {
 		return ""
 	}
 
@@ -65,19 +72,11 @@ func (amf *AMF) RegisteredSubscribersForRadio(radioName string) []string {
 	var imsis []string
 
 	for _, ue := range amf.UEs {
-		if ue.GetState() != Registered {
-			continue
-		}
-
-		if ue.ranUe == nil || ue.ranUe.Radio == nil {
-			continue
-		}
-
-		if ue.ranUe.Radio.Name != radioName {
-			continue
-		}
+		ue.Mutex.Lock()
+		match := ue.state == Registered && ue.ranUe != nil && ue.ranUe.Radio != nil && ue.ranUe.Radio.Name == radioName
+		ue.Mutex.Unlock()
 
-		if ue.Supi.IsValid() && ue.Supi.IsIMSI() {
+		if match && ue.Supi.IsValid() && ue.Supi.IsIMSI() {
 			imsis = append(imsis, ue.Supi.IMSI())
 		}
 	}