start-os/core/src/net/host/mod.rs

use std::collections::{BTreeMap, BTreeSet};
use std::future::Future;
use std::panic::RefUnwindSafe;

use clap::Parser;
use imbl_value::InternedString;
use itertools::Itertools;
use rpc_toolkit::{Context, Empty, HandlerExt, OrEmpty, ParentHandler, from_fn_async};
use serde::{Deserialize, Serialize};
use ts_rs::TS;

use crate::context::RpcContext;
use crate::db::model::DatabaseModel;
use crate::net::forward::AvailablePorts;
use crate::net::host::address::{HostAddress, PublicDomainConfig, address_api};
use crate::net::host::binding::{BindInfo, BindOptions, binding};
use crate::net::service_interface::HostnameInfo;
use crate::net::tor::OnionAddress;
use crate::prelude::*;
use crate::{HostId, PackageId};

pub mod address;
pub mod binding;

#[derive(Debug, Default, Deserialize, Serialize, HasModel, TS)]
#[serde(rename_all = "camelCase")]
#[model = "Model<Self>"]
#[ts(export)]
pub struct Host {
    pub bindings: BTreeMap<u16, BindInfo>,
    #[ts(type = "string[]")]
    pub onions: BTreeSet<OnionAddress>,
    pub public_domains: BTreeMap<InternedString, PublicDomainConfig>,
    pub private_domains: BTreeSet<InternedString>,
    /// COMPUTED: NetService::update
    pub hostname_info: BTreeMap<u16, Vec<HostnameInfo>>, // internal port -> Hostnames
}

impl AsRef<Host> for Host {
    fn as_ref(&self) -> &Host {
        self
    }
}
impl Host {
    pub fn new() -> Self {
        Self::default()
    }
    pub fn addresses<'a>(&'a self) -> impl Iterator<Item = HostAddress> + 'a {
        self.onions
            .iter()
            .cloned()
            .map(|address| HostAddress::Onion { address })
            .chain(
                self.public_domains
                    .iter()
                    .map(|(address, config)| HostAddress::Domain {
                        address: address.clone(),
                        public: Some(config.clone()),
                        private: self.private_domains.contains(address),
                    }),
            )
            .chain(
                self.private_domains
                    .iter()
                    .filter(|a| !self.public_domains.contains_key(*a))
                    .map(|address| HostAddress::Domain {
                        address: address.clone(),
                        public: None,
                        private: true,
                    }),
            )
    }
}

#[derive(Debug, Default, Deserialize, Serialize, HasModel, TS)]
#[model = "Model<Self>"]
#[ts(export)]
pub struct Hosts(pub BTreeMap<HostId, Host>);

impl Map for Hosts {
    type Key = HostId;
    type Value = Host;
    fn key_str(key: &Self::Key) -> Result<impl AsRef<str>, Error> {
        Ok(key)
    }
    fn key_string(key: &Self::Key) -> Result<InternedString, Error> {
        Ok(key.clone().into())
    }
}

pub fn host_for<'a>(
    db: &'a mut DatabaseModel,
    package_id: Option<&PackageId>,
    host_id: &HostId,
) -> Result<&'a mut Model<Host>, Error> {
    let Some(package_id) = package_id else {
        return Ok(db
            .as_public_mut()
            .as_server_info_mut()
            .as_network_mut()
            .as_host_mut());
    };
    fn host_info<'a>(
        db: &'a mut DatabaseModel,
        package_id: &PackageId,
    ) -> Result<&'a mut Model<Hosts>, Error> {
        Ok::<_, Error>(
            db.as_public_mut()
                .as_package_data_mut()
                .as_idx_mut(package_id)
                .or_not_found(package_id)?
                .as_hosts_mut(),
        )
    }
    let tor_key = if host_info(db, package_id)?.as_idx(host_id).is_none() {
        Some(
            db.as_private_mut()
                .as_key_store_mut()
                .as_onion_mut()
                .new_key()?,
        )
    } else {
        None
    };
    host_info(db, package_id)?.upsert(host_id, || {
        let mut h = Host::new();
        h.onions
            .insert(tor_key.or_not_found("generated tor key")?.onion_address());
        Ok(h)
    })
}

pub fn all_hosts(db: &mut DatabaseModel) -> impl Iterator<Item = Result<&mut Model<Host>, Error>> {
    use patch_db::DestructureMut;
    let destructured = db.as_public_mut().destructure_mut();
    [Ok(destructured.server_info.as_network_mut().as_host_mut())]
        .into_iter()
        .chain(
            [destructured.package_data.as_entries_mut()]
                .into_iter()
                .flatten_ok()
                .map(|entry| entry.and_then(|(_, v)| v.as_hosts_mut().as_entries_mut()))
                .flatten_ok()
                .map_ok(|(_, v)| v),
        )
}

impl Model<Host> {
    pub fn add_binding(
        &mut self,
        available_ports: &mut AvailablePorts,
        internal_port: u16,
        options: BindOptions,
    ) -> Result<(), Error> {
        self.as_bindings_mut().mutate(|b| {
            let info = if let Some(info) = b.remove(&internal_port) {
                info.update(available_ports, options)?
            } else {
                BindInfo::new(available_ports, options)?
            };
            b.insert(internal_port, info);
            Ok(())
        })
    }
}

#[derive(Deserialize, Serialize, Parser)]
pub struct RequiresPackageId {
    #[arg(help = "help.arg.package-id")]
    package: PackageId,
}

#[derive(Deserialize, Serialize, Parser)]
pub struct RequiresHostId {
    #[arg(help = "help.arg.host-id")]
    host: HostId,
}

pub trait HostApiKind: 'static {
    type Params: Send + Sync + 'static;
    type InheritedParams: Send + Sync + 'static;
    type Inheritance: RefUnwindSafe + OrEmpty<Self::Inheritance> + Send + Sync + 'static;
    fn inheritance(params: Self::Params, inherited: Self::InheritedParams) -> Self::Inheritance;
    fn host_for<'a>(
        inheritance: &Self::Inheritance,
        db: &'a mut DatabaseModel,
    ) -> Result<&'a mut Model<Host>, Error>;
    fn sync_host(
        ctx: &RpcContext,
        inheritance: Self::Inheritance,
    ) -> impl Future<Output = Result<(), Error>> + Send;
}
pub struct ForPackage;
impl HostApiKind for ForPackage {
    type Params = RequiresHostId;
    type InheritedParams = PackageId;
    type Inheritance = (PackageId, HostId);
    fn inheritance(
        RequiresHostId { host }: Self::Params,
        package: Self::InheritedParams,
    ) -> Self::Inheritance {
        (package, host)
    }
    fn host_for<'a>(
        (package, host): &Self::Inheritance,
        db: &'a mut DatabaseModel,
    ) -> Result<&'a mut Model<Host>, Error> {
        host_for(db, Some(package), host)
    }
    async fn sync_host(ctx: &RpcContext, (package, host): Self::Inheritance) -> Result<(), Error> {
        let service = ctx.services.get(&package).await;
        let service_ref = service.as_ref().or_not_found(&package)?;
        service_ref.sync_host(host).await?;
        Ok(())
    }
}
pub struct ForServer;
impl HostApiKind for ForServer {
    type Params = Empty;
    type InheritedParams = Empty;
    type Inheritance = Empty;
    fn inheritance(_: Self::Params, _: Self::InheritedParams) -> Self::Inheritance {
        Empty {}
    }
    fn host_for<'a>(
        _: &Self::Inheritance,
        db: &'a mut DatabaseModel,
    ) -> Result<&'a mut Model<Host>, Error> {
        host_for(db, None, &HostId::default())
    }
    async fn sync_host(ctx: &RpcContext, _: Self::Inheritance) -> Result<(), Error> {
        ctx.os_net_service.sync_host(HostId::default()).await
    }
}

pub fn host_api<C: Context>() -> ParentHandler<C, RequiresPackageId> {
    ParentHandler::<C, RequiresPackageId>::new()
        .subcommand(
            "list",
            from_fn_async(list_hosts)
                .with_inherited(|RequiresPackageId { package }, _| package)
                .with_custom_display_fn(|_, ids| {
                    for id in ids {
                        println!("{id}")
                    }
                    Ok(())
                })
                .with_about("about.list-host-ids-for-service"),
        )
        .subcommand(
            "address",
            address_api::<C, ForPackage>()
                .with_inherited(|RequiresPackageId { package }, _| package),
        )
        .subcommand(
            "binding",
            binding::<C, ForPackage>().with_inherited(|RequiresPackageId { package }, _| package),
        )
}

pub fn server_host_api<C: Context>() -> ParentHandler<C> {
    ParentHandler::<C>::new()
        .subcommand("address", address_api::<C, ForServer>())
        .subcommand("binding", binding::<C, ForServer>())
}

pub async fn list_hosts(
    ctx: RpcContext,
    _: Empty,
    package: PackageId,
) -> Result<BTreeSet<HostId>, Error> {
    ctx.db
        .peek()
        .await
        .into_public()
        .into_package_data()
        .into_idx(&package)
        .or_not_found(&package)?
        .into_hosts()
        .keys()
}