use std::collections::HashMap;
use std::sync::Arc;
use futures::future::BoxFuture;
use ntk_proto::domain::UnicastId;
use ntk_proto::v1::method_call::Call;
use ntk_proto::v1::{
Auth, CallerContext, ErrorDomain, MethodCall, RemoteError, ResponsePayload, TypedValue,
};
use ntk_rpc::RpcHandler;
use tokio::sync::RwLock;
#[derive(Debug)]
pub struct IdentityStack {
pub qspn: ntk_qspn::QspnRpcHandler,
pub peers: ntk_peerservices::PeersRpcHandler,
pub coordinator: ntk_coordinator::CoordinatorRpcHandler,
pub hooking: ntk_hooking::HookingRpcHandler,
}
pub struct Dispatcher {
neighborhood: ntk_neighborhood::NeighborhoodRpcHandler,
identity: ntk_identities::IdentityRpcHandler,
main_id: ntk_neighborhood::NodeId,
identity_stack: RwLock<Arc<IdentityStack>>,
secondary: RwLock<HashMap<ntk_neighborhood::NodeId, Arc<IdentityStack>>>,
}
impl Dispatcher {
#[must_use]
pub fn new(
neighborhood: ntk_neighborhood::NeighborhoodRpcHandler,
identity: ntk_identities::IdentityRpcHandler,
main_id: ntk_neighborhood::NodeId,
stack: IdentityStack,
) -> Self {
Self {
neighborhood,
identity,
main_id,
identity_stack: RwLock::new(Arc::new(stack)),
secondary: RwLock::new(HashMap::new()),
}
}
pub async fn replace_identity_stack(&self, stack: IdentityStack) {
*self.identity_stack.write().await = Arc::new(stack);
}
pub async fn register_identity(&self, id: ntk_neighborhood::NodeId, stack: IdentityStack) {
self.secondary.write().await.insert(id, Arc::new(stack));
}
pub async fn unregister_identity(&self, id: ntk_neighborhood::NodeId) {
self.secondary.write().await.remove(&id);
}
async fn resolve_stack(
&self,
unicast_id: &TypedValue,
) -> Result<Arc<IdentityStack>, RemoteError> {
match select_identity(unicast_id, self.main_id)? {
Selection::Main => Ok(self.identity_stack.read().await.clone()),
Selection::Secondary(id) => {
self.secondary
.read()
.await
.get(&id)
.cloned()
.ok_or_else(|| {
protocol_error(format!(
"Request.unicast_id: no locally-hosted identity with id {}",
id.get()
))
})
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Selection {
Main,
Secondary(ntk_neighborhood::NodeId),
}
fn select_identity(
unicast_id: &TypedValue,
main_id: ntk_neighborhood::NodeId,
) -> Result<Selection, RemoteError> {
match UnicastId::from_typed_value(unicast_id)
.map_err(|err| protocol_error(format!("Request.unicast_id: {err}")))?
{
UnicastId::MainIdentity => Ok(Selection::Main),
UnicastId::WholeNode(_) => Err(protocol_error(
"Request.unicast_id: WholeNode is not valid for an identity-scoped call",
)),
UnicastId::IdentityAware(id_tv) => {
let id = crate::node::registry::decode_caller_id(&id_tv).ok_or_else(|| {
protocol_error(
"Request.unicast_id: IdentityAware payload is not a valid identity id",
)
})?;
Ok(if id == main_id {
Selection::Main
} else {
Selection::Secondary(id)
})
}
}
}
impl std::fmt::Debug for Dispatcher {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Dispatcher").finish_non_exhaustive()
}
}
fn protocol_error(message: impl Into<String>) -> RemoteError {
RemoteError {
domain: ErrorDomain::Deserialize as i32,
message: message.into(),
}
}
fn misrouted() -> RemoteError {
protocol_error("MethodCall with no oneof arm set")
}
impl RpcHandler for Dispatcher {
fn handle<'a>(
&'a self,
caller: CallerContext,
unicast_id: TypedValue,
call: MethodCall,
auth: Option<Auth>,
) -> BoxFuture<'a, Result<ResponsePayload, RemoteError>> {
Box::pin(async move {
let Some(arm) = &call.call else {
return Err(misrouted());
};
match arm {
Call::NeighborhoodHereIAm(_)
| Call::NeighborhoodRequestArc(_)
| Call::NeighborhoodCanYouExport(_)
| Call::NeighborhoodRemoveArc(_)
| Call::NeighborhoodNop(_) => {
self.neighborhood
.handle(caller, unicast_id, call, auth)
.await
}
Call::IdentityMatchDuplication(_)
| Call::IdentityGetPeerMainId(_)
| Call::IdentityNotifyIdentityArcRemoved(_) => {
self.identity.handle(caller, unicast_id, call, auth).await
}
Call::QspnGetFullEtp(_)
| Call::QspnSendEtp(_)
| Call::QspnGotPrepareDestroy(_)
| Call::QspnGotDestroy(_) => {
let stack = self.resolve_stack(&unicast_id).await?;
stack.qspn.handle(caller, unicast_id, call, auth).await
}
Call::PeersForwardPeerMessage(_)
| Call::PeersGetRequest(_)
| Call::PeersSetResponse(_)
| Call::PeersSetRefuseMessage(_)
| Call::PeersSetRedoFromStart(_)
| Call::PeersSetNextDestination(_)
| Call::PeersSetFailure(_)
| Call::PeersSetNonParticipant(_)
| Call::PeersSetMissingOptionalMaps(_)
| Call::PeersSetParticipant(_)
| Call::PeersGiveParticipantMaps(_)
| Call::PeersAskParticipantMaps(_) => {
let stack = self.resolve_stack(&unicast_id).await?;
stack.peers.handle(caller, unicast_id, call, auth).await
}
Call::CoordinatorExecutePrepareMigration(_)
| Call::CoordinatorExecuteFinishMigration(_)
| Call::CoordinatorExecutePrepareEnter(_)
| Call::CoordinatorExecuteFinishEnter(_)
| Call::CoordinatorExecuteWeHaveSplitted(_) => {
let stack = self.resolve_stack(&unicast_id).await?;
stack
.coordinator
.handle(caller, unicast_id, call, auth)
.await
}
Call::HookingRetrieveNetworkData(_)
| Call::HookingSearchMigrationPath(_)
| Call::HookingRouteSearchRequest(_)
| Call::HookingRouteSearchError(_)
| Call::HookingRouteSearchResponse(_)
| Call::HookingRouteExploreRequest(_)
| Call::HookingRouteExploreResponse(_)
| Call::HookingRouteDeleteReserveRequest(_)
| Call::HookingRouteMigRequest(_)
| Call::HookingRouteMigResponse(_) => {
let stack = self.resolve_stack(&unicast_id).await?;
stack.hooking.handle(caller, unicast_id, call, auth).await
}
}
})
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use futures::future::BoxFuture;
use ntk_proto::v1::{
Auth, CallerContext, MethodCall, RemoteError, ResponsePayload, TypedValue,
};
use ntk_rpc::{FakeRpcClient, RpcClient, RpcHandler};
use tokio_util::sync::CancellationToken;
use super::{Dispatcher, IdentityStack, Selection, protocol_error, select_identity};
fn node_id(raw: i32) -> ntk_neighborhood::NodeId {
ntk_neighborhood::NodeId::from_raw(raw).unwrap()
}
#[test]
fn an_empty_unicast_id_selects_main() {
let main = node_id(1);
assert_eq!(
select_identity(&TypedValue::default(), main).unwrap(),
Selection::Main
);
}
#[test]
fn an_identity_aware_unicast_id_naming_main_selects_main() {
let main = node_id(1);
let tv = ntk_proto::domain::UnicastId::IdentityAware(
crate::node::registry::encode_caller_id(main),
)
.to_typed_value();
assert_eq!(select_identity(&tv, main).unwrap(), Selection::Main);
}
#[test]
fn an_identity_aware_unicast_id_naming_another_identity_selects_that_identity_not_main() {
let main = node_id(1);
let other = node_id(2);
let tv = ntk_proto::domain::UnicastId::IdentityAware(
crate::node::registry::encode_caller_id(other),
)
.to_typed_value();
assert_eq!(
select_identity(&tv, main).unwrap(),
Selection::Secondary(other)
);
}
#[test]
fn a_whole_node_unicast_id_is_rejected_on_an_identity_scoped_call() {
let main = node_id(1);
let tv =
ntk_proto::domain::UnicastId::WholeNode(crate::node::registry::encode_caller_id(main))
.to_typed_value();
assert!(select_identity(&tv, main).is_err());
}
struct NeverCalledHandler;
impl RpcHandler for NeverCalledHandler {
fn handle<'a>(
&'a self,
_caller: CallerContext,
_unicast_id: TypedValue,
_call: MethodCall,
_auth: Option<Auth>,
) -> BoxFuture<'a, Result<ResponsePayload, RemoteError>> {
Box::pin(async { Err(protocol_error("test fixture: never actually called")) })
}
}
fn never_called_client() -> Arc<dyn RpcClient> {
Arc::new(FakeRpcClient::new(Arc::new(NeverCalledHandler)))
}
#[derive(Debug)]
struct NoRealArcs;
impl ntk_neighborhood::NeighborhoodStubFactory for NoRealArcs {
fn broadcast(&self, _dev: &str) -> Arc<dyn RpcClient> {
never_called_client()
}
fn unicast(&self, _arc: &ntk_neighborhood::Arc) -> Arc<dyn RpcClient> {
never_called_client()
}
}
impl ntk_identities::IdentityStubFactory for NoRealArcs {
fn stub(&self, _arc: ntk_identities::ArcId) -> Arc<dyn RpcClient> {
never_called_client()
}
fn arc_for_caller(&self, _caller: &CallerContext) -> Option<ntk_identities::ArcId> {
None
}
}
struct NoArcResolver;
impl ntk_qspn::ArcResolver for NoArcResolver {
fn resolve(&self, _caller: &CallerContext) -> Option<ntk_qspn::ArcId> {
None
}
}
struct NoopRoutingEnv;
impl ntk_peerservices::RoutingEnv for NoopRoutingEnv {
fn gnode_exists(&self, _hc: ntk_common::HCoord) -> bool {
false
}
fn gateway(
&self,
_hc: ntk_common::HCoord,
_failed: Option<&Arc<dyn ntk_peerservices::PeersStub>>,
) -> Option<Arc<dyn ntk_peerservices::PeersStub>> {
None
}
fn dial(
&self,
_n: &ntk_peerservices::TupleNode,
) -> Option<Arc<dyn ntk_peerservices::PeersStub>> {
None
}
fn nodes_in_my_group(&self, _level: usize) -> usize {
0
}
fn neighbors(&self) -> Vec<Arc<dyn ntk_peerservices::PeersStub>> {
Vec::new()
}
}
struct NoopCoordinatorMap;
impl ntk_coordinator::CoordinatorMap for NoopCoordinatorMap {
fn n_nodes(&self) -> u64 {
1
}
fn free_positions(&self, _level: usize) -> Vec<u32> {
Vec::new()
}
fn can_reserve(&self, _level: usize) -> bool {
false
}
fn my_pos(&self, _level: usize) -> u32 {
0
}
fn fp_id(&self, _level: usize) -> i64 {
0
}
}
struct NoopPropagationHandler;
impl ntk_coordinator::PropagationHandler for NoopPropagationHandler {
fn prepare_migration(&self, _level: usize, _data: TypedValue) -> BoxFuture<'_, ()> {
Box::pin(async {})
}
fn finish_migration(&self, _level: usize, _data: TypedValue) -> BoxFuture<'_, ()> {
Box::pin(async {})
}
fn prepare_enter(&self, _level: usize, _data: TypedValue) -> BoxFuture<'_, ()> {
Box::pin(async {})
}
fn finish_enter(&self, _level: usize, _data: TypedValue) -> BoxFuture<'_, ()> {
Box::pin(async {})
}
fn we_have_splitted(&self, _level: usize, _data: TypedValue) -> BoxFuture<'_, ()> {
Box::pin(async {})
}
}
struct NoopEnterHandler;
impl ntk_coordinator::EvaluateEnterHandler for NoopEnterHandler {
fn evaluate_enter<'a>(
&'a self,
_top: usize,
_data: TypedValue,
_client_tuple: &'a [u32],
) -> BoxFuture<'a, TypedValue> {
Box::pin(async { TypedValue::default() })
}
}
impl ntk_coordinator::BeginEnterHandler for NoopEnterHandler {
fn begin_enter<'a>(
&'a self,
_top: usize,
_data: TypedValue,
_client_tuple: &'a [u32],
) -> BoxFuture<'a, TypedValue> {
Box::pin(async { TypedValue::default() })
}
}
impl ntk_coordinator::CompletedEnterHandler for NoopEnterHandler {
fn completed_enter<'a>(
&'a self,
_top: usize,
_data: TypedValue,
_client_tuple: &'a [u32],
) -> BoxFuture<'a, TypedValue> {
Box::pin(async { TypedValue::default() })
}
}
impl ntk_coordinator::AbortEnterHandler for NoopEnterHandler {
fn abort_enter<'a>(
&'a self,
_top: usize,
_data: TypedValue,
_client_tuple: &'a [u32],
) -> BoxFuture<'a, TypedValue> {
Box::pin(async { TypedValue::default() })
}
}
fn spawn_identity_stack(pos: u32) -> IdentityStack {
let topology = ntk_common::Topology::new([4u32]).unwrap();
let my_naddr = ntk_common::Naddr::new(topology.clone(), vec![pos]).unwrap();
let fingerprint = ntk_common::Fingerprint::new(vec![1u8], 0, vec![0]);
let (qspn_handle, _qspn_join) = ntk_qspn::spawn(
my_naddr.clone(),
fingerprint,
ntk_qspn::QspnConfig::default(),
Arc::new(ntk_qspn::FakeQspnStubFactory::new()),
Arc::new(ntk_qspn::FixedThreshold::default()),
Arc::new(ntk_qspn::DefaultArcIdSource::default()),
CancellationToken::new(),
);
let qspn = ntk_qspn::QspnRpcHandler::new(
qspn_handle,
Arc::new(NoArcResolver),
Duration::from_secs(1),
Duration::from_millis(1),
);
let (peers_manager, peers_handle) = ntk_peerservices::Manager::new(
topology.clone(),
my_naddr.clone(),
Arc::new(NoopRoutingEnv),
ntk_peerservices::Config::default(),
topology.levels(),
);
tokio::spawn(peers_manager.run(CancellationToken::new()));
let peers = ntk_peerservices::PeersRpcHandler::new(peers_handle);
let enter = Arc::new(NoopEnterHandler);
let (coordinator_manager, coordinator_handle) = ntk_coordinator::Manager::new(
topology.clone(),
Arc::new(NoopCoordinatorMap),
Arc::new(ntk_coordinator::FakeCoordinatorStubFactory::new(Vec::new())),
Arc::new(NoopPropagationHandler),
ntk_coordinator::EnterHandlers {
evaluate_enter: enter.clone(),
begin_enter: enter.clone(),
completed_enter: enter.clone(),
abort_enter: enter,
},
ntk_coordinator::Config::default(),
None,
);
tokio::spawn(coordinator_manager.run(CancellationToken::new()));
let coordinator = ntk_coordinator::CoordinatorRpcHandler::new(coordinator_handle);
let view = Arc::new(ntk_hooking::FakeQspnView::new(topology, vec![pos]));
let coord_client = Arc::new(ntk_hooking::FakeCoordinatorClient::default());
let router = Arc::new(ntk_hooking::MessageRouting::new(
view.clone(),
coord_client.clone(),
Arc::new(ntk_hooking::FakeHookingStubFactory::default()),
Duration::from_secs(1),
));
let hooking = ntk_hooking::HookingRpcHandler::new(view, coord_client, router);
IdentityStack {
qspn,
peers,
coordinator,
hooking,
}
}
fn spawn_neighborhood_rpc(
my_id: ntk_neighborhood::NodeId,
) -> ntk_neighborhood::NeighborhoodRpcHandler {
let (handle, _join) = ntk_neighborhood::Manager::spawn(
ntk_neighborhood::NeighborhoodConfig {
my_id,
max_arcs: 8,
kernel: ntk_netlink::FakeNetlink::new(),
stub_factory: Arc::new(NoRealArcs),
ip_route_manager: Arc::new(ntk_neighborhood::FakeIpRouteManager::new()),
rtt_probe: Arc::new(ntk_neighborhood::FixedRttProbe(None)),
timing: ntk_neighborhood::NeighborhoodTiming::default(),
new_linklocal_address: Box::new(|| "10.0.0.1".to_owned()),
signing_key: None,
require_auth: false,
},
CancellationToken::new(),
);
ntk_neighborhood::NeighborhoodRpcHandler::for_unicast(handle)
}
fn spawn_identity_rpc() -> ntk_identities::IdentityRpcHandler {
let (handle, _join) =
ntk_identities::Handle::spawn(None, Arc::new(NoRealArcs), CancellationToken::new());
ntk_identities::IdentityRpcHandler::new(handle, Arc::new(NoRealArcs))
}
fn dispatcher(main_id: ntk_neighborhood::NodeId) -> Dispatcher {
Dispatcher::new(
spawn_neighborhood_rpc(main_id),
spawn_identity_rpc(),
main_id,
spawn_identity_stack(0),
)
}
#[tokio::test]
async fn an_empty_unicast_id_resolves_the_main_stack() {
let main_id = node_id(1);
let d = dispatcher(main_id);
let main_stack = d.identity_stack.read().await.clone();
let resolved = d.resolve_stack(&TypedValue::default()).await.unwrap();
assert!(Arc::ptr_eq(&resolved, &main_stack));
}
#[tokio::test]
async fn a_registered_second_identity_resolves_to_its_own_stack_not_main() {
let main_id = node_id(1);
let other_id = node_id(2);
let d = dispatcher(main_id);
let main_stack = d.identity_stack.read().await.clone();
d.register_identity(other_id, spawn_identity_stack(1)).await;
let secondary_stack = d.secondary.read().await.get(&other_id).cloned().unwrap();
let tv = ntk_proto::domain::UnicastId::IdentityAware(
crate::node::registry::encode_caller_id(other_id),
)
.to_typed_value();
let resolved = d.resolve_stack(&tv).await.unwrap();
assert!(Arc::ptr_eq(&resolved, &secondary_stack));
assert!(!Arc::ptr_eq(&resolved, &main_stack));
}
#[tokio::test]
async fn an_identity_aware_unicast_id_naming_an_unregistered_identity_is_rejected() {
let main_id = node_id(1);
let unknown_id = node_id(99);
let d = dispatcher(main_id);
let tv = ntk_proto::domain::UnicastId::IdentityAware(
crate::node::registry::encode_caller_id(unknown_id),
)
.to_typed_value();
assert!(d.resolve_stack(&tv).await.is_err());
}
#[tokio::test]
async fn unregistering_a_second_identity_makes_it_unknown_again() {
let main_id = node_id(1);
let other_id = node_id(2);
let d = dispatcher(main_id);
d.register_identity(other_id, spawn_identity_stack(1)).await;
d.unregister_identity(other_id).await;
let tv = ntk_proto::domain::UnicastId::IdentityAware(
crate::node::registry::encode_caller_id(other_id),
)
.to_typed_value();
assert!(d.resolve_stack(&tv).await.is_err());
}
}