Allow more than 1 PDU Session per UE on each slice

internal/amf/n2_establishment_response.go

@@ -28,7 +28,7 @@ func (amf *Amf) N2EstablishmentResponse(c *gin.Context) {
 
 func (amf *Amf) HandleN2EstablishmentResponse(ps n1n2.N2PduSessionRespMsg) {
 	ctx := amf.Context()
-	pduSession, err := amf.smf.CreateSessionDownlinkContext(ctx, ps.UeInfo.Header.Ue, ps.UeInfo.Header.Dnn, ps.Gnb, ps.DownlinkTeid)
+	pduSession, err := amf.smf.CreateSessionDownlinkContext(ctx, ps.UeInfo.Header.Ue, ps.UeInfo.Addr, ps.UeInfo.Header.Dnn, ps.Gnb, ps.DownlinkTeid)
 	if err != nil {
 		logrus.WithError(err).WithFields(logrus.Fields{
 			"ue-ip-addr": ps.UeInfo.Addr,

internal/smf/sessions.go

@@ -6,9 +6,47 @@
 package smf
 
 import (
+	"net/netip"
 	"sync"
+
+	"github.com/nextmn/json-api/jsonapi"
 )
 
+type Sessions struct {
+	s []*PduSessionN3
+}
+
 type SessionsMap struct {
-	sync.Map // key: UE Ctrl ; value: *PduSessionN3
+	m map[jsonapi.ControlURI]*Sessions
+	sync.RWMutex
+}
+
+func NewSessionsMap() *SessionsMap {
+	return &SessionsMap{
+		m: make(map[jsonapi.ControlURI]*Sessions),
+	}
+}
+
+func (s *SessionsMap) Get(ueCtrl jsonapi.ControlURI, ueAddr netip.Addr) (*PduSessionN3, error) {
+	s.RLock()
+	defer s.RUnlock()
+	if sessions, ok := s.m[ueCtrl]; ok {
+		for _, session := range sessions.s {
+			if session.UeIpAddr == ueAddr {
+				return session, nil
+			}
+		}
+	}
+	return nil, ErrPDUSessionNotFound
+}
+
+func (s *SessionsMap) Add(ueCtrl jsonapi.ControlURI, session *PduSessionN3) {
+	s.Lock()
+	defer s.Unlock()
+	m, ok := s.m[ueCtrl]
+	if !ok {
+		s.m[ueCtrl] = &Sessions{s: []*PduSessionN3{session}}
+	} else {
+		m.s = append(m.s, session)
+	}
 }

internal/smf/slice.go

@@ -39,6 +39,6 @@ func NewSlice(pool netip.Prefix, upfs []netip.Addr) *Slice {
 	return &Slice{
 		Pool:     NewUeIpPool(pool),
 		Upfs:     upfs,
-		sessions: &SessionsMap{},
+		sessions: NewSessionsMap(),
 	}
 }

internal/smf/smf.go

@@ -101,11 +101,11 @@ func (smf *Smf) Context() context.Context {
 	return context.Background()
 }
 
-func (smf *Smf) CreateSessionDownlink(ueCtrl jsonapi.ControlURI, dnn string, gnb netip.Addr, gnb_teid uint32) (*PduSessionN3, error) {
-	return smf.CreateSessionDownlinkContext(smf.ctx, ueCtrl, dnn, gnb, gnb_teid)
+func (smf *Smf) CreateSessionDownlink(ueCtrl jsonapi.ControlURI, ueIp netip.Addr, dnn string, gnb netip.Addr, gnb_teid uint32) (*PduSessionN3, error) {
+	return smf.CreateSessionDownlinkContext(smf.ctx, ueCtrl, ueIp, dnn, gnb, gnb_teid)
 }
 
-func (smf *Smf) CreateSessionDownlinkContext(ctx context.Context, ueCtrl jsonapi.ControlURI, dnn string, gnb netip.Addr, gnb_teid uint32) (*PduSessionN3, error) {
+func (smf *Smf) CreateSessionDownlinkContext(ctx context.Context, ueCtrl jsonapi.ControlURI, ueIp netip.Addr, dnn string, gnb netip.Addr, gnb_teid uint32) (*PduSessionN3, error) {
 	if !smf.started {
 		return nil, ErrSmfNotStarted
 	}
@@ -127,11 +127,10 @@ func (smf *Smf) CreateSessionDownlinkContext(ctx context.Context, ueCtrl jsonapi
 		return nil, ErrDnnNotFound
 	}
 	slice := s.(*Slice)
-	session_any, ok := slice.sessions.Load(ueCtrl)
-	if !ok {
-		return nil, ErrPDUSessionNotFound
+	session, err := slice.sessions.Get(ueCtrl, ueIp)
+	if err != nil {
+		return nil, err
 	}
-	session := session_any.(*PduSessionN3)
 	session.DownlinkFteid = &Fteid{
 		Addr: gnb,
 		Teid: gnb_teid,
@@ -197,10 +196,6 @@ func (smf *Smf) CreateSessionUplinkContext(ctx context.Context, ueCtrl jsonapi.C
 		return nil, ErrDnnNotFound
 	}
 	slice := s.(*Slice)
-	_, ok = slice.sessions.Load(ueCtrl)
-	if ok {
-		return nil, ErrPDUSessionAlreadyExists
-	}
 	// create ue ip addr
 	ueIpAddr, err := slice.Pool.Next()
 	if err != nil {
@@ -294,7 +289,7 @@ func (smf *Smf) CreateSessionUplinkContext(ctx context.Context, ueCtrl jsonapi.C
 		UplinkFteid: last_fteid,
 	}
 	// store session
-	slice.sessions.Store(ueCtrl, &session)
+	slice.sessions.Add(ueCtrl, &session)
 	return &session, nil
 }