kftray-portforward 0.30.2

pub mod kubernetes;
pub mod models;
pub mod templates;
pub mod websocket_client;

use kftray_commons::models::{
    config_model::Config,
    config_state_model::ConfigState,
    response::CustomResponse,
};
use kftray_commons::utils::config_state::update_config_state_with_mode;
use kftray_commons::utils::db_mode::DatabaseMode;
use log::{
    error,
    info,
};
use tokio_util::sync::CancellationToken;

use crate::kube::shared_client::{
    SHARED_CLIENT_MANAGER,
    ServiceClientKey,
};

/// Start expose for given configs
pub async fn start_expose(
    configs: Vec<Config>, mode: DatabaseMode,
) -> Result<Vec<CustomResponse>, String> {
    let configs = configs
        .into_iter()
        .map(|mut config| {
            config.workload_type = Some("expose".to_string());
            config
        })
        .collect();
    crate::kube::start_port_forward_with_mode(configs, "tcp", mode, false).await
}

/// Deletes exactly what this attempt created and releases its cleanup record
/// only once that succeeded.
///
/// Deleting by label would also reach resources another installation created
/// for the same configuration id, which is local to each database.
async fn roll_back_exposure<T>(
    client: &kube::Client, config: &Config, resources: &models::ExposeResources,
    location: &kubernetes::ExposeLocation, mode: DatabaseMode,
    guard: crate::kube::stop::ClusterResourceGuard, reason: String,
) -> Result<T, String> {
    let config_id_str = config
        .id
        .map_or_else(|| "default".to_string(), |id| id.to_string());
    let (cleaned, message) = kubernetes::rollback_created_resources(
        client,
        &config.namespace,
        &resources.owned,
        &config_id_str,
        location,
        mode,
        reason,
    )
    .await;
    // The guard stays armed unless cleanup fully completed, which is what
    // keeps it retryable.
    if cleaned {
        guard.disarm().await;
    }
    Err(message)
}

pub(crate) async fn start_single_expose(
    config: Config, mode: DatabaseMode, cancellation: Option<&CancellationToken>,
) -> Result<CustomResponse, String> {
    use self::kubernetes::create_expose_resources;
    use self::websocket_client::WebSocketTunnelClient;
    use crate::kube::models::{
        NameSpace,
        Port,
        PortForward,
        Target,
        TargetSelector,
    };
    use crate::port_forward::CHILD_PROCESSES;

    let config_id = config.id.ok_or("Config has no ID")?;
    // Startup holds the per-config lifecycle lock, and a stop waits on that
    // same lock, so without observing cancellation here a stop would block for
    // the full readiness budget.
    let cancelled = || cancellation.is_some_and(CancellationToken::is_cancelled);
    if cancelled() {
        return Err(format!("Expose startup cancelled for config {config_id}"));
    }

    let client_key = ServiceClientKey::new(config.context.clone(), config.kubeconfig.clone());
    let connection = SHARED_CLIENT_MANAGER
        .get_connection(client_key)
        .await
        .map_err(|e| format!("Failed to get K8s client: {}", e))?;
    let client = connection.client.clone();
    let location = kubernetes::ExposeLocation::resolve(&connection.cluster_url, &config.namespace);
    // Normalized the way `delete_cluster_resources_inner` and the forwarder's
    // own destination check compare it: a scheme/host case difference or an
    // explicit default port must not read as the resources living elsewhere.
    let destination = crate::kube::client::cluster_identity(&connection.cluster_url);

    info!("Creating expose resources for config {}", config_id);
    // Armed before creation so a dropped startup future, or a create whose
    // response is lost, still leaves a trail for stop-all.
    let mut guard = crate::kube::stop::ClusterResourceGuard::arm(
        config_id,
        config.clone(),
        Some(destination.clone()),
        mode,
    )
    .await?;
    // The individual create calls are never raced against cancellation:
    // dropping one mid-flight would abandon a request whose outcome is
    // unknown, including whether `create_ingress` reached its actual create
    // call after recording that it might have. Letting each finish keeps
    // every outcome definitive: an ambiguous one leaves the guard uncertain
    // for a later reconciliation pass, and a confirmed one is rolled back
    // below by UID. The pre-start cleanup wait and the pod readiness wait
    // inside `create_expose_resources` are cancellable, since nothing there
    // is a create whose outcome would become ambiguous.
    let created = create_expose_resources(&connection, &config, mode, cancellation).await;
    // Confirmed only when the outcome is definitive. A transport failure or a
    // server-side timeout can be answered while the object is still being
    // applied, and rollback cannot name a resource whose create never returned.
    let resources = match created {
        Ok(resources) => {
            guard.confirm();
            if cancelled() {
                let reason = format!("Expose startup cancelled for config {config_id}");
                return roll_back_exposure(
                    &client, &config, &resources, &location, mode, guard, reason,
                )
                .await;
            }
            resources
        }
        Err(error) => {
            // Disarmed only when the outcome was definitive and UID-scoped
            // rollback removed everything this attempt created: leaving the
            // record would make a later stop delete by label, which can reach
            // another installation's resources for the same config id.
            if !error.ambiguous && error.rolled_back {
                guard.disarm().await;
            } else if !error.ambiguous {
                guard.confirm();
            }
            return Err(error.message);
        }
    };

    info!(
        "Resources created: deployment={}, service={}, pod={}",
        resources.deployment_name, resources.service_name, resources.pod_name
    );

    // Installation-scoped: the tunnel must reach this installation's relay, not
    // another one that happens to share the locally assigned configuration id.
    let label_selector = match kubernetes::expose_owner_selector(&config_id.to_string(), mode).await
    {
        Ok(selector) => selector,
        Err(error) => {
            return roll_back_exposure(&client, &config, &resources, &location, mode, guard, error)
                .await;
        }
    };
    let target = Target {
        selector: TargetSelector::PodLabel(label_selector),
        port: Port::Number(i32::from(resources.websocket_port)),
        namespace: NameSpace(Some(config.namespace.clone())),
    };

    // Pinned to the server the relay was created on: the forwarder resolves
    // the context again, and a cached client that expired during creation or
    // readiness could resolve it to another cluster with a relay under the
    // same labels.
    let port_forward = PortForward::new(
        target,
        Some(0),
        None,
        config.context.clone(),
        config.kubeconfig.clone(),
        config_id,
        "expose".to_string(),
    )
    .expecting_destination(Some(destination.clone()));

    // Never raced against cancellation: dropping this future mid-flight would
    // abandon a listener task `PortForwarder` has no `Drop` to stop, leaking
    // it. Instead this always runs to completion, and a cancellation that
    // arrived while it was starting is applied after, by aborting the
    // process it just handed back.
    let started = port_forward.port_forward_tcp(None).await;
    let (websocket_port, mut pf_process) = match started {
        Ok(started) => started,
        Err(error) => {
            let reason = format!("Failed to start port-forward: {error}");
            return roll_back_exposure(
                &client, &config, &resources, &location, mode, guard, reason,
            )
            .await;
        }
    };
    if cancelled() {
        let reason = format!("Expose startup cancelled for config {config_id}");
        pf_process.cleanup_and_abort().await;
        return roll_back_exposure(&client, &config, &resources, &location, mode, guard, reason)
            .await;
    }

    info!(
        "Port-forward established: localhost:{} → pod:{}",
        websocket_port, resources.websocket_port
    );

    let local_service_port = config.local_port.unwrap_or(8080);
    let local_service_address = config
        .local_address
        .clone()
        .unwrap_or_else(|| "127.0.0.1".to_string());
    let ws_client = WebSocketTunnelClient::new(
        websocket_port,
        local_service_address.clone(),
        local_service_port,
    );

    info!(
        "Starting WebSocket tunnel: pod → localhost:{} → {}:{}",
        websocket_port, local_service_address, local_service_port
    );

    let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
    let ws_handle = tokio::spawn(async move {
        if let Err(e) = ws_client.start(ready_tx).await {
            error!("WebSocket client error: {}", e);
            // The reconnect budget is exhausted and nothing else watches this
            // task: `proxy_recovery` has no expose participation, so without
            // this the config keeps reporting healthy while nothing forwards.
            // Stopping from a detached task, not inline, matters here:
            // `PortForwardProcess::cancel` aborts this very `JoinHandle`, and
            // self-aborting partway through a stop would leave it half done.
            tokio::spawn(async move {
                if let Err(stop_error) =
                    crate::kube::stop_port_forward_with_mode(config_id.to_string(), mode).await
                {
                    error!(
                        "Failed to clean up config {config_id} after its reverse tunnel ended: \
                         {stop_error}"
                    );
                }
            });
        }
    });
    // Attach before awaiting readiness: dropping this startup future must abort
    // the tunnel task through the process, not detach it.
    pf_process.set_ws_client_handle(ws_handle);

    let ready = async {
        match tokio::time::timeout(std::time::Duration::from_secs(30), ready_rx).await {
            Ok(Ok(result)) => result,
            Ok(Err(error)) => Err(format!(
                "Expose startup task ended before connecting: {error}"
            )),
            Err(_) => Err("Timed out connecting the reverse WebSocket tunnel".to_owned()),
        }
    };
    let startup = match cancellation {
        Some(token) => tokio::select! {
            biased;
            _ = token.cancelled() => Err(format!("Expose startup cancelled for config {config_id}")),
            startup = ready => startup,
        },
        None => ready.await,
    };
    if let Err(error) = startup {
        pf_process.cleanup_and_abort().await;
        return roll_back_exposure(&client, &config, &resources, &location, mode, guard, error)
            .await;
    }

    let config_state = ConfigState {
        id: None,
        config_id,
        is_running: true,
        process_id: Some(std::process::id()),
        is_retrying: false,
        retry_count: None,
        last_error: None,
    };
    let recorded =
        kftray_commons::utils::config_state::set_running_snapshot(config_id, &config, mode).await;
    if let Err(error) = match recorded {
        Ok(()) => update_config_state_with_mode(&config_state, mode).await,
        Err(error) => Err(error),
    } {
        pf_process.cleanup_and_abort().await;
        kftray_commons::utils::config_state::clear_running_snapshot(config_id, mode).await;
        return roll_back_exposure(&client, &config, &resources, &location, mode, guard, error)
            .await;
    }
    pf_process.set_config(config.clone());
    CHILD_PROCESSES.insert(config_id, pf_process);
    // The record stays, now settled rather than uncertain: the registered
    // process describes these resources only in memory, and a crash after the
    // row is edited would otherwise leave cleanup looking for them at the new
    // destination. It is dropped when cleanup actually succeeds.
    guard.confirm();
    drop(guard);

    info!("Expose tunnel fully established for config {}", config_id);

    Ok(CustomResponse {
        id: Some(config_id),
        service: config.service.unwrap_or(resources.service_name),
        namespace: config.namespace.clone(),
        local_port: local_service_port,
        remote_port: resources.websocket_port,
        context: config.context.unwrap_or_default(),
        stdout: String::new(),
        stderr: String::new(),
        status: 0,
        protocol: "tcp".to_string(),
    })
}

pub async fn stop_expose(config_id: i64, mode: DatabaseMode) -> Result<CustomResponse, String> {
    crate::kube::stop_port_forward_with_mode(config_id.to_string(), mode).await
}