temporary(datasecurity): fill all fields

Author: aimenebelfodil

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/build.rs

@@ -15,7 +15,24 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     println!("cargo:rerun-if-changed={proto_file}");
 
-    prost_build::compile_protos(&[proto_file], &[proto_root])?;
+    let mut config = prost_build::Config::new();
+
+    // TODO(DSEC): remove this serde derive once debugging ends — it only exists to
+    // emit the SDS result payload as JSON (see `send_sds_result_json`), which we do
+    // not need in production. `google.protobuf.Timestamp` (prost-types) has no serde
+    // support, so its fields use a local `serialize_with` helper.
+    config.type_attribute(".", "#[derive(serde::Serialize)]");
+    for field in [
+        ".datadog.sds.SdsResultPayload.ScanMetadata.ScanTaskMetadata.started_at",
+        ".datadog.sds.SdsResultPayload.ScanMetadata.ScanTaskMetadata.ended_at",
+    ] {
+        config.field_attribute(
+            field,
+            "#[serde(serialize_with = \"crate::proto::serialize_timestamp\")]",
+        );
+    }
+
+    config.compile_protos(&[proto_file], &[proto_root])?;
 
     Ok(())
 }

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/src/check.rs

@@ -1,9 +1,14 @@
+use std::time::{SystemTime, UNIX_EPOCH};
+
 use anyhow::{Context, Result, anyhow};
-use core::*;
+use shlib_core::*;
+use serde_json::Value;
 
 use crate::backend;
 use crate::config::{CheckConfig, SubTask};
-use crate::payload::{Match, ScanEventPayload, ScanStatus};
+use crate::constants::SDS_RESULT_EVENT_TYPE;
+use crate::payload::{Match, ScanStatus};
+use crate::proto::{self, ScanMetadata, ScanResult, ScanTaskMetadata, SdsResultPayload, Status};
 use crate::scanning::Scanner;
 
 /// Check entrypoint.
@@ -45,56 +50,195 @@ fn run_sub_task(
         sub_task.sub_task_id, sub_task.entity.platform
     );
 
+    // Time the scan so the payload can carry started_at / ended_at / duration.
+    let started_at = SystemTime::now();
+    let scan = run_scan(scanner, sub_task);
+    let ended_at = SystemTime::now();
+
     // A sub task failure is reported inside the payload (status=ERROR) rather
     // than aborting the check, so every sub task produces exactly one event.
-    let (status, failure_reason, matches) = match run_scan(scanner, sub_task) {
-        Ok(matches) => {
+    let (status, failure_reason, matches, scanned_row_count) = match scan {
+        Ok(out) => {
             println!(
                 "datasecurity: sub task succeeded ({} match(es))",
                 matches.len()
             );
-            (ScanStatus::Success, String::new(), matches)
+            (ScanStatus::Success, String::new(), out.matches, out.scanned_row_count)
         }
         Err(err) => {
             let reason = format!("{err:#}");
             eprintln!(
                 "datasecurity: sub task {} failed: {reason}",
                 sub_task.sub_task_id
             );
-            (ScanStatus::Error, reason, Vec::new())
+            (ScanStatus::Error, reason, Vec::new(), 0)
         }
     };
 
-    let payload = ScanEventPayload {
-        task_id: config.task_id.clone(),
-        sub_task_id: sub_task.sub_task_id.clone(),
+    // Build the SDS result protobuf: task metadata, timing, postgres location and
+    // matches, mirroring the Data Observability crawler payload.
+    let payload = build_sds_result(
+        config,
+        sub_task,
         status,
-        failure_reason,
-        matches,
-    };
+        &failure_reason,
+        &matches,
+        scanned_row_count,
+        started_at,
+        ended_at,
+    );
+
+    // Emit the protobuf on the `sds-result` event platform track.
+    if config.send_sds_result {
+        check.event_platform_event_bytes(&proto::encode(&payload), SDS_RESULT_EVENT_TYPE)?;
+    }
 
-    // TODO(DSEC-140): send sdsresult rather than an event
-    let payload_json =
-        serde_json::to_string(&payload).context("failed to serialize scan event payload")?;
-    check.event(
-        "datasecurity scan result",
-        &payload_json,
-        0,
-        "normal",
-        "",
-        &[],
-        "info",
-        "",
-        "datasecurity",
-        "",
-    )?;
+    // TODO(DSEC): remove this JSON event once the protobuf is validated end to
+    // end — we do not need to send SDS results as JSON. It is serialized from the
+    // same protobuf so the two representations cannot drift.
+    if config.send_sds_result_json {
+        let payload_json =
+            proto::to_json(&payload).context("failed to serialize sds result payload to json")?;
+        check.event(
+            "datasecurity scan result",
+            &payload_json,
+            0,
+            "normal",
+            "",
+            &[],
+            "info",
+            "",
+            "datasecurity",
+            "",
+        )?;
+    }
 
     Ok(())
 }
 
-/// Fetches the sub task's data and scans it, returning the matches.
-/// TODO(dsec-161): add tests for the scan.
-fn run_scan(scanner: &Scanner, sub_task: &SubTask) -> Result<Vec<Match>> {
+/// Builds the `SdsResultPayload` protobuf for one sub task.
+///
+/// Mirrors the Data Observability crawler payload (`Resource`, `RuleIds`,
+/// `ScanningSource`, `ScanResults`), swapping the snowflake location for a
+/// postgres one and adding the scan-task metadata block.
+#[allow(clippy::too_many_arguments)]
+fn build_sds_result(
+    config: &CheckConfig,
+    sub_task: &SubTask,
+    status: ScanStatus,
+    failure_reason: &str,
+    matches: &[Match],
+    scanned_row_count: i64,
+    started_at: SystemTime,
+    ended_at: SystemTime,
+) -> SdsResultPayload {
+    let entity = &sub_task.entity;
+    let duration_ms = ended_at
+        .duration_since(started_at)
+        .map(|d| d.as_millis() as i64)
+        .unwrap_or(0);
+
+    let location = proto::ScanLocation {
+        scan_location: Some(proto::scan_location::ScanLocation::PostgresTable(
+            proto::PostgresTable {
+                database_cluster_name: entity.database_cluster_name.clone(),
+                database_instance_name: entity.database_instance_name.clone(),
+                database_host_name: sub_task.connection.host.clone(),
+                database_name: entity.database.clone(),
+                schema_name: entity.schema.clone(),
+                table_name: entity.table.clone(),
+                scanned_row_count,
+                // TODO(DSEC): populate table_row_count (from DBM metadata) and
+                // scanned_columns.
+                ..Default::default()
+            },
+        )),
+        ..Default::default()
+    };
+
+    let scan_result = ScanResult {
+        table_matches: proto::table_matches(matches),
+        location: Some(location),
+        duration: duration_ms,
+        scan_metadata: Some(ScanMetadata {
+            scan_task_metadata: Some(ScanTaskMetadata {
+                task_id: config.task_id.clone(),
+                sub_task_id: sub_task.sub_task_id.clone(),
+                started_at: Some(proto::to_timestamp(started_at)),
+                ended_at: Some(proto::to_timestamp(ended_at)),
+                status: match status {
+                    ScanStatus::Success => Status::Success,
+                    ScanStatus::Error => Status::Error,
+                } as i32,
+                failure_reason: (!failure_reason.is_empty()).then(|| failure_reason.to_string()),
+            }),
+        }),
+        ..Default::default()
+    };
+
+    SdsResultPayload {
+        timestamp: now_unix_millis(),
+        resource: Some(proto::Resource {
+            r#type: "postgres_table".to_string(),
+            name: resource_name(sub_task),
+        }),
+        rule_ids: config.scanning_rules.iter().map(|rule| rule.id.clone()).collect(),
+        // The scanning source is the Agent. TODO(DSEC): populate hostname and
+        // agent version once the check receives them (not provided via config yet).
+        scanning_source: Some(proto::ScanningSource {
+            source: Some(proto::scanning_source::Source::Agent(
+                proto::scanning_source::Agent::default(),
+            )),
+        }),
+        scan_results: vec![scan_result],
+        ..Default::default()
+    }
+}
+
+/// Result of a successful sub task scan.
+struct ScanOutput {
+    matches: Vec<Match>,
+    scanned_row_count: i64,
+}
+
+/// Fetches the sub task's data and scans it, returning the matches and the
+/// number of rows scanned.
+fn run_scan(scanner: &Scanner, sub_task: &SubTask) -> Result<ScanOutput> {
     let data = backend::fetch_data(sub_task).context("fetching sub task data")?;
-    scanner.scan(data).context("scanning sub task data")
+    let matches = scanner.scan(data.clone()).context("scanning sub task data")?;
+    Ok(ScanOutput {
+        scanned_row_count: scanned_rows(&data),
+        matches,
+    })
+}
+
+/// Number of rows scanned: the longest column array in the `{ column: [values] }`
+/// map returned by the backend.
+fn scanned_rows(data: &Value) -> i64 {
+    data.as_object()
+        .and_then(|columns| {
+            columns
+                .values()
+                .filter_map(|value| value.as_array().map(Vec::len))
+                .max()
+        })
+        .unwrap_or(0) as i64
+}
+
+/// Current unix time in milliseconds, for the payload timestamp.
+fn now_unix_millis() -> i64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|d| d.as_millis() as i64)
+        .unwrap_or(0)
+}
+
+/// Resource name (`<instance_name>.<database>.<schema>.<table>`), following the
+/// DO crawler convention.
+fn resource_name(sub_task: &SubTask) -> String {
+    let entity = &sub_task.entity;
+    format!(
+        "{}.{}.{}.{}",
+        entity.database_instance_name, entity.database, entity.schema, entity.table
+    )
 }

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/src/config.rs

@@ -1,5 +1,5 @@
 use anyhow::{Context, Result};
-use core::AgentCheck;
+use shlib_core::AgentCheck;
 use serde::Deserialize;
 
 use crate::scanning::ScanningRule;
@@ -19,6 +19,11 @@ impl CheckConfig {
                 .instance
                 .get("scan_data")
                 .context("failed to read scan_data from instance config")?,
+            // Emit the SDS result protobuf on the `sds-result` track (default on).
+            send_sds_result: check.instance.get("send_sds_result").unwrap_or(true),
+            // TODO(DSEC): debug only, off by default. Remove once the protobuf is
+            // validated end to end.
+            send_sds_result_json: check.instance.get("send_sds_result_json").unwrap_or(false),
         })
     }
 }
@@ -29,6 +34,16 @@ pub struct CheckConfig {
     pub task_id: String,
     pub scanning_rules: Vec<ScanningRule>,
     pub scan_data: Vec<SubTask>,
+    /// Emit the SDS result protobuf on the `sds-result` event platform track.
+    #[serde(default)]
+    pub send_sds_result: bool,
+    /// Emit the SDS result payload as JSON (serialized from the same protobuf) as
+    /// a regular event.
+    ///
+    /// TODO(DSEC): debug only, remove once the protobuf is validated end to end —
+    /// we do not need to send SDS results as JSON.
+    #[serde(default)]
+    pub send_sds_result_json: bool,
 }
 
 /// A single scan sub task: a query to run against one data source.
@@ -47,13 +62,28 @@ pub struct SubTask {
     pub timeout_seconds: u64,
 }
 
-/// TODO(dsec-140): add the other entity values (scan location) when needed,
-/// e.g. database_cluster_name, database_instance_name, database, schema, table.
+/// The data asset a sub task targets, mirroring the Data Observability `entity`
+/// object (`comp/dataobs/queryactions`). Used to select the backend engine
+/// (`platform`) and to describe the scan location; not used to connect.
 #[derive(Debug, Default, Deserialize)]
 pub struct Entity {
     /// Data-source platform, used to select the backend engine (e.g. `postgres`).
     #[serde(default)]
     pub platform: String,
+    /// Cloud cluster identifier of the data source (e.g. an RDS/Aurora cluster).
+    /// Reported in the scan location.
+    #[serde(default)]
+    pub database_cluster_name: String,
+    /// Cloud instance identifier of the data source; the DO entity `account`.
+    /// Reported in the scan location.
+    #[serde(default, alias = "account")]
+    pub database_instance_name: String,
+    #[serde(default)]
+    pub database: String,
+    #[serde(default)]
+    pub schema: String,
+    #[serde(default)]
+    pub table: String,
 }
 
 fn default_timeout_seconds() -> u64 {

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/src/lib.rs

@@ -1,4 +1,4 @@
-use core::generate_ffi;
+use shlib_core::generate_ffi;
 
 mod backend;
 mod check;

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/src/payload.rs

@@ -1,29 +1,12 @@
-use serde::Serialize;
-
-/// JSON event payload emitted by the check, one per sub task.
-///
-/// TODO(dsec-140): replace with the SDS result payload (protoc generated code).
-#[derive(Debug, Serialize)]
-#[serde(rename_all = "snake_case")]
-pub struct ScanEventPayload {
-    pub task_id: String,
-    pub sub_task_id: String,
-    pub status: ScanStatus,
-    #[serde(skip_serializing_if = "String::is_empty")]
-    pub failure_reason: String,
-    pub matches: Vec<Match>,
-}
-
-/// TODO(dsec-140): replace with the SDS result payload task status (protoc generated code).
-#[derive(Debug, Clone, Copy, Serialize)]
-#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
+/// Outcome of a sub task scan, mirroring the proto `ScanStatus`.
+#[derive(Debug, Clone, Copy)]
 pub enum ScanStatus {
     Success,
     Error,
 }
 
-#[derive(Debug, Serialize, PartialEq)]
-#[serde(rename_all = "snake_case")]
+/// Aggregated scanner match for one column, converted into a proto `TableMatch`.
+#[derive(Debug, PartialEq)]
 pub struct Match {
     pub rule_id: String,
     pub column_name: String,