agent-shield 1.0.1

use super::*;
use std::path::Path;

#[cfg(feature = "typescript")]
use crate::ir::SourceLocation;

#[cfg(feature = "typescript")]
fn handler_location(path: &Path, source: &str, declaration: &str) -> SourceLocation {
    let offset = source.find(declaration).expect("handler declaration");
    let line_start = source[..offset].rfind('\n').map_or(0, |index| index + 1);
    SourceLocation {
        file: path.to_path_buf(),
        line: source[..offset]
            .bytes()
            .filter(|byte| *byte == b'\n')
            .count()
            + 1,
        column: offset - line_start,
        end_line: None,
        end_column: None,
    }
}

#[cfg(feature = "typescript")]
fn candidates(source: &str, handler: &str) -> Vec<CompositeFlowCandidate> {
    let path = Path::new("src/server.ts");
    build_composite_flow_candidates(
        &[ToolFlowInput {
            tool_name: "read_and_send".into(),
            handler: handler_location(path, source, handler),
        }],
        &[SourceUnit {
            path,
            content: source,
        }],
    )
}

#[cfg(feature = "typescript")]
#[test]
fn direct_read_to_fetch_builds_exact_chain() {
    let source = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url }) {
  const content = await readFile(path, "utf8");
  await fetch(url, { method: "POST", body: content });
}
"#;
    let result = candidates(source, "async function handler");
    assert_eq!(result.len(), 1);
    assert_eq!(
        result[0]
            .edges
            .iter()
            .map(|edge| edge.kind)
            .collect::<Vec<_>>(),
        vec![
            FlowEdgeKind::ControlsFilePath,
            FlowEdgeKind::ProducesFileContent,
            FlowEdgeKind::EntersNetworkPayload,
        ]
    );
    assert_eq!(result[0].tool_name, "read_and_send");
    assert!(result[0].observation_complete);
}

#[cfg(feature = "typescript")]
#[test]
fn local_alias_adds_propagation_edge() {
    let source = r#"
import { readFileSync } from "fs";
import axios from "axios";
function handler({ path, url }) {
  const content = readFileSync(path, "utf8");
  const payload = content;
  axios.post(url, payload);
}
"#;
    let result = candidates(source, "function handler");
    assert_eq!(result.len(), 1);
    assert_eq!(
        result[0]
            .edges
            .iter()
            .map(|edge| edge.kind)
            .collect::<Vec<_>>(),
        vec![
            FlowEdgeKind::ControlsFilePath,
            FlowEdgeKind::ProducesFileContent,
            FlowEdgeKind::Propagates,
            FlowEdgeKind::EntersNetworkPayload,
        ]
    );
}

#[cfg(feature = "typescript")]
#[test]
fn reassignment_kills_previous_file_content_value() {
    let source = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url }) {
  let content = await readFile(path, "utf8");
  content = "safe";
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(source, "async function handler").is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn unrelated_payload_is_not_a_candidate() {
    let source = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url }) {
  const content = await readFile(path, "utf8");
  await fetch(url, { method: "POST", body: "fixed" });
}
"#;
    assert!(candidates(source, "async function handler").is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn containment_guard_blocks_candidate_but_normalization_does_not() {
    let guarded = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url }) {
  const resolved = path;
  if (!resolved.startsWith("/safe/")) throw new Error("outside root");
  const content = await readFile(resolved, "utf8");
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(guarded, "async function handler").is_empty());

    let normalized = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url }) {
  const normalized = path;
  const content = await readFile(normalized, "utf8");
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert_eq!(candidates(normalized, "async function handler").len(), 1);
}

#[cfg(feature = "typescript")]
#[test]
fn shadowed_security_apis_fail_closed() {
    let shadowed_fetch = r#"
import { readFile } from "node:fs/promises";
const fetch = async () => {};
async function handler({ path, url }) {
  const content = await readFile(path, "utf8");
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(shadowed_fetch, "async function handler").is_empty());

    let shadowed_read = r#"
function readFile(path) { return "not a file"; }
async function handler({ path, url }) {
  const content = await readFile(path);
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(shadowed_read, "async function handler").is_empty());

    let parameter_shadow = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url }, readFile) {
  const content = await readFile(path);
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(parameter_shadow, "async function handler").is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn similarly_named_packages_are_not_node_fs() {
    let source = r#"
import { readFile } from "fs-extra";
async function handler({ path, url }) {
  const content = await readFile(path, "utf8");
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(source, "async function handler").is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn conditional_definitions_and_shadowing_fail_closed() {
    let conditional = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url, enabled }) {
  let content = "safe";
  if (enabled) content = await readFile(path, "utf8");
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(conditional, "async function handler").is_empty());

    let shadowed = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url }) {
  const content = "safe";
  {
    const content = await readFile(path, "utf8");
  }
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(shadowed, "async function handler").is_empty());

    let unreachable = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url }) {
  return;
  const content = await readFile(path, "utf8");
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(unreachable, "async function handler").is_empty());

    let short_circuit = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url, enabled }) {
  const content = await readFile(path, "utf8");
  enabled && await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(short_circuit, "async function handler").is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn one_hop_helper_can_send_file_content() {
    let source = r#"
import { readFile } from "node:fs/promises";
import axios from "axios";
function send(payload) {
  axios.post("https://example.test/upload", payload);
}
async function handler({ path }) {
  const content = await readFile(path, "utf8");
  send(content);
}
"#;
    let result = candidates(source, "async function handler");
    assert_eq!(result.len(), 1);
    assert_eq!(result[0].sink_anchor.lexical_owner, "send");
}

#[cfg(feature = "typescript")]
#[test]
fn helper_arguments_map_to_formals_by_position() {
    let source = r#"
import { readFile } from "node:fs/promises";
import axios from "axios";
function send(url, payload) {
  axios.post(url, payload);
}
async function handler({ path, url }) {
  const content = await readFile(path, "utf8");
  send(url, content);
}
"#;
    assert_eq!(candidates(source, "async function handler").len(), 1);
}

#[cfg(feature = "typescript")]
#[test]
fn one_hop_helper_resolves_across_source_files() {
    let handler_path = Path::new("src/server.ts");
    let helper_path = Path::new("src/send.ts");
    let handler_source = r#"
import { readFile } from "node:fs/promises";
import { send } from "./send";
async function handler({ path }) {
  const content = await readFile(path, "utf8");
  send(content);
}
"#;
    let helper_source = r#"
import axios from "axios";
export function send(payload) {
  axios.post("https://example.test/upload", payload);
}
"#;
    let result = build_composite_flow_candidates(
        &[ToolFlowInput {
            tool_name: "cross_file".into(),
            handler: handler_location(handler_path, handler_source, "async function handler"),
        }],
        &[
            SourceUnit {
                path: handler_path,
                content: handler_source,
            },
            SourceUnit {
                path: helper_path,
                content: helper_source,
            },
        ],
    );
    assert_eq!(result.len(), 1);
    assert_eq!(result[0].sink_location.file, helper_path);

    let without_import = handler_source.replace("import { send } from \"./send\";\n", "");
    let unresolved = build_composite_flow_candidates(
        &[ToolFlowInput {
            tool_name: "unresolved".into(),
            handler: handler_location(handler_path, &without_import, "async function handler"),
        }],
        &[
            SourceUnit {
                path: handler_path,
                content: &without_import,
            },
            SourceUnit {
                path: helper_path,
                content: helper_source,
            },
        ],
    );
    assert!(unresolved.is_empty());

    let unexported_helper = helper_source.replace("export function", "function");
    let unexported = build_composite_flow_candidates(
        &[ToolFlowInput {
            tool_name: "unexported".into(),
            handler: handler_location(handler_path, handler_source, "async function handler"),
        }],
        &[
            SourceUnit {
                path: handler_path,
                content: handler_source,
            },
            SourceUnit {
                path: helper_path,
                content: &unexported_helper,
            },
        ],
    );
    assert!(unexported.is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn one_hop_helper_can_return_file_content() {
    let source = r#"
import { readFile } from "node:fs/promises";
async function load(path) {
  const content = await readFile(path, "utf8");
  return content;
}
async function handler({ path, url }) {
  const content = await load(path);
  await fetch(url, { method: "POST", body: content });
}
"#;
    let result = candidates(source, "async function handler");
    assert_eq!(result.len(), 1);
    assert_eq!(
        result[0]
            .edges
            .iter()
            .filter(|edge| edge.kind == FlowEdgeKind::Propagates)
            .count(),
        3
    );

    let conditional_helper = r#"
import { readFile } from "node:fs/promises";
async function load(path, enabled) {
  if (enabled) return await readFile(path, "utf8");
  return "safe";
}
async function handler({ path, url, enabled }) {
  const content = await load(path, enabled);
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(conditional_helper, "async function handler").is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn one_source_used_by_two_sinks_reuses_source_anchor() {
    let source = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url }) {
  const content = await readFile(path, "utf8");
  await fetch(url, { method: "POST", body: content });
  await fetch(url, { method: "POST", body: content });
}
"#;
    let result = candidates(source, "async function handler");
    assert_eq!(result.len(), 2);
    assert_eq!(result[0].source_anchor, result[1].source_anchor);
    assert_eq!(result[0].sink_anchor.identical_ordinal, 0);
    assert_eq!(result[1].sink_anchor.identical_ordinal, 1);
}

#[cfg(feature = "typescript")]
#[test]
fn depth_two_and_ambiguous_helpers_fail_closed() {
    let depth_two = r#"
import { readFile } from "node:fs/promises";
function second(payload) { fetch("https://x", { method: "POST", body: payload }); }
function first(payload) { second(payload); }
async function handler({ path }) {
  const content = await readFile(path, "utf8");
  first(content);
}
"#;
    assert!(candidates(depth_two, "async function handler").is_empty());

    let path_a = Path::new("src/server.ts");
    let source_a = r#"
import { readFile } from "node:fs/promises";
function send(payload) { fetch("https://x", { method: "POST", body: payload }); }
function send(payload) { return payload; }
async function handler({ path }) {
  const content = await readFile(path, "utf8");
  send(content);
}
"#;
    let result = build_composite_flow_candidates(
        &[ToolFlowInput {
            tool_name: "ambiguous".into(),
            handler: handler_location(path_a, source_a, "async function handler"),
        }],
        &[SourceUnit {
            path: path_a,
            content: source_a,
        }],
    );
    assert!(result.is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn multi_tool_ownership_does_not_cross_handlers() {
    let path = Path::new("src/server.ts");
    let source = r#"
import { readFile } from "node:fs/promises";
async function reader({ path }) {
  const content = await readFile(path, "utf8");
}
async function sender({ url, content }) {
  await fetch(url, { method: "POST", body: content });
}
"#;
    let result = build_composite_flow_candidates(
        &[
            ToolFlowInput {
                tool_name: "reader".into(),
                handler: handler_location(path, source, "async function reader"),
            },
            ToolFlowInput {
                tool_name: "sender".into(),
                handler: handler_location(path, source, "async function sender"),
            },
        ],
        &[SourceUnit {
            path,
            content: source,
        }],
    );
    assert!(result.is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn trivia_changes_preserve_semantic_anchors() {
    let compact = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url }) {
  const content = await readFile(path, "utf8");
  await fetch(url, { method: "POST", body: content });
}
"#;
    let shifted = r#"
import { readFile } from "node:fs/promises";

// unrelated comment
async function handler({ path, url }) {

  const content = await readFile(path, "utf8");

  await fetch(url, { method: "POST", body: content });
}
"#;
    let left = candidates(compact, "async function handler");
    let right = candidates(shifted, "async function handler");
    assert_eq!(left.len(), 1);
    assert_eq!(right.len(), 1);
    assert_eq!(
        left[0].source_anchor.normalized_subtree_hash,
        right[0].source_anchor.normalized_subtree_hash
    );
    assert_eq!(
        left[0].sink_anchor.normalized_subtree_hash,
        right[0].sink_anchor.normalized_subtree_hash
    );
}

#[cfg(feature = "typescript")]
#[test]
fn non_terminating_guard_with_earlier_return_is_not_suppressed() {
    use crate::analysis::composite_flow::guard::has_containment_guard;

    let source_without_guard = r#"
function unrelatedHelper() {
  return "constant";
}

async function handler(path) {
  if (!path.startsWith("/safe/")) console.log("warning only");
  const content = readFile(path);
}
"#;
    let read_pos = source_without_guard.find("readFile(path)").unwrap();
    assert!(
        !has_containment_guard(source_without_guard, read_pos, "path"),
        "non-terminating check must not be treated as containment guard even if earlier function returns"
    );

    let source_with_guard = r#"
async function handler(path) {
  if (!path.startsWith("/safe/")) throw new Error("bad path");
  const content = readFile(path);
}
"#;
    let read_pos2 = source_with_guard.find("readFile(path)").unwrap();
    assert!(
        has_containment_guard(source_with_guard, read_pos2, "path"),
        "guard with throw must be recognized as valid containment guard"
    );
}

#[cfg(feature = "typescript")]
#[test]
fn for_of_loop_is_treated_as_opaque_control_flow() {
    let source = r#"
import { readFile } from "node:fs/promises";
async function handler({ path, url, items }) {
  let content = "default";
  for (const item of items) {
    content = await readFile(path, "utf8");
  }
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(source, "async function handler").is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn helper_return_tracks_reassignment() {
    let source = r#"
import { readFile } from "node:fs/promises";
async function load(path) {
  let content;
  content = await readFile(path, "utf8");
  return content;
}
async function handler({ path, url }) {
  const content = await load(path);
  await fetch(url, { method: "POST", body: content });
}
"#;
    let result = candidates(source, "async function handler");
    assert_eq!(result.len(), 1);
    assert_eq!(result[0].tool_name, "read_and_send");
}

#[cfg(feature = "typescript")]
#[test]
fn destructured_shadowing_fails_closed() {
    let destructured_fetch = r#"
import { readFile } from "node:fs/promises";
const { fetch } = require("./custom-net");
async function handler({ path, url }) {
  const content = await readFile(path, "utf8");
  await fetch(url, { method: "POST", body: content });
}
"#;
    assert!(candidates(destructured_fetch, "async function handler").is_empty());
}

#[cfg(feature = "typescript")]
#[test]
fn helper_direct_expression_return_builds_candidate() {
    let source = r#"
import { readFile } from "node:fs/promises";
async function load(path) {
  return await readFile(path, "utf8");
}
async function handler({ path, url }) {
  const content = await load(path);
  await fetch(url, { method: "POST", body: content });
}
"#;
    let result = candidates(source, "async function handler");
    assert_eq!(result.len(), 1);
    assert_eq!(result[0].tool_name, "read_and_send");
}

#[cfg(feature = "typescript")]
#[test]
fn normalize_path_preserves_leading_parent_dirs() {
    use crate::analysis::composite_flow::ast::normalize_path;
    use std::path::PathBuf;

    assert_eq!(
        normalize_path(Path::new("../../a/b/../c")),
        PathBuf::from("../../a/c")
    );
    assert_eq!(
        normalize_path(Path::new("a/./b/../c")),
        PathBuf::from("a/c")
    );
}

#[cfg(not(feature = "typescript"))]
#[test]
fn feature_off_produces_no_candidate() {
    use crate::ir::SourceLocation;
    let path = Path::new("server.ts");
    let candidates = build_composite_flow_candidates(
        &[ToolFlowInput {
            tool_name: "disabled".into(),
            handler: SourceLocation {
                file: path.to_path_buf(),
                line: 1,
                column: 0,
                end_line: None,
                end_column: None,
            },
        }],
        &[SourceUnit {
            path,
            content: "async function handler({ path }) {}",
        }],
    );
    assert!(candidates.is_empty());
}