aranya-daemon 6.0.0

//! Aranya graph actions/effects API.

use std::{future::Future, marker::PhantomData};

use anyhow::{Context, Result};
use aranya_crypto::{
    policy::{LabelId, RoleId},
    DeviceId, Random as _, Rng,
};
use aranya_daemon_api::Rank;
use aranya_keygen::PublicKeys;
use aranya_policy_ifgen::{Actionable, VmEffect};
use aranya_policy_text::Text;
#[cfg(feature = "afc")]
use aranya_runtime::NullSink;
use aranya_runtime::{GraphId, PolicyStore, Session, StorageProvider, VmPolicy};
use futures_util::TryFutureExt as _;
use tracing::{debug, instrument, warn, Instrument};

use crate::{
    aranya::Client,
    policy::{self, ChanOp, Effect, Perm, PublicKeyBundle},
    vm_policy::{MsgSink, VecSink},
};

/// Container for complex AQC channel creation results.
#[derive(Debug)]
pub struct SessionData {
    /// The serialized messages
    #[cfg(feature = "afc")]
    pub ctrl: Vec<Box<[u8]>>,
    /// The effects produced
    pub effects: Vec<Effect>,
}

/// Functions related to Aranya actions
impl<PS, SP, CE> Client<PS, SP>
where
    PS: PolicyStore<Policy = VmPolicy<CE>, Effect = VmEffect> + Send + 'static,
    SP: StorageProvider + Send + 'static,
    CE: aranya_crypto::Engine + Send + Sync + 'static,
{
    /// Creates the team.
    /// Creates a new graph, adds the `CreateTeam` command to the root of the graph.
    /// Returns the [`GraphId`] of the newly created graph.
    #[instrument(skip_all)]
    pub async fn create_team(
        &self,
        owner_keys: PublicKeyBundle,
        nonce: Option<&[u8]>,
    ) -> Result<(GraphId, Vec<Effect>)> {
        let mut sink = VecSink::new();
        let policy_data = &[0u8];
        let act = policy::create_team(
            owner_keys,
            nonce.unwrap_or(&<[u8; 16]>::random(Rng)).to_vec(),
        );
        let id = {
            let mut client = self.lock_aranya().await;
            act.with_action(|act| client.new_graph(policy_data, act, &mut sink))
                .context("unable to create new team")?
        };
        Ok((id, sink.collect()?))
    }

    /// Returns an implementation of [`Actions`] for a particular
    /// storage.
    #[instrument(skip_all, fields(%graph_id))]
    pub fn actions(&self, graph_id: GraphId) -> impl Actions<PS, SP, CE> {
        ActionsImpl {
            client: self.clone(),
            graph_id,
            _eng: PhantomData,
        }
    }

    /// Create new ephemeral Session.
    /// Once the Session has been created, call `session_receive` to add an ephemeral command to the Session.
    #[instrument(skip_all, fields(%graph_id))]
    pub(crate) async fn session_new(&self, graph_id: GraphId) -> Result<Session<SP, PS>> {
        let session = self.lock_aranya().await.session(graph_id)?;
        Ok(session)
    }

    /// Receives an ephemeral command from another ephemeral Session.
    /// Assumes an ephemeral Session has already been created before adding an ephemeral command to the Session.
    #[instrument(skip_all)]
    pub(crate) async fn session_receive(
        &self,
        session: &mut Session<SP, PS>,
        command: &[u8],
    ) -> Result<Vec<Effect>> {
        let client = self.lock_aranya().await;
        let mut sink = VecSink::new();
        session.receive(&client, &mut sink, command)?;
        Ok(sink.collect()?)
    }
}

/// Implements [`Actions`] for a particular storage.
struct ActionsImpl<PS, SP, CE> {
    /// Aranya client graph state.
    client: Client<PS, SP>,
    /// Aranya graph ID.
    graph_id: GraphId,
    /// Crypto engine.
    _eng: PhantomData<CE>,
}

impl<PS, SP, CE> Actions<PS, SP, CE> for ActionsImpl<PS, SP, CE>
where
    PS: PolicyStore<Policy = VmPolicy<CE>, Effect = VmEffect> + Send + 'static,
    SP: StorageProvider + Send + 'static,
    CE: aranya_crypto::Engine + Send + Sync + 'static,
{
    async fn call_persistent_action(
        &self,
        act: impl Actionable<Interface = policy::Persistent> + Send,
    ) -> Result<Vec<Effect>> {
        let mut sink = VecSink::new();
        // Make sure we drop the lock as quickly as possible.
        {
            let mut client = self.client.lock_aranya().await;
            act.with_action(|act| client.action(self.graph_id, &mut sink, act))?;
        }

        let total = sink.effects.len();
        for (i, effect) in sink.effects.iter().enumerate() {
            debug!(i, total, effect = effect.name.as_str());
        }

        Ok(sink.collect()?)
    }

    async fn call_session_action(
        &self,
        act: impl Actionable<Interface = policy::Ephemeral> + Send,
    ) -> Result<SessionData> {
        let mut sink = VecSink::new();
        let mut msg_sink = MsgSink::new();
        {
            let mut client = self.client.lock_aranya().await;
            let mut session = client.session(self.graph_id)?;
            act.with_action(|act| session.action(&client, &mut sink, &mut msg_sink, act))?;
        }
        Ok(SessionData {
            #[cfg(feature = "afc")]
            ctrl: msg_sink.into_cmds(),
            effects: sink.collect()?,
        })
    }
}

/// A programmatic API for policy actions.
pub trait Actions<PS, SP, CE>
where
    PS: PolicyStore<Policy = VmPolicy<CE>, Effect = VmEffect> + Send + 'static,
    SP: StorageProvider + Send + 'static,
    CE: aranya_crypto::Engine + Send + Sync + 'static,
{
    /// Perform a persistent action.
    fn call_persistent_action(
        &self,
        act: impl Actionable<Interface = policy::Persistent> + Send,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send;

    #[allow(clippy::type_complexity)]
    /// Performs a session action.
    fn call_session_action(
        &self,
        act: impl Actionable<Interface = policy::Ephemeral> + Send,
    ) -> impl Future<Output = Result<SessionData>> + Send;

    /// Invokes `add_device`.
    #[instrument(skip_all)]
    fn add_device(
        &self,
        keys: PublicKeyBundle,
        initial_role_id: Option<RoleId>,
        rank: Rank,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::add_device(
            keys,
            initial_role_id.map(|id| id.as_base()),
            rank.value(),
        ))
        .in_current_span()
    }

    /// Invokes `create_role`.
    #[instrument(skip(self))]
    fn create_role(
        &self,
        role_name: Text,
        rank: Rank,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::create_role(role_name, rank.value()))
            .in_current_span()
    }

    /// Invokes `delete_role`.
    #[instrument(skip(self))]
    fn delete_role(&self, role_id: RoleId) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::delete_role(role_id.as_base()))
            .in_current_span()
    }

    /// Invokes `add_perm_to_role`.
    #[instrument(skip(self), fields(%role_id))]
    fn add_perm_to_role(
        &self,
        role_id: RoleId,
        perm: Perm,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::add_perm_to_role(role_id.as_base(), perm))
            .in_current_span()
    }

    /// Invokes `remove_perm_from_role`.
    #[instrument(skip(self), fields(%role_id))]
    fn remove_perm_from_role(
        &self,
        role_id: RoleId,
        perm: Perm,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::remove_perm_from_role(role_id.as_base(), perm))
            .in_current_span()
    }

    /// Invokes `query_role_perms`.
    #[instrument(skip(self), fields(%role_id))]
    fn query_role_perms(
        &self,
        role_id: RoleId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_session_action(policy::query_role_perms(role_id.as_base()))
            .map_ok(|SessionData { effects, .. }| effects)
            .in_current_span()
    }

    /// Invokes `change_rank`.
    #[instrument(skip(self))]
    fn change_rank(
        &self,
        object_id: aranya_daemon_api::ObjectId,
        old_rank: Rank,
        new_rank: Rank,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::change_rank(
            object_id.as_base(),
            old_rank.value(),
            new_rank.value(),
        ))
        .in_current_span()
    }

    /// Invokes `query_device_generation`.
    #[cfg(feature = "test-utils")]
    #[instrument(skip(self))]
    fn query_device_generation(
        &self,
        device_id: DeviceId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_session_action(policy::query_device_generation(device_id.as_base()))
            .map_ok(|SessionData { effects, .. }| effects)
            .in_current_span()
    }

    /// Invokes `query_rank`.
    #[allow(clippy::type_complexity)]
    #[instrument(skip(self))]
    fn query_rank(
        &self,
        object_id: aranya_daemon_api::ObjectId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_session_action(policy::query_rank(object_id.as_base()))
            .map_ok(|SessionData { effects, .. }| effects)
            .in_current_span()
    }

    /// Invokes `assign_label_to_device`.
    #[instrument(skip(self), fields(%device_id, %label_id, %op))]
    fn assign_label_to_device(
        &self,
        device_id: DeviceId,
        label_id: LabelId,
        op: ChanOp,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::assign_label_to_device(
            device_id.as_base(),
            label_id.as_base(),
            op,
        ))
        .in_current_span()
    }

    /// Invokes `assign_role`.
    #[instrument(skip(self), fields(%device_id, %role_id))]
    fn assign_role(
        &self,
        device_id: DeviceId,
        role_id: RoleId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::assign_role(device_id.as_base(), role_id.as_base()))
            .in_current_span()
    }

    /// Invokes `change_role`.
    #[instrument(skip(self), fields(%device_id, %old_role_id, %new_role_id))]
    fn change_role(
        &self,
        device_id: DeviceId,
        old_role_id: RoleId,
        new_role_id: RoleId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::change_role(
            device_id.as_base(),
            old_role_id.as_base(),
            new_role_id.as_base(),
        ))
        .in_current_span()
    }

    /// Invokes `create_label`.
    #[instrument(skip(self), fields(%name))]
    fn create_label(
        &self,
        name: Text,
        rank: Rank,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::create_label(name, rank.value()))
            .in_current_span()
    }

    /// Creates a unidirectional AFC channel.
    #[cfg(feature = "afc")]
    #[allow(clippy::type_complexity)]
    #[instrument(skip(self), fields(open_id = %open_id, label_id = %label_id))]
    fn create_afc_uni_channel_off_graph(
        &self,
        open_id: DeviceId,
        label_id: LabelId,
    ) -> impl Future<Output = Result<SessionData>> + Send {
        self.call_session_action(policy::create_afc_uni_channel(
            open_id.as_base(),
            label_id.as_base(),
        ))
        .in_current_span()
    }

    /// Invokes `delete_label`.
    #[instrument(skip(self), fields(%label_id))]
    fn delete_label(&self, label_id: LabelId) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::delete_label(label_id.as_base()))
            .in_current_span()
    }

    /// Invokes `query_device_public_key_bundle`.
    #[allow(clippy::type_complexity)]
    #[instrument(skip(self), fields(%device_id))]
    fn query_device_public_key_bundle(
        &self,
        device_id: DeviceId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_session_action(policy::query_device_public_key_bundle(device_id.as_base()))
            .map_ok(|SessionData { effects, .. }| effects)
            .in_current_span()
    }

    /// Invokes `query_device_role`.
    #[allow(clippy::type_complexity)]
    #[instrument(skip(self), fields(%device_id))]
    fn query_device_role(
        &self,
        device_id: DeviceId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_session_action(policy::query_device_role(device_id.as_base()))
            .map_ok(|SessionData { effects, .. }| effects)
            .in_current_span()
    }

    /// Invokes `query_devices_on_team`.
    #[allow(clippy::type_complexity)]
    #[instrument(skip(self))]
    fn query_devices_on_team(&self) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_session_action(policy::query_devices_on_team())
            .map_ok(|SessionData { effects, .. }| effects)
            .in_current_span()
    }

    /// Invokes `query_label`.
    #[allow(clippy::type_complexity)]
    #[instrument(skip(self), fields(%label_id))]
    fn query_label(&self, label_id: LabelId) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_session_action(policy::query_label(label_id.as_base()))
            .map_ok(|SessionData { effects, .. }| effects)
            .in_current_span()
    }

    /// Invokes `query_labels`.
    #[allow(clippy::type_complexity)]
    #[instrument(skip(self))]
    fn query_labels(&self) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_session_action(policy::query_labels())
            .map_ok(|SessionData { effects, .. }| effects)
            .in_current_span()
    }

    /// Invokes `query_labels_assigned_to_device`.
    #[instrument(skip(self), fields(%device))]
    fn query_labels_assigned_to_device(
        &self,
        device: DeviceId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_session_action(policy::query_labels_assigned_to_device(device.as_base()))
            .map_ok(|SessionData { effects, .. }| effects)
            .in_current_span()
    }

    /// Invokes `query_team_roles`.
    #[instrument(skip(self))]
    fn query_team_roles(&self) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_session_action(policy::query_team_roles())
            .map_ok(|SessionData { effects, .. }| effects)
            .in_current_span()
    }

    /// Invokes `remove_device`.
    #[instrument(skip(self), fields(%device_id))]
    fn remove_device(
        &self,
        device_id: DeviceId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::remove_device(device_id.as_base()))
            .in_current_span()
    }

    /// Invokes `revoke_label_from_device`.
    #[instrument(skip(self), fields(%device_id, %label_id))]
    fn revoke_label_from_device(
        &self,
        device_id: DeviceId,
        label_id: LabelId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::revoke_label_from_device(
            device_id.as_base(),
            label_id.as_base(),
        ))
        .in_current_span()
    }

    /// Invokes `revoke_role`.
    #[instrument(skip(self), fields(%device_id, %role_id))]
    fn revoke_role(
        &self,
        device_id: DeviceId,
        role_id: RoleId,
    ) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::revoke_role(device_id.as_base(), role_id.as_base()))
            .in_current_span()
    }

    /// Invokes `setup_default_roles`.
    #[instrument(skip(self))]
    fn setup_default_roles(&self) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::setup_default_roles())
            .in_current_span()
    }

    /// Invokes `terminate_team`.
    #[instrument(skip(self), fields(%team_id))]
    fn terminate_team(&self, team_id: GraphId) -> impl Future<Output = Result<Vec<Effect>>> + Send {
        self.call_persistent_action(policy::terminate_team(team_id.as_base()))
            .in_current_span()
    }
}

impl<CS: aranya_crypto::CipherSuite> TryFrom<&PublicKeys<CS>> for PublicKeyBundle {
    type Error = postcard::Error;
    fn try_from(pk: &PublicKeys<CS>) -> Result<Self, Self::Error> {
        Ok(Self {
            ident_key: postcard::to_allocvec(&pk.ident_pk)?,
            enc_key: postcard::to_allocvec(&pk.enc_pk)?,
            sign_key: postcard::to_allocvec(&pk.sign_pk)?,
        })
    }
}

#[cfg(feature = "afc")]
pub(crate) fn query_afc_channel_is_valid<PS, SP, CE>(
    aranya: &mut aranya_runtime::ClientState<PS, SP>,
    graph_id: GraphId,
    sender_id: DeviceId,
    receiver_id: DeviceId,
    label_id: LabelId,
) -> Result<bool>
where
    PS: PolicyStore<Policy = VmPolicy<CE>, Effect = VmEffect>,
    SP: StorageProvider,
    CE: aranya_crypto::Engine,
{
    let mut session = aranya.session(graph_id)?;
    let mut sink = VecSink::new();
    policy::query_afc_channel_is_valid(
        sender_id.as_base(),
        receiver_id.as_base(),
        label_id.as_base(),
    )
    .with_action(|act| session.action(aranya, &mut sink, &mut NullSink, act))?;
    let effects = sink.collect()?;
    Ok(effects.iter().any(|e| {
        if let Effect::QueryAfcChannelIsValidResult(e) = e {
            return e.is_valid;
        }
        false
    }))
}