kindly-guard-server 0.11.14

KindlyGuard MCP server - Enterprise-grade security for AI model interactions
Documentation
// Copyright 2025 Kindly Software Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! Enhanced telemetry implementation with advanced features
//! This module provides optimized telemetry using advanced techniques

use super::*;
use crate::traits::SecurityEventProcessor;
use std::sync::Arc;
use tracing::{debug, info};

// Stub for advanced telemetry buffer that would use advanced tech
struct AdvancedTelemetryBuffer;

impl AdvancedTelemetryBuffer {
    fn new(_buffer_size: usize, _sampling_rate: f64) -> Result<Self> {
        Ok(Self)
    }

    fn track_span_start(&self, _name: &str, _timestamp: u64) {
        // Stub implementation
    }

    fn track_span_end(&self, _name: &str, _duration_ns: u64, _attributes: Vec<(String, String)>) {
        // Stub implementation
    }

    fn increment_counter(&self, _name: &str, _value: u64, _labels: Vec<(String, String)>) {
        // Stub implementation
    }

    fn record_gauge(&self, _name: &str, _value: f64, _labels: Vec<(String, String)>) {
        // Stub implementation
    }

    fn record_histogram(&self, _name: &str, _value: f64, _labels: Vec<(String, String)>) {
        // Stub implementation
    }

    fn add_event(&self, _name: &str, _attributes: Vec<(&str, &str)>) {
        // Stub implementation
    }

    async fn flush(&self) -> Result<()> {
        Ok(())
    }

    async fn shutdown(&self) -> Result<()> {
        Ok(())
    }

    fn set_span_status(&self, _name: &str, _is_error: bool, _message: Option<&str>) {
        // Stub implementation
    }

    fn prepare_export_batch(&self) -> Result<ExportBatch> {
        Ok(ExportBatch::default())
    }

    async fn export_batch(&self, _batch: ExportBatch, _endpoint: &str) -> Result<()> {
        Ok(())
    }
}

// Stub for export batch
#[derive(Default)]
struct ExportBatch;
impl ExportBatch {
    fn span_count(&self) -> usize {
        0
    }
    fn metric_count(&self) -> usize {
        0
    }
}

/// Enhanced telemetry provider using advanced buffering
pub struct EnhancedTelemetryProvider {
    config: TelemetryConfig,
    buffer: Arc<AdvancedTelemetryBuffer>,
    event_processor: Arc<dyn SecurityEventProcessor>,
}

impl EnhancedTelemetryProvider {
    pub fn new(
        config: TelemetryConfig,
        event_processor: Arc<dyn SecurityEventProcessor>,
    ) -> Result<Self> {
        let buffer = AdvancedTelemetryBuffer::new(
            config.export_interval_seconds as usize * 1000, // Convert to buffer size
            config.sampling_rate,
        )?;

        Ok(Self {
            config,
            buffer: Arc::new(buffer),
            event_processor,
        })
    }
}

#[async_trait]
impl TelemetryProvider for EnhancedTelemetryProvider {
    fn start_span(&self, name: &str) -> TelemetrySpan {
        let span = TelemetrySpan {
            name: name.to_string(),
            start_time: std::time::Instant::now(),
            attributes: vec![
                ("service.name".to_string(), self.config.service_name.clone()),
                (
                    "service.version".to_string(),
                    self.config.service_version.clone(),
                ),
                ("telemetry.mode".to_string(), "enhanced".to_string()),
            ],
        };

        // Use advanced buffer for efficient span tracking
        self.buffer
            .track_span_start(&span.name, span.start_time.elapsed().as_nanos() as u64);

        span
    }

    fn end_span(&self, span: TelemetrySpan) {
        if !self.config.tracing_enabled {
            return;
        }

        let duration = span.start_time.elapsed();

        // Use optimized buffer for high-performance recording
        self.buffer
            .track_span_end(&span.name, duration.as_nanos() as u64, span.attributes);

        debug!(
            "Enhanced span completed: {} (duration: {:?})",
            span.name, duration
        );
    }

    fn record_metric(&self, metric: TelemetryMetric) {
        if !self.config.metrics_enabled {
            return;
        }

        // Use optimized metric recording
        match metric.value {
            MetricValue::Counter(value) => {
                self.buffer
                    .increment_counter(&metric.name, value, metric.labels);
            },
            MetricValue::Gauge(value) => {
                self.buffer.record_gauge(&metric.name, value, metric.labels);
            },
            MetricValue::Histogram(value) => {
                self.buffer
                    .record_histogram(&metric.name, value, metric.labels);
            },
        }
    }

    fn add_event(&self, name: &str, attributes: Vec<(&str, &str)>) {
        if !self.config.tracing_enabled {
            return;
        }

        // Convert to owned strings for buffer
        let attrs: Vec<(String, String)> = attributes
            .into_iter()
            .map(|(k, v)| (k.to_string(), v.to_string()))
            .collect();

        // Use event processor for correlation
        let event = crate::traits::SecurityEvent {
            event_type: name.to_string(),
            client_id: attrs
                .iter()
                .find(|(k, _)| k == "client.id")
                .map(|(_, v)| v.clone())
                .unwrap_or_else(|| "unknown".to_string()),
            timestamp: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_secs(),
            metadata: serde_json::json!(attrs),
        };

        // Fire and forget for performance
        let processor = self.event_processor.clone();
        tokio::spawn(async move {
            let _ = processor.process_event(event).await;
        });
    }

    fn set_status(&self, span: &TelemetrySpan, is_error: bool, message: Option<&str>) {
        if !self.config.tracing_enabled {
            return;
        }

        self.buffer.set_span_status(&span.name, is_error, message);
    }

    async fn flush(&self) -> Result<()> {
        info!("Flushing enhanced telemetry data");

        // Export using optimized batching
        if let Some(endpoint) = &self.config.export_endpoint {
            let batch = self.buffer.prepare_export_batch()?;

            // In production, this would use optimized OTLP export
            debug!(
                "Enhanced export to {}: {} spans, {} metrics",
                endpoint,
                batch.span_count(),
                batch.metric_count()
            );

            self.buffer.export_batch(batch, endpoint).await?;
        }

        Ok(())
    }

    async fn shutdown(&self) -> Result<()> {
        info!("Shutting down enhanced telemetry provider");

        // Final flush
        self.flush().await?;

        // Graceful shutdown of buffer
        self.buffer.shutdown().await?;

        Ok(())
    }
}

/// Factory for enhanced telemetry
pub struct EnhancedTelemetryFactory {
    event_processor: Arc<dyn SecurityEventProcessor>,
}

impl EnhancedTelemetryFactory {
    pub fn new(event_processor: Arc<dyn SecurityEventProcessor>) -> Self {
        Self { event_processor }
    }
}

impl TelemetryProviderFactory for EnhancedTelemetryFactory {
    fn create(&self, config: &TelemetryConfig) -> Result<Arc<dyn TelemetryProvider>> {
        Ok(Arc::new(EnhancedTelemetryProvider::new(
            config.clone(),
            self.event_processor.clone(),
        )?))
    }
}