fix: free old GUTI after Registration Complete and other quality fixes in AMF (#1174)

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

Author: gruyaume

internal/amf/amf.go

@@ -31,7 +31,7 @@ import (
 // Authenticator is the interface the AMF requires from the AUSF.
 // *ausf.AUSF satisfies this interface directly.
 type Authenticator interface {
-	Authenticate(ctx context.Context, suci, servingNetwork string, resync *ausf.ResyncInfo) (*ausf.AuthResult, error)
+	Authenticate(ctx context.Context, suci string, plmn models.PlmnID, resync *ausf.ResyncInfo) (*ausf.AuthResult, error)
 	Confirm(ctx context.Context, resStar, suci string) (etsi.SUPI, string, error)
 }
 
@@ -241,10 +241,11 @@ func (amf *AMF) NewRadio(conn *sctp.SCTPConn) (*Radio, error) {
 		SupportedTAIs: make([]SupportedTAI, 0),
 		Conn:          conn,
 		ConnectedAt:   now,
-		LastSeenAt:    now,
 		Log:           logger.AmfLog.With(logger.RanAddr(remoteAddr.String())),
 	}
 
+	radio.SetLastSeenAt(now)
+
 	amf.mu.Lock()
 	defer amf.mu.Unlock()
 
@@ -524,7 +525,7 @@ func (amf *AMF) SendPaging(ctx context.Context, ue *AmfUe, ngapBuf []byte) error
 			for _, ran := range amf.ListRadios() {
 				for _, item := range ran.SupportedTAIs {
 					if InTaiList(item.Tai, taiList) {
-						err := ran.NGAPSender.SendToRan(ctx, ngapBuf, send.NGAPProcedurePaging)
+						err := ran.NGAPSender.SendToRan(context.Background(), ngapBuf, send.NGAPProcedurePaging)
 						if err != nil {
 							ue.Log.Error("failed to send paging", zap.Error(err))
 							continue

internal/amf/amf_ran.go

@@ -10,6 +10,7 @@ package amf
 import (
 	"context"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/ellanetworks/core/internal/amf/ngap/send"
@@ -64,7 +65,7 @@ type Radio struct {
 	Name          string
 	Conn          *sctp.SCTPConn
 	ConnectedAt   time.Time
-	LastSeenAt    time.Time
+	lastSeenAt    atomic.Int64 // Unix nanoseconds; use GetLastSeenAt()/TouchLastSeen()
 	SupportedTAIs []SupportedTAI
 	mu            sync.RWMutex     // protects RanUEs
 	RanUEs        map[int64]*RanUe // Key: RanUeNgapID
@@ -127,7 +128,22 @@ func (r *Radio) SetRanID(ranNodeID *ngapType.GlobalRANNodeID) {
 }
 
 func (r *Radio) TouchLastSeen() {
-	r.LastSeenAt = time.Now()
+	r.lastSeenAt.Store(time.Now().UnixNano())
+}
+
+// GetLastSeenAt returns the last-seen timestamp. Safe for concurrent use.
+func (r *Radio) GetLastSeenAt() time.Time {
+	ns := r.lastSeenAt.Load()
+	if ns == 0 {
+		return time.Time{}
+	}
+
+	return time.Unix(0, ns)
+}
+
+// SetLastSeenAt sets the last-seen timestamp. Safe for concurrent use.
+func (r *Radio) SetLastSeenAt(t time.Time) {
+	r.lastSeenAt.Store(t.UnixNano())
 }
 
 // NodeID returns the RAN node identifier string regardless of radio type.

internal/amf/amf_ran_test.go

@@ -35,18 +35,18 @@ func TestRadioRanNodeTypeName(t *testing.T) {
 }
 
 func TestRadioTouchLastSeen(t *testing.T) {
-	radio := &amf.Radio{
-		LastSeenAt: time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC),
-	}
+	radio := &amf.Radio{}
+	radio.SetLastSeenAt(time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC))
 
 	before := time.Now()
 
 	radio.TouchLastSeen()
 
 	after := time.Now()
 
-	if radio.LastSeenAt.Before(before) || radio.LastSeenAt.After(after) {
-		t.Fatalf("expected LastSeenAt between %v and %v, got %v", before, after, radio.LastSeenAt)
+	lastSeen := radio.GetLastSeenAt()
+	if lastSeen.Before(before) || lastSeen.After(after) {
+		t.Fatalf("expected LastSeenAt between %v and %v, got %v", before, after, lastSeen)
 	}
 }
 
@@ -57,23 +57,25 @@ func TestRadioTimestampsSetOnCreation(t *testing.T) {
 		t.Fatal("expected ConnectedAt to be zero on a blank Radio")
 	}
 
-	if !blank.LastSeenAt.IsZero() {
+	if !blank.GetLastSeenAt().IsZero() {
 		t.Fatal("expected LastSeenAt to be zero on a blank Radio")
 	}
 
 	now := time.Now()
 
 	radio := &amf.Radio{
 		ConnectedAt: now,
-		LastSeenAt:  now,
 	}
 
+	radio.SetLastSeenAt(now)
+
 	if radio.ConnectedAt.IsZero() || radio.ConnectedAt != now {
 		t.Fatalf("expected ConnectedAt to be %v, got %v", now, radio.ConnectedAt)
 	}
 
-	if radio.LastSeenAt.IsZero() || radio.LastSeenAt != now {
-		t.Fatalf("expected LastSeenAt to be %v, got %v", now, radio.LastSeenAt)
+	lastSeen := radio.GetLastSeenAt()
+	if lastSeen.IsZero() || !lastSeen.Equal(now) {
+		t.Fatalf("expected LastSeenAt to be %v, got %v", now, lastSeen)
 	}
 }
 

internal/amf/amf_ue.go

@@ -13,7 +13,6 @@ import (
 	"encoding/binary"
 	"encoding/hex"
 	"fmt"
-	"reflect"
 	"slices"
 	"sync"
 	"time"
@@ -275,15 +274,15 @@ func (ue *AmfUe) AllocateRegistrationArea(supportedTais []models.Tai) {
 	copy(taiList, supportedTais)
 
 	for _, supportTai := range taiList {
-		if reflect.DeepEqual(supportTai, ue.Tai) {
+		if supportTai.Equal(ue.Tai) {
 			ue.RegistrationArea = append(ue.RegistrationArea, supportTai)
 			break
 		}
 	}
 }
 
 func (ue *AmfUe) IsAllowedNssai(targetSNssai *models.Snssai) bool {
-	return reflect.DeepEqual(*ue.AllowedNssai, *targetSNssai)
+	return ue.AllowedNssai.Equal(*targetSNssai)
 }
 
 func (ue *AmfUe) SecurityContextIsValid() bool {
@@ -584,6 +583,9 @@ func (ue *AmfUe) CreateSmContext(pduSessionID uint8, ref string, snssai *models.
 		return fmt.Errorf("invalid PDU session ID %d: must be in range 1-15 per TS 24.501", pduSessionID)
 	}
 
+	ue.Mutex.Lock()
+	defer ue.Mutex.Unlock()
+
 	ue.SmContextList[pduSessionID] = &SmContext{
 		Ref:    ref,
 		Snssai: snssai,
@@ -592,12 +594,35 @@ func (ue *AmfUe) CreateSmContext(pduSessionID uint8, ref string, snssai *models.
 	return nil
 }
 
+func (ue *AmfUe) DeleteSmContext(pduSessionID uint8) {
+	ue.Mutex.Lock()
+	defer ue.Mutex.Unlock()
+
+	delete(ue.SmContextList, pduSessionID)
+}
+
 func (ue *AmfUe) SmContextFindByPDUSessionID(pduSessionID uint8) (*SmContext, bool) {
+	ue.Mutex.Lock()
+	defer ue.Mutex.Unlock()
+
 	smContext, ok := ue.SmContextList[pduSessionID]
+
 	return smContext, ok
 }
 
+func (ue *AmfUe) SetSmContextInactive(pduSessionID uint8) {
+	ue.Mutex.Lock()
+	defer ue.Mutex.Unlock()
+
+	if sc, ok := ue.SmContextList[pduSessionID]; ok {
+		sc.PduSessionInactive = true
+	}
+}
+
 func (ue *AmfUe) HasActivePduSessions() bool {
+	ue.Mutex.Lock()
+	defer ue.Mutex.Unlock()
+
 	for _, smContext := range ue.SmContextList {
 		if !smContext.PduSessionInactive {
 			return true
@@ -933,8 +958,19 @@ func (ue *AmfUe) releaseSmContexts(ctx context.Context) {
 		return
 	}
 
+	// Copy refs under lock, then release lock before external SMF calls.
+	ue.Mutex.Lock()
+
+	smContextRefs := make([]string, 0, len(ue.SmContextList))
 	for _, smContext := range ue.SmContextList {
-		err := ue.smf.ReleaseSmContext(ctx, smContext.Ref)
+		smContextRefs = append(smContextRefs, smContext.Ref)
+	}
+
+	ue.SmContextList = make(map[uint8]*SmContext)
+	ue.Mutex.Unlock()
+
+	for _, smContextRef := range smContextRefs {
+		err := ue.smf.ReleaseSmContext(ctx, smContextRef)
 		if err != nil {
 			ue.Log.Error("Release SmContext Error", zap.Error(err))
 		}

internal/amf/nas/gmm/authentication_procedure.go

@@ -3,7 +3,6 @@ package gmm
 import (
 	"context"
 	"fmt"
-	"strconv"
 
 	"github.com/ellanetworks/core/internal/amf"
 	"github.com/ellanetworks/core/internal/amf/nas/gmm/message"
@@ -16,14 +15,7 @@ func sendUEAuthenticationAuthenticateRequest(ctx context.Context, amfInstance *a
 		return nil, fmt.Errorf("tai is not available in UE context")
 	}
 
-	mnc, err := strconv.Atoi(ue.Tai.PlmnID.Mnc)
-	if err != nil {
-		return nil, fmt.Errorf("could not convert mnc to int: %v", err)
-	}
-
-	snName := fmt.Sprintf("5G:mnc%03d.mcc%s.3gppnetwork.org", mnc, ue.Tai.PlmnID.Mcc)
-
-	ueAuthenticationCtx, err := amfInstance.Ausf.Authenticate(ctx, ue.Suci, snName, resyncInfo)
+	ueAuthenticationCtx, err := amfInstance.Ausf.Authenticate(ctx, ue.Suci, *ue.Tai.PlmnID, resyncInfo)
 	if err != nil {
 		return nil, fmt.Errorf("ausf UE Authentication Authenticate Request failed: %s", err.Error())
 	}

internal/amf/nas/gmm/handle_registration_complete.go

@@ -22,6 +22,9 @@ func handleRegistrationComplete(ctx context.Context, amfInstance *amf.AMF, ue *a
 		ue.T3550 = nil // clear the timer
 	}
 
+	// UE confirmed receipt of the new GUTI — free the old one (TS 24.501 5.5.1.2.4 step 20)
+	amfInstance.FreeOldGuti(ue)
+
 	// Send NITZ (network name + timezone) to UE per TS 24.501
 	message.SendConfigurationUpdateCommand(ctx, amfInstance, ue, false)
 

internal/amf/nas/gmm/handle_service_request.go

@@ -112,7 +112,24 @@ func sendServiceAccept(
 
 // TS 24501 5.6.1
 func handleServiceRequest(ctx context.Context, amfInstance *amf.AMF, ue *amf.AmfUe, msg *nasMessage.ServiceRequest) error {
-	if ue.GetState() != amf.Deregistered && ue.GetState() != amf.Registered {
+	// TS 24.501 5.6.1.1: reject service request from deregistered UE
+	if ue.GetState() == amf.Deregistered {
+		err := message.SendServiceReject(ctx, ue.RanUe(), nasMessage.Cause5GMMUEIdentityCannotBeDerivedByTheNetwork)
+		if err != nil {
+			return fmt.Errorf("error sending service reject: %v", err)
+		}
+
+		ue.RanUe().ReleaseAction = amf.UeContextN2NormalRelease
+
+		err = ue.RanUe().SendUEContextReleaseCommand(ctx, ngapType.CausePresentNas, ngapType.CauseNasPresentNormalRelease)
+		if err != nil {
+			return fmt.Errorf("error sending ue context release command: %v", err)
+		}
+
+		return nil
+	}
+
+	if ue.GetState() != amf.Registered {
 		return fmt.Errorf("state mismatch: receive Service Request message in state %s", ue.GetState())
 	}
 
@@ -165,8 +182,8 @@ func handleServiceRequest(ctx context.Context, amfInstance *amf.AMF, ue *amf.Amf
 		ue.RetransmissionOfInitialNASMsg = ue.MacFailed
 	}
 
-	// Service Reject if the SecurityContext is invalid or the UE is Deregistered
-	if !ue.SecurityContextIsValid() || ue.GetState() == amf.Deregistered {
+	// Service Reject if the SecurityContext is invalid
+	if !ue.SecurityContextIsValid() {
 		ue.Log.Warn("No security context", logger.SUPI(ue.Supi.String()))
 		ue.SecurityContextAvailable = false
 
@@ -221,12 +238,21 @@ func handleServiceRequest(ctx context.Context, amfInstance *amf.AMF, ue *amf.Amf
 		}
 	}
 
+	// Copy SmContextList under lock for safe concurrent iteration.
+	ue.Mutex.Lock()
+
+	smContextSnapshot := make(map[uint8]*amf.SmContext, len(ue.SmContextList))
+	for id, sc := range ue.SmContextList {
+		smContextSnapshot[id] = sc
+	}
+	ue.Mutex.Unlock()
+
 	// If the UE has uplink data pending for some PDU sessions, we need to activate them
 	if msg.UplinkDataStatus != nil {
 		uplinkDataPsi := nasConvert.PSIToBooleanArray(msg.UplinkDataStatus.Buffer)
 		reactivationResult = new([16]bool)
 
-		for pduSessionID, smContext := range ue.SmContextList {
+		for pduSessionID, smContext := range smContextSnapshot {
 			if int(pduSessionID) >= len(uplinkDataPsi) {
 				ue.Log.Warn("Ignoring out-of-range PDU session ID in UplinkDataStatus processing", zap.Uint8("pduSessionID", pduSessionID))
 				continue
@@ -255,7 +281,7 @@ func handleServiceRequest(ctx context.Context, amfInstance *amf.AMF, ue *amf.Amf
 		acceptPduSessionPsi = new([16]bool)
 
 		psiArray := nasConvert.PSIToBooleanArray(msg.PDUSessionStatus.Buffer)
-		for pduSessionID, smContext := range ue.SmContextList {
+		for pduSessionID, smContext := range smContextSnapshot {
 			if int(pduSessionID) >= len(psiArray) {
 				ue.Log.Warn("Ignoring out-of-range PDU session ID in PDUSessionStatus processing", zap.Uint8("pduSessionID", pduSessionID))
 				continue

internal/amf/nas/gmm/handle_test.go

@@ -261,7 +261,7 @@ type FakeAusf struct {
 	AvKgAka *ausf.AuthResult
 }
 
-func (a *FakeAusf) Authenticate(ctx context.Context, suci string, servingNetwork string, resync *ausf.ResyncInfo) (*ausf.AuthResult, error) {
+func (a *FakeAusf) Authenticate(ctx context.Context, suci string, plmn models.PlmnID, resync *ausf.ResyncInfo) (*ausf.AuthResult, error) {
 	if a.Error != nil {
 		return nil, a.Error
 	}

internal/amf/nas/gmm/handle_ul_nas_transport.go

@@ -125,7 +125,7 @@ func transport5GSMMessage(ctx context.Context, amfInstance *amf.AMF, ue *amf.Amf
 	if smContextExist && requestType != nil {
 		/* AMF releases context locally as this is duplicate pdu session */
 		if requestType.GetRequestTypeValue() == nasMessage.ULNASTransportRequestTypeInitialRequest {
-			delete(ue.SmContextList, pduSessionID)
+			ue.DeleteSmContext(pduSessionID)
 
 			smContextExist = false
 		}

internal/amf/nas/gmm/message/send.go

@@ -110,7 +110,7 @@ func SendAuthenticationRequest(ctx context.Context, amfInstance *amf.AMF, ue *am
 		amfUe.T3560 = amf.NewTimer(cfg.ExpireTime, cfg.MaxRetryTimes, func(expireTimes int32) {
 			amfUe.Log.Warn("T3560 expires, retransmit Authentication Request", zap.Any("expireTimes", expireTimes))
 
-			err := ue.SendDownlinkNasTransport(ctx, nasMsg, nil)
+			err := ue.SendDownlinkNasTransport(context.Background(), nasMsg, nil)
 			if err != nil {
 				amfUe.Log.Error("could not send downlink NAS transport message", zap.Error(err))
 				return
@@ -268,7 +268,7 @@ func SendSecurityModeCommand(ctx context.Context, amfInstance *amf.AMF, ue *amf.
 		amfUe.T3560 = amf.NewTimer(cfg.ExpireTime, cfg.MaxRetryTimes, func(expireTimes int32) {
 			amfUe.Log.Warn("T3560 expires, retransmit Security Mode Command", zap.Any("expireTimes", expireTimes))
 
-			err = ue.SendDownlinkNasTransport(ctx, nasMsg, nil)
+			err = ue.SendDownlinkNasTransport(context.Background(), nasMsg, nil)
 			if err != nil {
 				amfUe.Log.Error("could not send downlink NAS transport message", zap.Error(err))
 				return
@@ -277,7 +277,6 @@ func SendSecurityModeCommand(ctx context.Context, amfInstance *amf.AMF, ue *amf.
 			amfUe.Log.Info("sent security mode command")
 		}, func() {
 			amfUe.Log.Warn("T3560 Expires, abort security mode control procedure", zap.Any("expireTimes", cfg.MaxRetryTimes))
-			// amfUe.Remove()
 			amfInstance.DeregisterAndRemoveAMFUE(context.Background(), amfUe)
 		})
 	}
@@ -380,7 +379,7 @@ func SendRegistrationAccept(
 			} else {
 				if ue.RanUe().UeContextRequest && !ue.RanUe().RecvdInitialContextSetupResponse {
 					err = ue.RanUe().SendInitialContextSetupRequest(
-						ctx,
+						context.Background(),
 						ue.Ambr.Uplink,
 						ue.Ambr.Downlink,
 						ue.AllowedNssai,
@@ -402,7 +401,7 @@ func SendRegistrationAccept(
 				} else {
 					ue.Log.Warn("T3550 expires, retransmit Registration Accept", zap.Any("expireTimes", expireTimes))
 
-					err = ue.RanUe().SendDownlinkNasTransport(ctx, nasMsg, nil)
+					err = ue.RanUe().SendDownlinkNasTransport(context.Background(), nasMsg, nil)
 					if err != nil {
 						ue.Log.Error("could not send downlink NAS transport message", zap.Error(err))
 					}
@@ -485,7 +484,7 @@ func SendConfigurationUpdateCommand(ctx context.Context, amfInstance *amf.AMF, a
 				return
 			}
 
-			err = amfUe.RanUe().SendDownlinkNasTransport(ctx, nasMsg, mobilityRestrictionList)
+			err = amfUe.RanUe().SendDownlinkNasTransport(context.Background(), nasMsg, mobilityRestrictionList)
 			if err != nil {
 				amfUe.Log.Error("could not send configuration update command", zap.Error(err))
 			}