ziti-sdk 0.3.0

//! Ziti dial functionality
//!
//! Provides the main dial function for establishing connections to Ziti services.

use crate::context::Context;
use crate::connection::ZitiStream;
use crate::error::{ZitiError, ZitiResult};
use crate::service::list_services;
use crate::transport::http::controller_client;
use crate::transport::protocol::{ContentType, ZitiMessage};
use crate::transport::{TlsConfig, WebSocketTransport};
use bytes::Bytes;
use tokio_tungstenite::tungstenite::Message;
use url::Url;

/// Dial a Ziti service by name
///
/// This function orchestrates the complete process of connecting to a Ziti service:
/// 1. Look up service details using the service discovery API
/// 2. Request network session and obtain edge router information
/// 3. Choose an edge router and establish WebSocket connection
/// 4. Perform Ziti protocol handshake
/// 5. Return a ZitiStream for communication
///
/// # Arguments
///
/// * `service_name` - The name of the service to connect to
/// * `context` - The Ziti context containing identity and session managers
///
/// # Returns
///
/// * `ZitiResult<ZitiStream>` - A stream for communicating with the service
///
/// # Example
///
/// ```rust,no_run
/// use ziti_sdk::{Context, connection::dial};
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
///     let context = Context::from_file("identity.json").await?;
///     let stream = dial("echo-service", &context).await?;
///     // Use stream for communication
///     Ok(())
/// }
/// ```
pub async fn dial(service_name: &str, context: &Context) -> ZitiResult<ZitiStream> {
    // Step 1: Find the service by name using list_services
    let services = list_services(context.session_manager()).await?;
    let service = services
        .iter()
        .find(|s| s.name == service_name)
        .ok_or_else(|| ZitiError::ServiceNotFound {
            service_name: service_name.to_string(),
        })?;

    // Step 2: Request network session and get edge router information
    let edge_routers = get_service_terminators(&service.id, context).await?;
    
    if edge_routers.is_empty() {
        return Err(ZitiError::ConnectionFailed(
            format!("No edge routers available for service '{}'", service_name)
        ));
    }

    // Step 3: Choose an edge router (for now, just pick the first one)
    let edge_router = &edge_routers[0];
    
    // Step 4: Create TLS config from identity
    let tls_config = TlsConfig::from_identity(context.identity_manager())?;
    
    // Step 5: Establish WebSocket connection to edge router
    let ws_url = Url::parse(&format!("wss://{}:{}/ws", edge_router.hostname, edge_router.port))
        .map_err(|e| ZitiError::ConfigError(format!("Invalid edge router URL: {}", e)))?;
    
    let mut transport = WebSocketTransport::connect(ws_url, tls_config).await?;
    
    // Step 6: Perform Ziti connection handshake
    perform_ziti_handshake(&mut transport, &service.id).await?;
    
    // Step 7: Return ZitiStream wrapping the transport
    Ok(ZitiStream::from_transport(transport))
}

/// Edge router information
#[derive(Debug, Clone, serde::Deserialize)]
pub struct EdgeRouter {
    pub hostname: String,
    pub port: u16,
    #[serde(default)]
    pub supported_protocols: Vec<String>,
}

/// Terminator information from the controller
#[derive(Debug, serde::Deserialize)]
struct Terminator {
    #[allow(dead_code)]
    pub id: String,
    pub router: Option<EdgeRouter>,
}

/// Response structure for terminators API call
#[derive(Debug, serde::Deserialize)]
struct TerminatorsResponse {
    data: Vec<Terminator>,
}

/// Get terminators (edge routers) for a service
async fn get_service_terminators(service_id: &str, context: &Context) -> ZitiResult<Vec<EdgeRouter>> {
    let api_session = context.session_manager().get_api_session().await?;
    let client = controller_client(context.identity_manager(), context.connect_timeout()).await?;

    let terminators_url = format!(
        "{}/services/{}/terminators",
        context.identity_manager().zt_api().trim_end_matches('/'),
        service_id
    );

    // Send request to get terminators
    let response = client
        .get(&terminators_url)
        .header("Content-Type", "application/json")
        .header("zt-session", &api_session.token)
        .send()
        .await
        .map_err(|e| {
            ZitiError::ConnectionFailed(format!("Failed to get service terminators: {}", e))
        })?;

    if !response.status().is_success() {
        let status = response.status();
        let error_text = response
            .text()
            .await
            .unwrap_or_else(|_| "Unknown error".to_string());
        return Err(ZitiError::ProtocolError {
            message: format!(
                "Terminators request failed with status {}: {}",
                status, error_text
            ),
        });
    }

    // Parse terminators response
    let terminators_response: TerminatorsResponse = response
        .json()
        .await
        .map_err(|e| ZitiError::ProtocolError {
            message: format!("Failed to parse terminators response: {}", e),
        })?;

    // Convert terminators to edge routers
    let edge_routers = terminators_response
        .data
        .into_iter()
        .filter_map(|t| t.router)
        .collect();

    Ok(edge_routers)
}

/// Perform Ziti connection handshake using the binary [`ZitiMessage`] framing.
async fn perform_ziti_handshake(
    transport: &mut WebSocketTransport,
    service_id: &str,
) -> ZitiResult<()> {
    let hello_bytes = build_hello_message(service_id)?;
    transport.send(Message::Binary(hello_bytes)).await?;

    match transport.receive().await? {
        Some(Message::Binary(data)) => parse_hello_response(&data),
        Some(_) => Err(ZitiError::ProtocolError {
            message: "Unexpected message type in handshake response".to_string(),
        }),
        None => Err(ZitiError::ConnectionFailed(
            "Connection closed during handshake".to_string(),
        )),
    }
}

/// Build a `Hello` [`ZitiMessage`] for the dial handshake.
fn build_hello_message(service_id: &str) -> ZitiResult<Bytes> {
    let mut hello = ZitiMessage::new(ContentType::Hello, 1, Bytes::new());
    hello.header
        .add_header("service_id".to_string(), service_id.to_string());
    hello.header
        .add_header("version".to_string(), "1.0".to_string());
    hello.serialize()
}

/// Parse and validate a `Hello` handshake response from the edge router.
fn parse_hello_response(data: &[u8]) -> ZitiResult<()> {
    let msg = ZitiMessage::deserialize(Bytes::copy_from_slice(data))?;
    if msg.content_type() != ContentType::Hello {
        return Err(ZitiError::ProtocolError {
            message: format!(
                "Expected Hello response, got content_type={:?}",
                msg.content_type()
            ),
        });
    }
    crate::connection::listen::check_response_status(msg, "hello")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_edge_router_deserialization() {
        let json = r#"{
            "hostname": "edge-router.example.com",
            "port": 443,
            "supported_protocols": ["tls", "ws"]
        }"#;

        let edge_router: EdgeRouter = serde_json::from_str(json).unwrap();
        assert_eq!(edge_router.hostname, "edge-router.example.com");
        assert_eq!(edge_router.port, 443);
        assert_eq!(edge_router.supported_protocols, vec!["tls", "ws"]);
    }

    #[test]
    fn test_hello_message_roundtrip() {
        let bytes = build_hello_message("test-service").unwrap();
        let parsed = ZitiMessage::deserialize(bytes).unwrap();
        assert_eq!(parsed.content_type(), ContentType::Hello);
        assert_eq!(parsed.header.get_header("service_id").unwrap(), "test-service");
        assert_eq!(parsed.header.get_header("version").unwrap(), "1.0");
    }

    #[test]
    fn test_parse_hello_response_ok() {
        let mut msg = ZitiMessage::new(ContentType::Hello, 2, Bytes::new());
        msg.header
            .add_header("status".to_string(), "ok".to_string());
        let bytes = msg.serialize().unwrap();
        assert!(parse_hello_response(&bytes).is_ok());
    }

    #[test]
    fn test_parse_hello_response_error_carries_detail() {
        let mut msg = ZitiMessage::new(ContentType::Hello, 2, Bytes::new());
        msg.header
            .add_header("status".to_string(), "error".to_string());
        msg.header
            .add_header("error".to_string(), "policy denied".to_string());
        let bytes = msg.serialize().unwrap();
        let err = parse_hello_response(&bytes).unwrap_err();
        let s = format!("{}", err);
        assert!(s.contains("policy denied"), "got: {}", s);
    }

    #[test]
    fn test_parse_hello_response_wrong_type() {
        let msg = ZitiMessage::new(ContentType::Data, 1, Bytes::from_static(b"x"));
        let bytes = msg.serialize().unwrap();
        assert!(parse_hello_response(&bytes).is_err());
    }
}