main.rs

use core::pin::Pin;
use core::task::{Context, Poll};
use std::collections::HashMap;
use std::io;
use std::io::{Read, Write};
use std::net::SocketAddr;
use std::sync::OnceLock;

use anyhow::anyhow;
use csv::{Reader, ReaderBuilder, StringRecord};
use env_logger::Env;
use futures::Stream;
use log::debug;
use maplit::hashmap;
use pact_matching::matchingrules::DoMatch;
use pact_models::matchingrules::{MatchingRule, RuleList, RuleLogic};
use pact_models::prelude::ContentType;
use serde_json::Value;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;
use tonic::{transport::Server, Response};
use uuid::Uuid;

use prost_types;

use crate::csv_content::{generate_csv_content, has_headers, setup_csv_contents};
use crate::proto::body::ContentTypeHint;
use crate::proto::catalogue_entry::EntryType;
use crate::proto::pact_plugin_server::{PactPlugin, PactPluginServer};
use crate::proto::plugin_host_client::PluginHostClient;
use crate::proto::to_object;

mod csv_content;
mod parser;
mod proto;
mod utils;

/// Plugin instance UUID set by the driver in InitPluginRequest; included in every Log RPC call
static PLUGIN_INSTANCE_ID: OnceLock<String> = OnceLock::new();

/// Test run ID from testContext["testRunId"]; scoped per gRPC request task
tokio::task_local! {
  static TEST_RUN_ID: String;
}

fn extract_test_run_id(ctx: Option<&prost_types::Struct>) -> String {
  ctx
    .and_then(|s| s.fields.get("testRunId"))
    .and_then(|v| match &v.kind {
      Some(prost_types::value::Kind::StringValue(s)) => Some(s.clone()),
      _ => None,
    })
    .unwrap_or_default()
}

fn is_transport_target(target: &str) -> bool {
  target.starts_with("h2::") || target.starts_with("tower::") ||
  target.starts_with("tonic::") || target.starts_with("hyper_util::") ||
  target.starts_with("hyper::")
}

/// Wraps env_logger and also forwards records to the driver via the PluginHost Log RPC
struct PluginHostLogger {
  inner: env_logger::Logger,
  host_tx: Option<mpsc::UnboundedSender<proto::LogMessage>>,
}

impl log::Log for PluginHostLogger {
  fn enabled(&self, metadata: &log::Metadata) -> bool {
    self.inner.enabled(metadata)
  }

  fn log(&self, record: &log::Record) {
    if !self.inner.enabled(record.metadata()) {
      return;
    }
    if let Some(tx) = &self.host_tx {
      if record.level() != log::Level::Trace && !is_transport_target(record.target()) {
        let instance_id = PLUGIN_INSTANCE_ID.get().cloned().unwrap_or_default();
        let test_run_id = TEST_RUN_ID.try_with(|id| id.clone()).unwrap_or_default();
        let timestamp_ms = std::time::SystemTime::now()
          .duration_since(std::time::UNIX_EPOCH)
          .map(|d| d.as_millis() as i64)
          .unwrap_or(0);
        let _ = tx.send(proto::LogMessage {
          plugin_instance_id: instance_id,
          test_run_id,
          level: record.level().to_string().to_uppercase(),
          message: record.args().to_string(),
          target: record.target().to_string(),
          timestamp_ms,
        });
      }
    }
    self.inner.log(record);
  }

  fn flush(&self) {
    self.inner.flush();
  }
}

async fn init_logging() {
  if let Ok(id) = std::env::var("PACT_PLUGIN_INSTANCE_ID") {
    PLUGIN_INSTANCE_ID.set(id).ok();
  }

  let mut builder = env_logger::Builder::from_env(Env::new().filter("LOG_LEVEL"));
  builder.format(|buf, record| {
    use std::io::Write;
    let test_run_id = TEST_RUN_ID.try_with(|id| id.clone()).unwrap_or_default();
    if !test_run_id.is_empty() {
      writeln!(buf, "[{} {} {}] {}", record.level(), record.target(), test_run_id, record.args())
    } else {
      writeln!(buf, "[{} {}] {}", record.level(), record.target(), record.args())
    }
  });
  let inner = builder.build();

  let host_tx = match std::env::var("PACT_PLUGIN_HOST") {
    Ok(host_addr) => {
      let endpoint = format!("http://{}", host_addr);
      match PluginHostClient::connect(endpoint).await {
        Ok(mut client) => {
          let (tx, mut rx) = mpsc::unbounded_channel::<proto::LogMessage>();
          tokio::spawn(async move {
            while let Some(msg) = rx.recv().await {
              let _ = client.log(tonic::Request::new(msg)).await;
            }
          });
          Some(tx)
        }
        Err(err) => {
          eprintln!("Could not connect to PluginHost, log forwarding disabled: {err}");
          None
        }
      }
    }
    Err(_) => None,
  };

  let logger = PluginHostLogger { inner, host_tx };
  if let Err(err) = log::set_boxed_logger(Box::new(logger)) {
    eprintln!("Failed to install logger: {err}");
  }
  log::set_max_level(log::LevelFilter::Trace);
}

#[derive(Debug, Default)]
pub struct CsvPactPlugin {}

#[tonic::async_trait]
impl PactPlugin for CsvPactPlugin {

  // Returns the catalogue entries for CSV
  async fn init_plugin(
    &self,
    request: tonic::Request<proto::InitPluginRequest>,
  ) -> Result<tonic::Response<proto::InitPluginResponse>, tonic::Status> {
    let message = request.get_ref();
    PLUGIN_INSTANCE_ID.set(message.plugin_instance_id.clone()).ok();
    debug!("Init request from {}/{}", message.implementation, message.version);
    Ok(Response::new(proto::InitPluginResponse {
      response: Some(proto::init_plugin_response::Response::Success(
        proto::InitPluginSuccess {
          catalogue: vec![
            proto::CatalogueEntry {
              r#type: EntryType::ContentMatcher as i32,
              key: "csv".to_string(),
              values: hashmap! {
                "content-types".to_string() => "text/csv;application/csv".to_string()
              },
            },
            proto::CatalogueEntry {
              r#type: EntryType::ContentGenerator as i32,
              key: "csv".to_string(),
              values: hashmap! {
                "content-types".to_string() => "text/csv;application/csv".to_string()
              },
            },
          ],
          plugin_capabilities: vec![],
        },
      )),
    }))
  }

  // Not used
  async fn update_catalogue(
    &self,
    _request: tonic::Request<proto::Catalogue>,
  ) -> Result<tonic::Response<()>, tonic::Status> {
    debug!("Update catalogue request, ignoring");
    Ok(Response::new(()))
  }

  // Request to compare the CSV contents
  async fn compare_contents(
    &self,
    request: tonic::Request<proto::CompareContentsRequest>,
  ) -> Result<tonic::Response<proto::CompareContentsResponse>, tonic::Status> {
    let test_run_id = extract_test_run_id(request.get_ref().test_context.as_ref());
    TEST_RUN_ID.scope(test_run_id, async move {
      let request = request.get_ref();
      debug!("compare_contents request - {:?}", request);

      let has_headers = has_headers(&request.plugin_configuration);

      match (request.expected.as_ref(), request.actual.as_ref()) {
        (Some(expected), Some(actual)) => {
          let expected_csv_data = std::str::from_utf8(expected.content.as_ref().unwrap())
            .map_err(|err| tonic::Status::aborted(format!("Failed to compare CSV contents: {}", err)))?;
          let mut expected_rdr = ReaderBuilder::new().has_headers(has_headers)
            .from_reader(expected_csv_data.as_bytes());
          let actual_csv_data = actual.content.as_ref().unwrap();
          let mut actual_rdr = ReaderBuilder::new().has_headers(has_headers)
            .from_reader(actual_csv_data.as_slice());

          let rules = request.rules.iter()
            .map(|(key, rules)| {
              let rules = rules.rule.iter().fold(RuleList::empty(RuleLogic::And), |mut list, rule| {
                match to_object(&rule.values.as_ref().unwrap()) {
                  Value::Object(mut map) => {
                    map.insert("match".to_string(), Value::String(rule.r#type.clone()));
                    debug!("Creating matching rule with {:?}", map);
                    list.add_rule(&MatchingRule::from_json(&Value::Object(map)).unwrap());
                  }
                  _ => {}
                }
                list
              });
              (key.clone(), rules)
            }).collect();
          compare_contents(has_headers, &mut expected_rdr, &mut actual_rdr,
                           request.allow_unexpected_keys, rules)
            .map_err(|err| tonic::Status::aborted(format!("Failed to compare CSV contents: {}", err)))
        }
        (None, Some(actual)) => {
          let contents = actual.content.as_ref().unwrap();
          Ok(Response::new(proto::CompareContentsResponse {
            error: String::default(),
            type_mismatch: None,
            results: hashmap! {
              String::default() => proto::ContentMismatches {
                mismatches: vec![
                  proto::ContentMismatch {
                    expected: None,
                    actual: Some(contents.clone()),
                    mismatch: format!("Expected no CSV content, but got {} bytes", contents.len()),
                    path: "".to_string(),
                    diff: "".to_string(),
                    mismatch_type: String::default(),
                  }
                ]
              }
            }
          }))
        }
        (Some(expected), None) => {
          let contents = expected.content.as_ref().unwrap();
          Ok(Response::new(proto::CompareContentsResponse {
            error: String::default(),
            type_mismatch: None,
            results: hashmap! {
              String::default() => proto::ContentMismatches {
                mismatches: vec![
                  proto::ContentMismatch {
                    expected: Some(contents.clone()),
                    actual: None,
                    mismatch: format!("Expected CSV content, but did not get any"),
                    path: "".to_string(),
                    diff: "".to_string(),
                    mismatch_type: String::default(),
                  }
                ]
              }
            }
          }))
        }
        (None, None) => {
          Ok(Response::new(proto::CompareContentsResponse {
            error: String::default(),
            type_mismatch: None,
            results: hashmap!{}
          }))
        }
      }
    }).await
  }

  // Request to configure the interaction with CSV contents
  // Example definition we should receive:
  // "column:1", "matching(type,'Name')",
  // "column:2", "matching(number,100)",
  // "column:3", "matching(datetime, 'yyyy-MM-dd','2000-01-01')"
  async fn configure_interaction(
    &self,
    request: tonic::Request<proto::ConfigureInteractionRequest>,
  ) -> Result<tonic::Response<proto::ConfigureInteractionResponse>, tonic::Status> {
    let test_run_id = extract_test_run_id(request.get_ref().test_context.as_ref());
    TEST_RUN_ID.scope(test_run_id, async move {
      debug!("Received configure_contents request for '{}'", request.get_ref().content_type);
      setup_csv_contents(&request)
        .map_err(|err| tonic::Status::aborted(format!("Invalid column definition: {}", err)))
    }).await
  }

  // Request to generate CSV contents
  async fn generate_content(
    &self,
    request: tonic::Request<proto::GenerateContentRequest>,
  ) -> Result<tonic::Response<proto::GenerateContentResponse>, tonic::Status> {
    debug!("Received generate_content request");
    generate_csv_content(&request)
      .map(|contents| {
        debug!("Generated contents: {}", contents);
        Response::new(proto::GenerateContentResponse {
          contents: Some(proto::Body {
            content_type: contents.content_type().unwrap_or(ContentType::from("text/csv")).to_string(),
            content: Some(contents.value().unwrap().to_vec()),
            content_type_hint: ContentTypeHint::Default as i32
          })
        })
      })
      .map_err(|err| tonic::Status::aborted(format!("Failed to generate CSV contents: {}", err)))
  }
}

fn compare_contents<R: Read>(
  has_headers: bool,
  expected: &mut Reader<R>,
  actual: &mut Reader<R>,
  allow_unexpected_keys: bool,
  rules: HashMap<String, RuleList>
) -> anyhow::Result<tonic::Response<proto::CompareContentsResponse>> {
  debug!("Comparing contents using allow_unexpected_keys ({}) and rules ({:?})", allow_unexpected_keys, rules);

  let mut results = vec![];

  let expected_headers = if has_headers {
    match expected.headers() {
      Ok(headers) => headers.clone(),
      Err(err) => return Err(anyhow!("Failed to read the expected headers: {}", err)),
    }
  } else {
    StringRecord::default()
  };
  let actual_headers = if has_headers {
    match actual.headers() {
      Ok(headers) => headers.clone(),
      Err(err) => return Err(anyhow!("Failed to read the actual headers: {}", err)),
    }
  } else {
    StringRecord::default()
  };
  let actual_headers: HashMap<&str, usize> = actual_headers
    .iter()
    .enumerate()
    .map(|(col, hdr)| (hdr, col))
    .collect();

  if has_headers {
    for header in expected_headers.iter() {
      if !actual_headers.contains_key(header) {
        results.push(proto::ContentMismatch {
          expected: Some(header.as_bytes().to_vec()),
          actual: None,
          mismatch: format!("Expected columns '{}', but was missing", header),
          path: String::default(),
          diff: String::default(),
          mismatch_type: String::default(),
        });
      }
    }
  }

  let mut expected_records = expected.records();
  let mut actual_records = actual.records();

  let expected_row = expected_records.next()
    .ok_or_else(|| anyhow!("Could not read the expected content"))??;
  let actual_row = actual_records.next()
    .ok_or_else(|| anyhow!("Could not read the actual content"))??;

  if !has_headers {
    if actual_row.len() < expected_row.len() {
      results.push(proto::ContentMismatch {
        expected: Some(format!("{} columns", expected_row.len()).as_bytes().to_vec()),
        actual: Some(format!("{} columns", actual_row.len()).as_bytes().to_vec()),
        mismatch: format!("Expected {} columns, but got {}", expected_row.len(), actual_row.len()),
        path: String::default(),
        diff: String::default(),
        mismatch_type: String::default(),
      });
    } else if actual_row.len() > expected_row.len() && !allow_unexpected_keys {
      results.push(proto::ContentMismatch {
        expected: Some(format!("{} columns", expected_row.len()).as_bytes().to_vec()),
        actual: Some(format!("{} columns", actual_row.len()).as_bytes().to_vec()),
        mismatch: format!("Expected at least {} columns, but got {}", expected_row.len(), actual_row.len()),
        path: String::default(),
        diff: String::default(),
        mismatch_type: String::default(),
      });
    }
  }

  compare_row(&expected_row, &actual_row, &rules, has_headers, &expected_headers, &actual_headers, &mut results);
  for row in actual_records {
    compare_row(&expected_row, &row?, &rules, has_headers, &expected_headers, &actual_headers, &mut results);
  }

  Ok(Response::new(proto::CompareContentsResponse {
    error: String::default(),
    type_mismatch: None,
    results: hashmap! {
      String::default() => proto::ContentMismatches {
        mismatches: results
      }
    }
  }))
}

fn compare_row(
  expected_row: &StringRecord,
  actual_row: &StringRecord,
  rules: &HashMap<String, RuleList>,
  has_headers: bool,
  expected_headers: &StringRecord,
  actual_headers: &HashMap<&str, usize>,
  results: &mut Vec<proto::ContentMismatch>) {
  for (index, expected_item) in expected_row.iter().enumerate() {
    let header = expected_headers.get(index).unwrap_or_default();
    let item = if has_headers {
      match actual_headers.get(header) {
        Some(actual_index) => actual_row.get(*actual_index).unwrap_or_default(),
        None => ""
      }
    } else {
      actual_row.get(index).unwrap_or_default()
    };

    let path = format!("column:{}", index + 1);
    let header_path = format!("column:{}", header);

    if let Some(rules) = rules.get(&path).or_else(|| rules.get(header_path.as_str())) {
      for rule in &rules.rules {
        if let Err(err) = rule.match_value(expected_item, item, false, false) {
          results.push(proto::ContentMismatch {
            expected: Some(expected_item.as_bytes().to_vec()),
            actual: Some(item.as_bytes().to_vec()),
            mismatch: err.to_string(),
            path: format!(
              "row:{:5}, column:{:2}",
              actual_row.position().unwrap().line(),
              index
            ),
            diff: String::default(),
            mismatch_type: String::default()
          });
        }
      }
    } else if item != expected_item {
      results.push(proto::ContentMismatch {
        expected: Some(expected_item.as_bytes().to_vec()),
        actual: Some(item.as_bytes().to_vec()),
        mismatch: format!(
          "Expected column {} value to equal '{}', but got '{}'",
          index, expected_item, item
        ),
        path: format!(
          "row:{:5}, column:{:2}",
          actual_row.position().unwrap().line(),
          index
        ),
        diff: String::default(),
        mismatch_type: String::default(),
      });
    }
  }
}

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

  #[test]
  fn compare_contents_handles_single_row_without_headers() {
    let mut expected = ReaderBuilder::new()
      .has_headers(false)
      .from_reader("Name,100,2000-01-01\n".as_bytes());
    let mut actual = ReaderBuilder::new()
      .has_headers(false)
      .from_reader("Alice,123,2001-02-03\n".as_bytes());

    let response = compare_contents(false, &mut expected, &mut actual, false, HashMap::new())
      .expect("compare_contents should succeed");

    assert!(response.get_ref().error.is_empty());
  }

  #[tokio::test]
  async fn init_plugin_returns_csv_catalogue_entries() {
    let plugin = CsvPactPlugin::default();

    let response = plugin
      .init_plugin(Request::new(proto::InitPluginRequest {
        implementation: "plugin-driver-rust".to_string(),
        version: "1.0.0-beta.1".to_string(),
        host_capabilities: vec![],
        ..Default::default()
      }))
      .await
      .unwrap()
      .into_inner();

    match response.response.unwrap() {
      proto::init_plugin_response::Response::Success(success) => {
        assert_eq!(success.catalogue.len(), 2);
        assert!(success.plugin_capabilities.is_empty());
      }
      _ => panic!("expected init success"),
    }
  }
}

struct TcpIncoming {
  inner: TcpListener
}

impl Stream for TcpIncoming {
  type Item = Result<TcpStream, std::io::Error>;

  fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
    Pin::new(&mut self.inner).poll_accept(cx)
      .map_ok(|(stream, _)| stream).map(|v| Some(v))
  }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
  init_logging().await;

  let addr: SocketAddr = "0.0.0.0:0".parse()?;
  let listener = TcpListener::bind(addr).await?;
  let address = listener.local_addr()?;

  let server_key = Uuid::new_v4().to_string();
  println!("{{\"port\":{}, \"serverKey\":\"{}\"}}", address.port(), server_key);
  let _ = io::stdout().flush();

  let plugin = CsvPactPlugin::default();
  Server::builder()
    .add_service(PactPluginServer::new(plugin))
    .serve_with_incoming(TcpIncoming { inner: listener }).await?;

  Ok(())
}