product_os_command_control/

registry.rs

//! Node registry module
//!
//! Manages the registry of Product OS nodes in the cluster, including
//! node information, capabilities, services, and features.

use core::str::FromStr;
use std::prelude::v1::*;

use std::collections::BTreeMap;
use std::sync::Arc;
use serde::{ Deserialize, Serialize };

use product_os_capabilities::{Features, ServiceError, Services, What};
use product_os_security::{AsByteVector, DHKeyStore, RandomGenerator, certificates::Certificates, RandomGeneratorTemplate, RNG, StdRng, SeedableRng};
use product_os_store::ProductOSKeyValueStore;

use chrono::{DateTime, Utc };
use parking_lot::Mutex;
use product_os_request::Uri;

/// Represents a single server instance (node) in the Product OS cluster.
///
/// Each node has a unique identifier, machine ID, URI, and tracks its own
/// capabilities, services, features, and failure count.
#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct Node {
    id: uuid::Uuid,
    machine_id: String,

    uri: String,
    process_id: u32,

    certificate: Vec<u8>,

    capabilities: Vec<String>,
    services: Services,
    features: Features,

    failures: u8,
    created_at: DateTime<Utc>,
    updated_at: DateTime<Utc>
}

impl AsByteVector for &Node {
    fn as_byte_vector(&self) -> Vec<u8> {
        let mut bytes = vec!();

        bytes.extend_from_slice(self.id.as_bytes());
        bytes.extend_from_slice(self.machine_id.as_bytes());
        bytes.extend_from_slice(self.uri.as_bytes());
        bytes.extend_from_slice(&self.certificate);
        bytes.extend_from_slice(&[self.failures]);
        bytes.extend_from_slice(self.created_at.to_string().as_bytes());
        bytes.extend_from_slice(self.updated_at.to_string().as_bytes());

        bytes
    }
}


impl Node {
    /// Creates a new `Node` instance with the given configuration and certificates.
    ///
    /// Generates a unique UUID, retrieves the machine ID, and initialises
    /// the node with the configured URL address.
    ///
    /// # Panics
    ///
    /// Panics if the machine UID cannot be retrieved or if the configured
    /// URL address is not a valid URI.
    pub fn new(config: &crate::config::CommandControl, certificates: Certificates) -> Self {
        let machine_uid = match machine_uid::get() {
            Ok(uid) => uid,
            Err(e) => panic!("Unable to generate machine id: {}", e)
        };

        Self {
            id: uuid::Uuid::new_v4(),
            uri: Uri::from_str(config.url_address.as_str()).unwrap().to_string(),
            process_id: std::process::id(),
            machine_id: product_os_security::create_string_hash(machine_uid.as_str()),
            certificate: certificates.certificates.first().unwrap().to_owned(),
            capabilities: Vec::new(),
            services: Services::new(),
            features: Features::new(),
            failures: 0,
            created_at: Utc::now(),
            updated_at: Utc::now()
        }
    }

    /// Creates a new `Node` instance with the given configuration and certificates.
    #[deprecated(since = "0.0.28", note = "Renamed to Node::new() to follow Rust conventions")]
    pub fn default(config: &crate::config::CommandControl, certificates: Certificates) -> Self {
        Self::new(config, certificates)
    }

    /// Returns the unique identifier for this node as a string.
    pub fn get_identifier(&self) -> String {
        self.id.to_string()
    }

    /// Returns the URI protocol scheme (e.g., "https") for this node.
    ///
    /// Returns an empty string if the URI has no scheme.
    #[deprecated(since = "0.0.28", note = "Use try_get_protocol() to avoid potential panics")]
    pub fn get_protocol(&self) -> String {
        let uri = Uri::from_str(self.uri.as_str()).unwrap();
        match uri.scheme() {
            None => String::new(),
            Some(s) => s.to_string()
        }
    }

    /// Returns the URI protocol scheme (e.g., "https") for this node, or `None` if the URI is invalid.
    pub fn try_get_protocol(&self) -> Option<String> {
        Uri::from_str(self.uri.as_str())
            .ok()
            .and_then(|uri| uri.scheme().map(|s| s.to_string()))
    }

    /// Returns the parsed URI address for this node.
    ///
    /// # Panics
    ///
    /// Panics if the stored URI string is not valid. This should not happen
    /// if the node was constructed via `Node::new()`.
    #[deprecated(since = "0.0.28", note = "Use try_get_address() to avoid potential panics")]
    pub fn get_address(&self) -> Uri {
        Uri::from_str(self.uri.as_str()).unwrap()
    }

    /// Returns the parsed URI address for this node, or `None` if the URI is invalid.
    pub fn try_get_address(&self) -> Option<Uri> {
        Uri::from_str(self.uri.as_str()).ok()
    }

    /// Returns the OS process ID of this node.
    pub fn get_process_id(&self) -> u32 {
        self.process_id
    }

    /// Returns a copy of the certificate bytes for this node.
    pub fn get_certificate(&self) -> Vec<u8> {
        self.certificate.to_owned()
    }

    /// Returns the current failure count for this node.
    pub fn get_failures(&self) -> u8 {
        self.failures
    }

    /// Returns a reference to this node's registered features.
    pub fn get_features(&self) -> &Features {
        &self.features
    }

    /// Returns a reference to this node's registered services.
    pub fn get_services(&self) -> &Services {
        &self.services
    }

    /// Tests whether this node matches a single selector/value pair.
    ///
    /// Supported selectors:
    /// - `"feature"` - checks if the node has the named feature
    /// - `"capability"` - checks if the node has the named capability (empty capabilities match nothing)
    /// - `"service.kind"` - checks if the node has a service of the given kind
    /// - `"service.enabled"` - checks if all services match the enabled state
    /// - `"service.active"` - checks if all services match the active state
    pub fn match_node(&self, selector: &str, search_value: &str) -> bool {
        match selector {
            "feature" => {
                self.features.get(search_value).is_some()
            },
            "capability" => {
                self.capabilities.contains(&search_value.to_string())
            },
            "service.kind" => {
                self.services.find(search_value).is_some()
            },
            "service.enabled" => {
                self.services.list().all(|(_, service)| service.enabled.to_string() == search_value)
            },
            "service.active" => {
                self.services.list().all(|(_, service)| service.active.to_string() == search_value)
            },
            _ => true
        }
    }

    /// Tests whether this node matches all selector/value pairs in the query.
    ///
    /// Returns `true` only if every entry in the query matches.
    /// An empty query always matches.
    pub fn match_node_query(&self, query: &BTreeMap<&str, &str>) -> bool {
        query.iter().all(|(selector, value)| self.match_node(selector, value))
    }

    /// Returns the timestamp when this node was created.
    pub fn get_created_at(&self) -> DateTime<Utc> {
        self.created_at
    }

    /// Returns the timestamp when this node was last updated.
    pub fn get_last_updated_at(&self) -> DateTime<Utc> {
        self.updated_at
    }
}



/// Manages the collection of known nodes in the Product OS cluster.
///
/// The registry tracks the local node ("me"), remote nodes discovered
/// from the key-value store, and cryptographic key sessions for
/// authenticated communication.
pub struct Registry {
    me: Node,
    nodes: BTreeMap<String, Node>,
    key_store: DHKeyStore,

    store: Arc<ProductOSKeyValueStore>,

    max_failures: u8
}

impl Registry {
    /// Creates a new registry with the given configuration and certificates.
    ///
    /// Initialises the local node and an empty set of remote nodes.
    pub fn new(config: &crate::config::CommandControl, key_value_store: Arc<ProductOSKeyValueStore>, certificates: Certificates) -> Self {
        let me = Node::new(config, certificates);

        Registry {
            me,
            nodes: BTreeMap::new(),
            key_store: DHKeyStore::new(),
            store: key_value_store,
            max_failures: config.max_failures,
        }
    }

    /// Returns the maximum number of failures before a node is removed.
    pub fn get_max_failures(&self) -> u8 {
        self.max_failures
    }

    async fn upsert_me_remote(&mut self) {
        self.store.group_set(self.me.id.to_string().as_str(), serde_json::to_string(&self.me).unwrap().as_str()).unwrap_or_default()
    }

    async fn upsert_node_remote(&mut self, node: &Node) {
        tracing::info!("Upserting node: {:?}", node);
        self.store.group_set(node.id.to_string().as_str(), serde_json::to_string(node).unwrap().as_str()).unwrap_or_default();
    }

    async fn remove_node_remote(&mut self, identifier: &str) {
        self.store.group_remove(identifier).unwrap_or_default()
    }

    async fn get_node_remote(&mut self, id: &str) -> Option<Node> {
        match self.store.group_get(id) {
            Ok(v) => {
                let mut node: Node = serde_json::from_str(v.as_str()).unwrap();
                let _ = node.features.setup_router();
                Some(node)
            },
            Err(_) => None
        }
    }

    /// Checks if the local node still exists in the remote store.
    ///
    /// Returns `Some` reference to self if the remote node matches,
    /// `None` if the node was lost or has a mismatched ID.
    pub async fn check_me_remote(&mut self) -> Option<&Node> {
        let id = self.me.id.to_string();

        match self.get_node_remote(id.as_str()).await {
            Some(node) => {
                if node.id != self.me.id {
                    None
                }
                else {
                    Some(&self.me)
                }
            },
            None => { None }
        }
    }

    /// Returns a reference to the local node.
    pub fn get_me(&self) -> &Node {
        &self.me
    }

    /// Persists the local node's current state to the remote store.
    pub async fn update_me(&mut self) {
        self.upsert_me_remote().await;
    }

    /// Updates the local node's failure status and returns whether the node is still alive.
    ///
    /// If `success` is `true`, resets the failure counter. If `false`, increments it.
    /// Returns `false` (dead) when the failure count reaches `max_failures`.
    pub fn update_me_status(&mut self, success: bool) -> bool {
        let failures = if success { 0 } else { self.me.failures + 1 };

        if failures < self.max_failures {
            self.me.failures = failures;
            self.me.updated_at = Utc::now();

            true
        }
        else {
            false
        }
    }

    /// Updates a remote node's pulse status and returns whether the node is still alive.
    ///
    /// Fetches the node from the remote store, updates its failure count,
    /// and removes it if the failure threshold is exceeded.
    pub async fn update_pulse_status(&mut self, id: &str, success: bool) -> bool {
        match self.get_node_remote(id).await {
            Some(mut node) => {
                let failures = if success { 0 } else { node.failures + 1 };

                if failures < self.max_failures {
                    node.failures = failures;
                    node.updated_at = Utc::now();

                    self.upsert_node_remote(&node).await;
                    self.nodes.insert(node.id.to_string(), node);

                    true
                }
                else {
                    tracing::info!("Removing node due to failures count {}: {:?}", failures, node);
                    self.remove_node(node.id.to_string().as_str()).await;

                    false
                }
            },
            None => false
        }
    }

    /// Inserts or updates a node in the local registry (not persisted to remote store).
    pub fn upsert_node_local(&mut self, identifier: String, mut node: Node) {
        let _ = node.features.setup_router();
        self.nodes.insert(identifier, node);
    }

    /// Finds all nodes matching the given query, optionally excluding the local node.
    pub fn find_nodes(&self, query: BTreeMap<&str, &str>, exclude_me: bool) -> BTreeMap<String, &Node> {
        let mut result = BTreeMap::new();
        let me = self.me.id.to_string();

        for (id, node) in &self.nodes {
            if (!exclude_me || !me.eq(id)) && node.match_node_query(&query) {
                result.insert(id.to_string(), node);
            }
        }

        result
    }

    /// Returns a reference to the node with the given ID, if present in the local registry.
    pub fn get_node(&self, id: &str) -> Option<&Node> {
        self.nodes.get(id)
    }

    /// Returns a map of nodes, optionally skipping the first `skip` entries and excluding the local node.
    pub fn get_nodes(&self, skip: u8, exclude_me: bool) -> BTreeMap<String, &Node> {
        let me = self.me.id.to_string();

        self.nodes.iter()
            .filter(|(id, _)| !exclude_me || !me.eq(*id))
            .skip(skip as usize)
            .map(|(id, node)| (id.to_string(), node))
            .collect()
    }

    /// Returns certificate bytes for nodes, optionally skipping entries and excluding the local node.
    #[deprecated(since = "0.0.28", note = "Use get_nodes_raw_certificates instead, which has identical behavior")]
    pub fn get_nodes_certificates(&self, skip: u8, exclude_me: bool) -> Vec<Vec<u8>> {
        self.get_nodes_raw_certificates(skip, exclude_me)
    }

    /// Returns raw certificate bytes for nodes, optionally skipping entries and excluding the local node.
    pub fn get_nodes_raw_certificates(&self, skip: u8, exclude_me: bool) -> Vec<Vec<u8>> {
        let me = self.me.id.to_string();

        self.nodes.iter()
            .filter(|(id, _)| !exclude_me || !me.eq(*id))
            .skip(skip as usize)
            .map(|(_, node)| node.get_certificate())
            .collect()
    }

    /// Returns a map of node endpoints (URI + optional key), skipping entries and optionally excluding the local node.
    pub fn get_nodes_endpoints(&self, skip: u8, exclude_me: bool) -> BTreeMap<String, (Uri, Option<Vec<u8>>)> {
        let me = self.me.id.to_string();

        #[allow(deprecated)]
        self.nodes.iter()
            .filter(|(id, _)| !exclude_me || !me.eq(*id))
            .skip(skip as usize)
            .map(|(id, node)| {
                let address = node.get_address();
                let key = self.get_key(id);
                (id.to_string(), (address, key))
            })
            .collect()
    }

    /// Removes a node from both the local registry and the remote store.
    pub async fn remove_node(&mut self, identifier: &str) -> Option<Node> {
        match self.nodes.remove(identifier) {
            Some(node) => {
                self.remove_node_remote(node.id.to_string().as_str()).await;
                Some(node)
            },
            None => None
        }
    }

    /// Picks a random node from those matching the query.
    pub fn pick_node(&self, query: BTreeMap<&str, &str>) -> Option<&Node> {
        let eligible_nodes: Vec<&Node> = self.find_nodes(query, false).values().copied().collect();
        let select = RandomGenerator::new(Some(RNG::Std(StdRng::from_entropy()))).get_random_usize(0, eligible_nodes.len());

        eligible_nodes.get(select).copied()
    }

    /// Picks a random node that has the specified capability.
    pub fn pick_node_for_capability(&self, capability: &str) -> Option<&Node> {
        let mut query = BTreeMap::new();
        query.insert("capability", capability);
        self.pick_node(query)
    }

    /// Registers a feature on the local node and updates the remote store.
    pub async fn add_feature(&mut self, feature: Arc<dyn product_os_capabilities::Feature>, base_path: String, router: &mut product_os_router::ProductOSRouter) {
        let _ = self.me.features.add(feature, base_path, router).await;
        self.update_me().await;
    }

    /// Registers a mutable feature on the local node and updates the remote store.
    pub async fn add_feature_mut(&mut self, feature: Arc<Mutex<dyn product_os_capabilities::Feature>>, base_path: String, router: &mut product_os_router::ProductOSRouter) {
        let _ = self.me.features.add_mut(feature, base_path, router).await;
        self.update_me().await;
    }

    /// Picks a random node that has the specified feature.
    pub fn pick_node_for_feature(&self, feature: &str) -> Option<&Node> {
        let mut query = BTreeMap::new();
        query.insert("feature", feature);
        self.pick_node(query)
    }

    /// Removes a feature from the local node and updates the remote store.
    pub async fn remove_feature(&mut self, identifier: &str) {
        let _ = self.me.features.remove(identifier);
        self.update_me().await;
    }

    /// Registers a service on the local node and updates the remote store.
    pub async fn add_service(&mut self, service: Arc<dyn product_os_capabilities::Service>) {
        self.me.services.add(service).await;
        self.update_me().await;
    }

    /// Registers a mutable service on the local node and updates the remote store.
    pub async fn add_service_mut(&mut self, service: Arc<Mutex<dyn product_os_capabilities::Service>>) {
        let _ = self.me.services.add_mut(service).await;
        self.update_me().await;
    }

    /// Sets the active status of a service on the local node and updates the remote store.
    pub async fn set_service_active(&mut self, identifier: String, status: bool) {
        let id = identifier.as_str();
        match self.me.services.get_mut(id) {
            None => (),
            Some(s) => {
                s.active = status;
                self.update_me().await;
            }
        }
    }

    /// Removes a service from the local node and updates the remote store.
    pub async fn remove_service(&mut self, identifier: &str) {
        self.me.services.remove(identifier);
        self.update_me().await;
    }

    /// Removes all inactive services matching the query and updates the remote store.
    pub async fn remove_inactive_services(&mut self, query: BTreeMap<&str, &str>) {
        let mut matches = Vec::new();

        for (identifier, _) in self.find_nodes(query, true) {
            matches.push(identifier.to_owned());
        }

        for identifier in &matches {
            self.remove_service(identifier).await;
        }

        if !matches.is_empty() { self.update_me().await };
    }

    /// Starts all registered services on the local node.
    pub async fn start_services(&mut self) -> Result<(), ()> {
        for (_, service) in self.me.services.list_mut() {
            match service.start().await {
                Ok(_) => {}
                Err(_) => return Err(())
            }
        }
        
        Ok(())
    }

    /// Starts the service with the given identifier on the local node.
    pub async fn start_service(&mut self, identifier: &str) -> Result<(), ()> {
        match self.me.services.get_mut(identifier) {
            None => Err(()),
            Some(s) => s.start().await
        }
    }
    
    /// Stops the service with the given identifier on the local node.
    pub async fn stop_service(&mut self, identifier: &str) -> Result<(), ()> {
        match self.me.services.get_mut(identifier) {
            None => Err(()),
            Some(s) => s.stop().await
        }
    }

    /// Restarts the service with the given identifier on the local node.
    pub async fn restart_service(&mut self, identifier: &str) -> Result<(), ()> {
        match self.me.services.get_mut(identifier) {
            None => Err(()),
            Some(s) => s.restart().await
        }
    }

    /// Calls a service action with the given input on the local node.
    pub async fn call_service(&mut self, identifier: &str, action: &What, input: &Option<serde_json::Value>) -> Result<Option<serde_json::Value>, ServiceError> {
        match self.me.services.get_mut(identifier) {
            None => Err(ServiceError::GenericError(format!("Service {} not found", identifier))),
            Some(s) => s.call(action, input).await
        }
    }

    /// Discovers nodes from the remote key-value store and imports them into the local registry.
    pub async fn discover_nodes(&mut self) {
        let mut nodes = BTreeMap::new();

        match self.store.group_find(None) {
            Ok(ns) => {
                nodes = ns;
            },
            Err(_) => {
                tracing::error!("Error getting nodes from store");
            }
        }

        for (id, node) in nodes {
            match serde_json::from_str(node.as_str()) {
                Ok(n) => {
                    let mut node: Node = n;
                    let _ = node.features.setup_router();
                    tracing::trace!("Importing remote node: {:?}", node.id);
                    self.upsert_node_local(node.id.to_string(), node);
                },
                Err(e) => {
                    tracing::error!("Error importing remote node {} - purging: {:?}", id, e);
                    self.remove_node_remote(id.as_str()).await;
                }
            }
        }
    }

    /// Returns the cryptographic key associated with the given node identifier.
    pub fn get_key(&self, identifier: &str) -> Option<Vec<u8>> {
        self.key_store.get_key(identifier).map(|k| k.to_vec())
    }

    /// Creates a new Diffie-Hellman key session and returns the session ID and public key.
    pub fn create_key_session(&mut self) -> (String, [u8; 32]) {
        self.key_store.create_session()
    }

    /// Generates a shared key from a Diffie-Hellman exchange session.
    pub fn generate_key(&mut self, session_identifier: &str, remote_public_key: &[u8], association: String, remote_session_identifier: Option<String>) {
        self.key_store.generate_key(session_identifier, remote_public_key, association, remote_session_identifier);
    }
}