1use crate::prelude::*;
66use crate::test_common::wait_for_peers;
67use citadel_io::tokio::sync::Mutex;
68use citadel_user::prelude::UserIdentifierExt;
69use futures::{Future, StreamExt};
70use std::marker::PhantomData;
71use std::pin::Pin;
72use std::sync::atomic::{AtomicBool, Ordering};
73use std::sync::Arc;
74use uuid::Uuid;
75
76pub struct BroadcastKernel<'a, F, Fut, R: Ratchet> {
82 inner_kernel: Box<dyn NetKernel<R> + 'a>,
83 shared: Arc<BroadcastShared>,
84 _pd: PhantomData<fn() -> (F, Fut)>,
85}
86
87pub struct BroadcastShared {
88 route_registers: AtomicBool,
89 register_rx:
90 citadel_io::Mutex<Option<citadel_io::tokio::sync::mpsc::UnboundedReceiver<PeerSignal>>>,
91 register_tx: citadel_io::tokio::sync::mpsc::UnboundedSender<PeerSignal>,
92}
93
94pub enum GroupInitRequestType {
102 Create {
107 local_user: UserIdentifier,
108 invite_list: Vec<UserIdentifier>,
109 group_id: Uuid,
110 accept_registrations: bool,
111 },
112 Join {
120 local_user: UserIdentifier,
121 owner: UserIdentifier,
122 group_id: Uuid,
123 do_peer_register: bool,
124 },
125}
126
127#[async_trait]
128impl<'a, F, Fut, R: Ratchet> PrefabFunctions<'a, GroupInitRequestType, R>
129 for BroadcastKernel<'a, F, Fut, R>
130where
131 F: FnOnce(GroupChannel, CitadelClientServerConnection<R>) -> Fut + Send + 'a,
132 Fut: Future<Output = Result<(), NetworkError>> + Send + 'a,
133{
134 type UserLevelInputFunction = F;
135 type SharedBundle = Arc<BroadcastShared>;
136
137 fn get_shared_bundle(&self) -> Self::SharedBundle {
138 self.shared.clone()
139 }
140
141 #[allow(unreachable_code, clippy::blocks_in_conditions)]
142 #[cfg_attr(
143 feature = "localhost-testing",
144 tracing::instrument(level = "trace", target = "citadel", skip_all, ret, err(Debug))
145 )]
146 async fn on_c2s_channel_received(
147 connect_success: CitadelClientServerConnection<R>,
148 arg: GroupInitRequestType,
149 fx: Self::UserLevelInputFunction,
150 shared: Arc<BroadcastShared>,
151 ) -> Result<(), NetworkError> {
152 let session_cid = connect_success.cid;
153 wait_for_peers().await;
154 let mut creator_only_accept_inbound_registers = false;
155
156 let mut is_owner = false;
157 let request = match arg {
158 GroupInitRequestType::Create {
159 local_user,
160 invite_list,
161 group_id,
162 accept_registrations,
163 } => {
164 is_owner = true;
165 let mut peers_registered = vec![];
167
168 for peer in &invite_list {
169 let peer = peer
170 .search_peer(session_cid, connect_success.account_manager())
171 .await?
172 .ok_or_else(|| {
173 citadel_io::error!(
174 citadel_io::ErrorCode::BroadcastCreateUserNotRegistered,
175 format!("{peer:?}"),
176 format!("{local_user:?}")
177 )
178 })?;
179
180 peers_registered.push(peer.cid)
181 }
182
183 creator_only_accept_inbound_registers = accept_registrations;
184
185 GroupBroadcast::Create {
186 initial_invitees: peers_registered,
187 options: MessageGroupOptions {
188 group_type: GroupType::Public,
189 id: group_id.as_u128(),
190 ..Default::default()
191 },
192 }
193 }
194
195 GroupInitRequestType::Join {
196 local_user,
197 owner,
198 group_id,
199 do_peer_register,
200 } => {
201 let owner_orig = owner;
203 let owner_find = owner_orig
204 .search_peer(session_cid, connect_success.account_manager())
205 .await?;
206
207 let owner = if let Some(owner) = owner_find {
208 Some(owner)
209 } else if do_peer_register {
210 let handle = connect_success
211 .propose_target(local_user.clone(), owner_orig.clone())
212 .await?;
213 let registration = handle.register_to_peer().await?;
217 if let Some(reason) = registration.refusal_reason() {
218 return Err(NetworkError::generic(format!(
219 "Cannot join the group: registering with its owner failed: {reason}"
220 )));
221 }
222 owner_orig
224 .search_peer(session_cid, connect_success.account_manager())
225 .await?
226 } else {
227 None
228 };
229
230 let owner = owner.ok_or_else(|| {
231 citadel_io::error!(
232 citadel_io::ErrorCode::BroadcastJoinUserNotRegistered,
233 format!("{owner_orig:?}"),
234 format!("{local_user:?}")
235 )
236 })?;
237
238 let expected_message_group_key = MessageGroupKey {
239 cid: owner.cid,
240 mgid: group_id.as_u128(),
241 };
242
243 crate::prefabs::client::broadcast_join::wait_for_owner_group(
244 &connect_success,
245 &local_user,
246 &owner,
247 group_id,
248 expected_message_group_key,
249 )
250 .await?;
251
252 GroupBroadcast::RequestJoin {
253 sender: local_user.get_cid(),
254 key: expected_message_group_key,
255 }
256 }
257 };
258
259 let request = NodeRequest::GroupBroadcastCommand(GroupBroadcastCommand {
260 session_cid,
261 command: request,
262 });
263
264 let subscription = &Mutex::new(Some(
265 connect_success.send_callback_subscription(request).await?,
266 ));
267
268 log::trace!(target: "citadel", "Peer {session_cid} is attempting to join group");
269 let acceptor_task = if creator_only_accept_inbound_registers {
270 shared.route_registers.store(true, Ordering::Relaxed);
271 let mut reg_rx = shared.register_rx.lock().take().unwrap();
272 let remote = connect_success.remote_ref().clone();
273 Box::pin(async move {
274 let mut subscription = subscription.lock().await.take().unwrap();
275 let mut count_registered = 0;
277 loop {
278 let post_register = citadel_io::tokio::select! {
279 reg_request = reg_rx.recv() => {
280 reg_request.ok_or_else(|| citadel_io::error!(citadel_io::ErrorCode::BroadcastStreamEndedUnexpectedly, "reg_rx"))?
281 },
282
283 reg_request2 = subscription.next() => {
284 let signal = reg_request2.ok_or_else(|| citadel_io::error!(citadel_io::ErrorCode::BroadcastStreamEndedUnexpectedly, "subscription"))?;
285 if let NodeResult::PeerEvent(PeerEvent { event: sig @ PeerSignal::PostRegister { .. }, .. }) = &signal {
286 sig.clone()
287 } else {
288 continue;
289 }
290 }
291 };
292
293 log::trace!(target: "citadel", "ACCEPTOR {session_cid} RECV reg_request: {post_register:?}");
294 if let PeerSignal::PostRegister {
295 peer_conn_type: peer_conn,
296 inviter_username: _,
297 invitee_username: _,
298 ticket_opt: _,
299 invitee_response: None,
300 } = &post_register
301 {
302 let cid = peer_conn.get_original_target_cid();
303 if cid != session_cid {
304 log::warn!(target: "citadel", "Received the wrong CID. Will not accept request");
305 continue;
306 }
307
308 let _ = responses::peer_register(post_register, true, &remote).await?;
309 if cfg!(feature = "localhost-testing") {
310 count_registered += 1;
311 if count_registered == crate::test_common::num_local_test_peers() - 1 {
312 break;
314 }
315 }
316 }
317 }
318
319 Ok::<_, NetworkError>(())
320 })
321 as Pin<
322 Box<
323 dyn futures::Future<
324 Output = Result<(), citadel_proto::prelude::NetworkError>,
325 > + Send,
326 >,
327 >
328 } else {
329 Box::pin(async move { Ok::<_, NetworkError>(()) })
330 as Pin<
331 Box<
332 dyn futures::Future<
333 Output = Result<(), citadel_proto::prelude::NetworkError>,
334 > + Send,
335 >,
336 >
337 };
338
339 let mut lock = subscription.lock().await;
340 let subscription = lock.as_mut().unwrap();
341 while let Some(event) = subscription.next().await {
342 match event.into_result()? {
343 NodeResult::PeerEvent(PeerEvent {
344 event: ref ps @ PeerSignal::PostRegister { .. },
345 ..
346 }) => {
347 shared
348 .register_tx
349 .send(ps.clone())
350 .map_err(|err| NetworkError::generic(err.to_string()))?;
351 }
352 NodeResult::GroupChannelCreated(GroupChannelCreated {
353 ticket: _,
354 channel,
355 session_cid: _,
356 }) => {
357 drop(lock);
360 return if is_owner {
361 citadel_io::tokio::try_join!(fx(channel, connect_success), acceptor_task)
362 .map(|_| ())
363 } else {
364 fx(channel, connect_success).await.map(|_| ())
365 };
366 }
367
368 NodeResult::GroupEvent(GroupEvent {
369 session_cid: _,
370 ticket: _,
371 event: GroupBroadcast::CreateResponse { key: None },
372 }) => {
373 return Err(citadel_io::error!(
374 citadel_io::ErrorCode::BroadcastCreateGroupFailed
375 ))
376 }
377
378 _ => {}
379 }
380 }
381
382 Ok(())
383 }
384
385 fn construct(kernel: Box<dyn NetKernel<R> + 'a>) -> Self {
386 let (tx, rx) = citadel_io::tokio::sync::mpsc::unbounded_channel();
387 Self {
388 shared: Arc::new(BroadcastShared {
389 route_registers: AtomicBool::new(false),
390 register_rx: citadel_io::Mutex::new(Some(rx)),
391 register_tx: tx,
392 }),
393 inner_kernel: kernel,
394 _pd: Default::default(),
395 }
396 }
397}
398
399#[async_trait]
400impl<F, Fut, R: Ratchet> NetKernel<R> for BroadcastKernel<'_, F, Fut, R> {
401 fn load_remote(&mut self, node_remote: NodeRemote<R>) -> Result<(), NetworkError> {
402 self.inner_kernel.load_remote(node_remote)
403 }
404
405 async fn on_start(&self) -> Result<(), NetworkError> {
406 self.inner_kernel.on_start().await
407 }
408
409 async fn on_node_event_received(&self, message: NodeResult<R>) -> Result<(), NetworkError> {
410 if let NodeResult::PeerEvent(PeerEvent {
411 event: ps @ PeerSignal::PostRegister { .. },
412 ..
413 }) = &message
414 {
415 if self.shared.route_registers.load(Ordering::Relaxed) {
416 return self
417 .shared
418 .register_tx
419 .send(ps.clone())
420 .map_err(|err| NetworkError::generic(err.to_string()));
421 }
422 }
423
424 self.inner_kernel.on_node_event_received(message).await
425 }
426
427 async fn on_stop(&mut self) -> Result<(), NetworkError> {
428 self.inner_kernel.on_stop().await
429 }
430}
431
432#[cfg(all(test, feature = "localhost-testing"))]
433mod tests {
434 use crate::prefabs::client::broadcast::{BroadcastKernel, GroupInitRequestType};
435 use crate::prefabs::client::peer_connection::PeerConnectionKernel;
436 use crate::prefabs::client::DefaultServerConnectionSettingsBuilder;
437 use crate::prelude::*;
438 use crate::test_common::{server_info, wait_for_peers, TestBarrier};
439 use citadel_io::tokio;
440 use futures::prelude::stream::FuturesUnordered;
441 use futures::TryStreamExt;
442 use rstest::rstest;
443 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
444 use uuid::Uuid;
445
446 #[citadel_io::tokio::test(flavor = "multi_thread")]
447 async fn group_connect_list_members() -> Result<(), Box<dyn std::error::Error>> {
448 let peer_count = 3;
449 assert!(peer_count > 1);
450 citadel_logging::setup_log();
451 TestBarrier::setup(peer_count);
452
453 let client_success = &AtomicUsize::new(0);
454 let (server, server_addr) = server_info::<StackedRatchet>();
455
456 let client_kernels = FuturesUnordered::new();
457 let total_peers = (0..peer_count)
458 .map(|_| Uuid::new_v4())
459 .collect::<Vec<Uuid>>();
460 let group_id = Uuid::new_v4();
461
462 for idx in 0..peer_count {
463 let uuid = total_peers.get(idx).cloned().unwrap();
464
465 let request = if idx == 0 {
466 GroupInitRequestType::Create {
468 local_user: UserIdentifier::from(uuid),
469 invite_list: vec![],
470 group_id,
471 accept_registrations: true,
472 }
473 } else {
474 GroupInitRequestType::Join {
475 local_user: UserIdentifier::from(uuid),
476 owner: total_peers.first().cloned().unwrap().into(),
477 group_id,
478 do_peer_register: true,
479 }
480 };
481
482 let server_connection_settings =
483 DefaultServerConnectionSettingsBuilder::transient_with_id(server_addr, uuid)
484 .build()
485 .unwrap();
486
487 let client_kernel = BroadcastKernel::new(
488 server_connection_settings,
489 request,
490 move |channel, connection| async move {
491 wait_for_peers().await;
492 log::trace!(target: "citadel", "***GROUP PEER {}={}={} CONNECT SUCCESS***", idx, uuid, connection.conn_type.get_session_cid());
493
494 let owned_groups = connection.list_owned_groups().await.unwrap();
495
496 if idx == 0 {
497 assert_eq!(owned_groups.len(), 1);
498 } else {
499 assert_eq!(owned_groups.len(), 0);
500 }
501
502 log::trace!(target: "citadel", "Peer {idx}={} is COMPLETE!", connection.conn_type.get_session_cid());
503
504 let _ = client_success.fetch_add(1, Ordering::Relaxed);
505 wait_for_peers().await;
506 drop(channel);
507 connection.shutdown_kernel().await
508 },
509 );
510
511 let client = DefaultNodeBuilder::default()
512 .with_stun_servers(crate::test_common::LOCALHOST_STUN_SERVERS)
513 .build(client_kernel)
514 .unwrap();
515
516 client_kernels.push(async move { client.await.map(|_| ()) });
517 }
518
519 let clients = Box::pin(async move { client_kernels.try_collect::<()>().await.map(|_| ()) });
520
521 let res = futures::future::try_select(server, clients).await;
522 if let Err(err) = res {
523 return match err {
524 futures::future::Either::Left(left) => Err(left.0.into_string().into()),
525 futures::future::Either::Right(right) => Err(right.0.into_string().into()),
526 };
527 }
528
529 assert_eq!(client_success.load(Ordering::Relaxed), peer_count);
530 Ok(())
531 }
532
533 #[rstest]
534 #[case(2)]
535 #[timeout(std::time::Duration::from_secs(90))]
536 #[citadel_io::tokio::test(flavor = "multi_thread")]
537 async fn test_manual_group_connect(
538 #[case] peer_count: usize,
539 ) -> Result<(), Box<dyn std::error::Error>> {
540 assert!(peer_count > 1);
545 citadel_logging::setup_log();
546 TestBarrier::setup(peer_count);
547
548 let client_success = &AtomicBool::new(false);
549 let receiver_success = &AtomicBool::new(false);
550
551 let (server, server_addr) = server_info::<StackedRatchet>();
552
553 let client_kernels = FuturesUnordered::new();
554 let total_peers = (0..peer_count)
555 .map(|_| Uuid::new_v4())
556 .collect::<Vec<Uuid>>();
557
558 for idx in 0..peer_count {
559 let uuid = total_peers.get(idx).cloned().unwrap();
560 let peers = total_peers
561 .clone()
562 .into_iter()
563 .filter(|r| r != &uuid)
564 .map(UserIdentifier::from)
565 .collect::<Vec<UserIdentifier>>();
566
567 let server_connection_settings =
568 DefaultServerConnectionSettingsBuilder::transient_with_id(server_addr, uuid)
569 .build()
570 .unwrap();
571
572 let client_kernel = PeerConnectionKernel::new(
573 server_connection_settings,
574 peers,
575 move |mut results, remote| async move {
576 let _sender = remote.conn_type.get_session_cid();
577 let mut signals = remote.get_unprocessed_signals_receiver().unwrap();
578
579 wait_for_peers().await;
580 let conn = results.recv().await.unwrap()?;
581 log::trace!(target: "citadel", "User {uuid} received {conn:?}");
582
583 if idx == 0 {
585 let _channel = remote
586 .create_group(Some(vec![conn.channel.get_peer_cid().into()]))
587 .await?;
588 log::info!(target: "citadel", "The designated node has finished creating a group");
589
590 wait_for_peers().await;
591 client_success.store(true, Ordering::Relaxed);
592 return remote.shutdown_kernel().await;
593 } else {
594 while let Some(evt) = signals.recv().await {
596 log::info!(target: "citadel", "Received unprocessed signal: {evt:?}");
597 match evt {
598 NodeResult::GroupEvent(GroupEvent {
599 session_cid: _,
600 ticket: _,
601 event:
602 GroupBroadcast::Invitation {
603 sender: _,
604 key: _key,
605 },
606 }) => {
607 let _ =
608 crate::responses::group_invite(evt, true, &remote.inner)
609 .await?;
610 }
611
612 NodeResult::GroupChannelCreated(GroupChannelCreated {
613 ticket: _,
614 channel: _chan,
615 session_cid: _,
616 }) => {
617 receiver_success.store(true, Ordering::Relaxed);
618 log::trace!(target: "citadel", "***PEER {uuid} CONNECT***");
619 wait_for_peers().await;
620 return remote.shutdown_kernel().await;
621 }
622
623 val => {
624 log::warn!(target: "citadel", "Unhandled response: {val:?}")
625 }
626 }
627 }
628 }
629
630 Err(citadel_io::error!(
631 citadel_io::ErrorCode::BroadcastStreamEndedUnexpectedly,
632 "signals_recv"
633 ))
634 },
635 );
636
637 let client = DefaultNodeBuilder::default()
638 .with_stun_servers(crate::test_common::LOCALHOST_STUN_SERVERS)
639 .build(client_kernel)
640 .unwrap();
641 client_kernels.push(async move { client.await.map(|_| ()) });
642 }
643
644 let clients = Box::pin(async move { client_kernels.try_collect::<()>().await.map(|_| ()) });
645
646 if let Err(err) = futures::future::try_select(server, clients).await {
647 return match err {
648 futures::future::Either::Left(res) => Err(res.0.into_string().into()),
649 futures::future::Either::Right(res) => Err(res.0.into_string().into()),
650 };
651 }
652
653 assert!(client_success.load(Ordering::Relaxed));
654 assert!(receiver_success.load(Ordering::Relaxed));
655 Ok(())
656 }
657
658 #[citadel_io::tokio::test(flavor = "multi_thread")]
663 async fn group_command_hierarchy_superior_reads_subordinate(
664 ) -> Result<(), Box<dyn std::error::Error>> {
665 use crate::prelude::GroupBroadcastPayload;
666 use citadel_types::crypto::SecBuffer;
667 use citadel_types::proto::{
668 CommandPath, GroupHierarchyMode, MessageGroupOptions, ReadPolicy,
669 };
670 use std::collections::HashMap;
671
672 let peer_count = 2;
673 citadel_logging::setup_log();
674 TestBarrier::setup(peer_count);
675
676 let owner_read = &AtomicBool::new(false);
677 let (server, server_addr) = server_info::<StackedRatchet>();
678 let client_kernels = FuturesUnordered::new();
679 let total_peers = (0..peer_count)
680 .map(|_| Uuid::new_v4())
681 .collect::<Vec<Uuid>>();
682
683 for idx in 0..peer_count {
684 let uuid = total_peers.get(idx).cloned().unwrap();
685 let peers = total_peers
686 .clone()
687 .into_iter()
688 .filter(|r| r != &uuid)
689 .map(UserIdentifier::from)
690 .collect::<Vec<UserIdentifier>>();
691 let server_connection_settings =
692 DefaultServerConnectionSettingsBuilder::transient_with_id(server_addr, uuid)
693 .build()
694 .unwrap();
695
696 let client_kernel = PeerConnectionKernel::new(
697 server_connection_settings,
698 peers,
699 move |mut results, remote| async move {
700 let mut signals = remote.get_unprocessed_signals_receiver().unwrap();
701 wait_for_peers().await;
702 let conn = results.recv().await.unwrap()?;
703
704 if idx == 0 {
705 let sub_cid: u64 = conn.channel.get_peer_cid();
707 let mut ranks = HashMap::new();
708 let _ = ranks.insert(sub_cid, CommandPath::parse("/alpha"));
709 let options = MessageGroupOptions {
710 hierarchy: GroupHierarchyMode::CommandHierarchy {
711 read_policy: ReadPolicy::SuperiorOnly,
712 ranks,
713 },
714 ..Default::default()
715 };
716 let mut channel = remote
717 .create_group_with_options(Some(vec![sub_cid.into()]), options)
718 .await?;
719
720 loop {
722 match channel.recv().await {
723 Some(GroupBroadcastPayload::Message { payload, sender: _ }) => {
724 assert_eq!(payload.as_ref(), b"sitrep from subordinate");
725 owner_read.store(true, Ordering::Relaxed);
726 break;
727 }
728 Some(_) => continue,
729 None => break,
730 }
731 }
732 wait_for_peers().await;
733 return remote.shutdown_kernel().await;
734 }
735
736 while let Some(evt) = signals.recv().await {
738 match evt {
739 NodeResult::GroupEvent(GroupEvent {
740 event: GroupBroadcast::Invitation { .. },
741 ..
742 }) => {
743 let _ = crate::responses::group_invite(evt, true, &remote.inner)
744 .await?;
745 }
746 NodeResult::GroupChannelCreated(GroupChannelCreated {
747 channel,
748 ..
749 }) => {
750 channel
751 .send_message(SecBuffer::from(
752 b"sitrep from subordinate".to_vec(),
753 ))
754 .await?;
755 wait_for_peers().await;
756 return remote.shutdown_kernel().await;
757 }
758 _ => {}
759 }
760 }
761
762 Err(citadel_io::error!(
763 citadel_io::ErrorCode::BroadcastStreamEndedUnexpectedly,
764 "signals"
765 ))
766 },
767 );
768 let client = DefaultNodeBuilder::default()
769 .with_stun_servers(crate::test_common::LOCALHOST_STUN_SERVERS)
770 .build(client_kernel)
771 .unwrap();
772 client_kernels.push(async move { client.await.map(|_| ()) });
773 }
774
775 let clients = Box::pin(async move { client_kernels.try_collect::<()>().await.map(|_| ()) });
776 if let Err(err) = futures::future::try_select(server, clients).await {
777 return match err {
778 futures::future::Either::Left(res) => Err(res.0.into_string().into()),
779 futures::future::Either::Right(res) => Err(res.0.into_string().into()),
780 };
781 }
782
783 assert!(
784 owner_read.load(Ordering::Relaxed),
785 "owner (hierarchy root) must read the subordinate's DHE message"
786 );
787 Ok(())
788 }
789}