adding client and nas architecture

Cargo.toml

@@ -49,7 +49,6 @@ http-body-util = "0.1.2"
 tower-reqwest = "0.4.0"
 color-eyre = "0.6"
 tracing-error = "0.2.1"
-
 asn1_codecs_derive = "0.7.0"
 asn1-codecs = "0.7.0"
 smart-default = "0.7.1"
@@ -68,14 +67,13 @@ tokio-sctp = "0.2.0"
 valuable = "0.1.0"
 statig = {version = "0.3.0", features = ["async"]}
 non-empty-string = "0.2.6"
-
-# Depenedency of tokio-sctp
-socket2 = "0.4"
 derive-new = "0.7"
 faster-hex = "0.10.0"
 nonempty = { version = "0.8.1", features = ["serialize"] }
 bitvec = "1.0.1"
 ascii = "1.1.0"
+atomic_enum = "0.3.0"
+socket2 = "0.4" # Depenedency of tokio-sctp
 
 oasbi = { git = "https://github.com/UnifyAir/open-api.git/", package = "oasbi", branch = "master" }
 openapi-smf = { git = "https://github.com/UnifyAir/open-api.git/", package = "openapi-smf", features = [
@@ -96,4 +94,3 @@ openapi-nrf = { git = "https://github.com/UnifyAir/open-api.git/", package = "op
 ngap-models = { git = "https://github.com/UnifyAir/asn-models.git/", package = "ngap", branch = "master" }
 asn1-per = { git = "https://github.com/UnifyAir/asn-models.git/", package = "asn1-per", branch = "master" }
 nas-models = { git = "https://github.com/UnifyAir/nas-models.git/", package = "nas-models", branch = "master" }
-

lightning-nf/omnipath/app/Cargo.toml

@@ -57,8 +57,10 @@ valuable.workspace = true
 ascii.workspace = true
 non-empty-string.workspace = true
 statig.workspace = true
+atomic_enum.workspace = true
 
 counter = { path = "../../../utils/counter" }
 client = { path = "../../../utils/client" }
 nf-base = { path = "../../../utils/nf-base" }
+atomic-handle = {path = "../../../utils/atomic-handle"}
 

lightning-nf/omnipath/app/src/context/gnb_context.rs

@@ -2,13 +2,16 @@ use std::sync::Arc;
 
 use counter::CounterU64;
 use derive_new::new;
-use ngap_models::{GlobalRanNodeId, PagingDrx};
+use ngap_models::{FiveGSTmsi, GlobalRanNodeId, PagingDrx, RanUeNgapId};
 use nonempty::NonEmpty;
 use oasbi::common::{Snssai, Tai};
+use scc::hash_map::HashMap as SccHashMap;
 use tokio_util::sync::CancellationToken;
 
+use atomic_handle::AtomicHandle;
+
 use crate::{
-	context::ue_context::UeContext,
+	context::ue_context::{GmmStateField, UeContext},
 	ngap::{manager::ContextManager, network::TnlaAssociation},
 };
 
@@ -22,6 +25,13 @@ pub struct GnbContext {
 	#[new(value = "ContextManager::new()")]
 	pub ue_context_manager: ContextManager<UeContext>,
 
+	// List of registered ues who might be paged later
+	#[new(default)]
+	pub idle_ues: SccHashMap<FiveGSTmsi, UeContext>,
+
+	#[new(default)]
+	pub ue_states: SccHashMap<RanUeNgapId, AtomicHandle<UeContext, GmmStateField>>,
+
 	#[new(default)]
 	pub name: String,
 

lightning-nf/omnipath/app/src/context/mod.rs

@@ -1,9 +1,13 @@
 pub mod app_context;
 mod gnb_context;
 mod ngap_context;
+mod state;
 mod ue_context;
+mod nas_context;
 
 pub use app_context::AppContext;
 pub use gnb_context::{GnbContext, SupportedTai};
 pub use ngap_context::NgapContext;
+pub use state::{AtomicGmmState, GmmState};
 pub use ue_context::UeContext;
+pub use nas_context::NasContext;

lightning-nf/omnipath/app/src/context/nas_context.rs

@@ -0,0 +1,25 @@
+use std::{cell::UnsafeCell, num::NonZeroU32};
+
+use derive_new::new;
+use ngap_models::RrcEstablishmentCause;
+use non_empty_string::NonEmptyString;
+
+use crate::utils::models::FiveGSTmsi;
+
+#[derive(new)]
+pub struct NasContext {
+	pub rrc_establishment_cause: RrcEstablishmentCause,
+	pub five_g_s_tmsi: Option<FiveGSTmsi>,
+	#[new(default)]
+	pub tmsi: Option<NonZeroU32>,
+	#[new(default)]
+	pub guti: Option<NonEmptyString>,
+	#[new(default)]
+	pub suci: Option<NonEmptyString>,
+	#[new(default)]
+	pub pei: Option<NonEmptyString>,
+	#[new(default)]
+	pub mac_addr: Option<NonEmptyString>,
+	#[new(default)]
+	pub plmn_id: Option<NonEmptyString>,
+}

lightning-nf/omnipath/app/src/context/state/atomic_gmm_state.rs

@@ -0,0 +1,218 @@
+use std::{
+	fmt, ops::Deref, sync::atomic::{AtomicUsize, Ordering}
+};
+
+use super::GmmState;
+
+/// A custom atomic wrapper for GmmState that performs operations directly on
+/// the enum's discriminant to avoid the potential overhead of `match`
+/// statements.
+///
+/// The safety of the `unsafe` transmutation relies on the invariant that this
+/// struct only ever stores values that are valid `GmmState` discriminants.
+#[derive(Debug)]
+pub struct AtomicGmmState {
+	state: AtomicUsize,
+}
+
+impl AtomicGmmState {
+	/// Creates a new AtomicGmmState with the given initial state.
+	pub const fn new(initial_state: GmmState) -> Self {
+		Self {
+			state: AtomicUsize::new(initial_state as usize),
+		}
+	}
+
+	/// Loads the current state using `std::mem::transmute`.
+	#[inline]
+	pub fn load(
+		&self,
+		ordering: Ordering,
+	) -> GmmState {
+		let discriminant = self.state.load(ordering);
+		// This assertion ensures that, in debug builds, we panic if the state
+		// ever holds an invalid discriminant, which would be undefined behavior.
+		debug_assert!(
+			discriminant <= GmmState::MAX_DISCRIMINANT as usize,
+			"Invalid GmmState discriminant: {}",
+			discriminant
+		);
+		// Safety: The `debug_assert` and disciplined use of the `store` methods
+		// ensure that `state` only contains valid discriminants for `GmmState`.
+		// `GmmState` is `#[repr(u8)]`, so transmutation from its discriminant is sound.
+		unsafe { std::mem::transmute(discriminant as u8) }
+	}
+
+	/// Stores a new state.
+	#[inline]
+	pub fn store(
+		&self,
+		new_state: GmmState,
+		ordering: Ordering,
+	) {
+		self.state.store(new_state as usize, ordering);
+	}
+
+	/// Atomically swaps the state and returns the previous state.
+	#[inline]
+	pub fn swap(
+		&self,
+		new_state: GmmState,
+		ordering: Ordering,
+	) -> GmmState {
+		let old_discriminant = self.state.swap(new_state as usize, ordering);
+		debug_assert!(
+			old_discriminant <= GmmState::MAX_DISCRIMINANT as usize,
+			"Invalid GmmState discriminant: {}",
+			old_discriminant
+		);
+		// Safety: Same justification as `load`.
+		unsafe { std::mem::transmute(old_discriminant as u8) }
+	}
+
+	/// A helper function to reduce code duplication in `compare_exchange`
+	/// methods.
+	#[inline]
+	fn transmute_result(result: Result<usize, usize>) -> Result<GmmState, GmmState> {
+		match result {
+			Ok(prev) => {
+				debug_assert!(prev <= GmmState::MAX_DISCRIMINANT as usize);
+				Ok(unsafe { std::mem::transmute(prev as u8) })
+			}
+			Err(actual) => {
+				debug_assert!(actual <= GmmState::MAX_DISCRIMINANT as usize);
+				Err(unsafe { std::mem::transmute(actual as u8) })
+			}
+		}
+	}
+
+	/// Atomically compares the current state with `current` and, if they match,
+	/// replaces it with `new`.
+	#[inline]
+	pub fn compare_exchange(
+		&self,
+		current: GmmState,
+		new: GmmState,
+		success: Ordering,
+		failure: Ordering,
+	) -> Result<GmmState, GmmState> {
+		let result = self
+			.state
+			.compare_exchange(current as usize, new as usize, success, failure);
+		Self::transmute_result(result)
+	}
+
+	/// Performs a weak compare-and-exchange operation.
+	#[inline]
+	pub fn compare_exchange_weak(
+		&self,
+		current: GmmState,
+		new: GmmState,
+		success: Ordering,
+		failure: Ordering,
+	) -> Result<GmmState, GmmState> {
+		let result =
+			self.state
+				.compare_exchange_weak(current as usize, new as usize, success, failure);
+		Self::transmute_result(result)
+	}
+
+	/// Atomically modifies the state with a given function.
+	#[inline]
+	pub fn fetch_update<F>(
+		&self,
+		set_order: Ordering,
+		fetch_order: Ordering,
+		mut f: F,
+	) -> Result<GmmState, GmmState>
+	where
+		F: FnMut(GmmState) -> Option<GmmState>,
+	{
+		let result = self
+			.state
+			.fetch_update(set_order, fetch_order, |discriminant| {
+				debug_assert!(discriminant <= GmmState::MAX_DISCRIMINANT as usize);
+				// Safety: Same justification as `load`.
+				let current_state = unsafe { std::mem::transmute(discriminant as u8) };
+				f(current_state).map(|new_state| new_state as usize)
+			});
+		Self::transmute_result(result)
+	}
+
+	// --- Convenience Methods ---
+
+	/// Gets the current state using `Acquire` ordering.
+	#[inline]
+	pub fn get(&self) -> GmmState {
+		self.load(Ordering::Acquire)
+	}
+
+	/// Sets the current state using `Release` ordering.
+	#[inline]
+	pub fn set(
+		&self,
+		state: GmmState,
+	) {
+		self.store(state, Ordering::Release);
+	}
+
+	/// Checks if the current state matches the given state.
+	#[inline]
+	pub fn is(
+		&self,
+		expected: GmmState,
+	) -> bool {
+		self.get() == expected
+	}
+
+	/// Checks if the UE can access services in the current state.
+	#[inline]
+	pub fn can_access_services(&self) -> bool {
+		self.get().can_access_services()
+	}
+
+	/// Checks if the current state is transitional.
+	#[inline]
+	pub fn is_transitional(&self) -> bool {
+		self.get().is_transitional()
+	}
+
+}
+
+impl Default for AtomicGmmState {
+	fn default() -> Self {
+		Self::new(GmmState::default())
+	}
+}
+
+impl Deref for AtomicGmmState {
+	type Target = AtomicUsize;
+
+	fn deref(&self) -> &Self::Target {
+		&self.state	
+	}
+}
+
+impl Clone for AtomicGmmState {
+	/// Clones the `AtomicGmmState` by creating a new atomic variable
+	/// initialized with the current state's value.
+	fn clone(&self) -> Self {
+		Self::new(self.get())
+	}
+}
+
+impl From<GmmState> for AtomicGmmState {
+	fn from(state: GmmState) -> Self {
+		Self::new(state)
+	}
+}
+
+impl fmt::Display for AtomicGmmState {
+	fn fmt(
+		&self,
+		f: &mut fmt::Formatter<'_>,
+	) -> fmt::Result {
+		// Delegate formatting to the underlying GmmState's Debug or Display impl.
+		write!(f, "{:?}", self.get())
+	}
+}