santh-writ 0.1.1

//! Per-sink-class concrete witness construction.
//!
//! Some `SinkConstraintKind` values have a witness shape that is
//! mechanically derivable from path-bound variables and does not
//! require an SMT solver:
//!
//! - **strcpy / unbounded string copy**: any byte sequence of length
//!   `> sizeof(dst)` overflows. The witness is `dst_size + 1` bytes.
//! - **gets**: any input of length `> sizeof(dst)`. Same witness as
//!   strcpy.
//! - **sprintf**: a formatted argument that, after %-substitution,
//!   exceeds `sizeof(dst)`. Witness depends on the format string.
//!
//! Other sink kinds (`heap_overflow_alloc_arith_wraps`, command
//! injection, SQL injection) need real SMT  -  those route through
//! `z3_backend` once per-language statement encoders land.
//!
//! This module ships the trivial encoders today so the simplest
//! Class 1 shapes have working witnesses at launch. The structure
//! mirrors the eventual SMT-backed encoders so swapping in real
//! solving is additive, not replacement.

use crate::{ConcreteInput, EntrypointKind, ShapeParams, SinkConstraintKind, WitnessRequest};

/// Bytes appended past `dst_capacity` to clobber the saved frame pointer +
/// return address on x86-64 SysV. Single owner (ONE-PLACE) so the byte length
/// and every render string that reports it can never drift apart.
const CANARY_PADDING: usize = 16;

/// Upper bound on a synthesized stack-overflow witness. A real stack-resident
/// destination is at most a few MiB (the default stack rlimit is 8 MiB); a
/// `dst_capacity` beyond this is a corrupt/implausible shape parameter, not a
/// witnessable buffer. We refuse rather than allocate a multi-gigabyte `Vec`
/// (OOM) or, near `usize::MAX`, wrap the length and emit a silently-short
/// payload that would not actually overflow.
const MAX_WITNESS_BYTES: usize = 64 * 1024 * 1024;

/// Attempt a trivial witness construction without invoking the SMT
/// solver. Returns `Some(...)` for shapes where the witness is a
/// closed-form function of statically-known sink parameters; returns
/// `None` for shapes that require constraint solving  -  those route
/// to `Z3Backend`.
#[must_use]
pub fn try_trivial_witness(request: &WitnessRequest) -> Option<ConcreteInput> {
    match request.sink_kind {
        // Memory-corruption shapes (Class 1 launch surface). Each
        // takes per-shape parameters carried in `shape_params`; if
        // the parameters are missing or of the wrong variant the
        // encoder cannot construct a deterministic witness and we
        // fall through to `None`, leaving the finding at Class 2.
        SinkConstraintKind::StackOverflowStrcpy => match request.shape_params.as_ref()? {
            ShapeParams::StackOverflow { dst_capacity } => stack_overflow_witness(*dst_capacity),
            _ => None,
        },
        SinkConstraintKind::StackOverflowSprintfUnbounded => {
            match request.shape_params.as_ref()? {
                ShapeParams::StackOverflow { dst_capacity } => {
                    sprintf_unbounded_witness(*dst_capacity)
                }
                _ => None,
            }
        }
        SinkConstraintKind::StackOverflowGets => match request.shape_params.as_ref()? {
            ShapeParams::StackOverflow { dst_capacity } => gets_witness(*dst_capacity),
            _ => None,
        },
        SinkConstraintKind::HeapOverflowAllocArithWraps => match request.shape_params.as_ref()? {
            ShapeParams::AllocArithWraps {
                struct_size,
                ptr_width_bytes,
            } => alloc_arith_wraps_witness(*struct_size, *ptr_width_bytes),
            _ => None,
        },
        SinkConstraintKind::OobWriteUnboundedIndex => match request.shape_params.as_ref()? {
            ShapeParams::OobWriteIndex {
                capacity,
                signed,
                index_width_bytes,
            } => oob_write_index_witness(*capacity, *signed, *index_width_bytes),
            _ => None,
        },
        SinkConstraintKind::FormatStringUserControlled => Some(format_string_witness()),

        // Solver-required shapes  -  route through Z3Backend.
        SinkConstraintKind::CommandInjection
        | SinkConstraintKind::Ssrf
        | SinkConstraintKind::Xxe
        | SinkConstraintKind::PathTraversal
        | SinkConstraintKind::Deserialization
        | SinkConstraintKind::TemplateInjection
        | SinkConstraintKind::CodeInjection
        | SinkConstraintKind::SqlInjection => {
            // These shapes need real SMT to construct payloads that
            // both reach the sink AND satisfy the exploit constraint.
            // Trivial-encoder doesn't ship them; Z3Backend does.
            None
        }
    }
}

/// Construct a stack-overflow-shape witness from a known destination
/// capacity. The witness is a byte sequence of length
/// `dst_capacity + canary_padding`, sufficient to overflow into the
/// next stack frame's saved RBP / return address.
///
/// This is the witness shape for `stack_overflow_strcpy`,
/// `stack_overflow_sprintf_unbounded`, and `stack_overflow_gets`.
/// The rule's `$dst_size` bound variable is the input here; the
/// canary padding (default 16 bytes) gives reliable RBP / saved-RIP
/// overrun on x86-64 SysV.
///
/// Returns `None` when `dst_capacity + CANARY_PADDING` overflows `usize` or
/// exceeds [`MAX_WITNESS_BYTES`]  -  an implausibly large destination is a
/// corrupt shape parameter, not a witnessable stack buffer, and materializing
/// its payload would OOM or (on wrap) emit a silently-short payload that does
/// not overflow. Mirrors the overflow guard in `alloc_arith_wraps_witness`.
#[must_use]
pub fn stack_overflow_witness(dst_capacity: usize) -> Option<ConcreteInput> {
    let total_len = dst_capacity.checked_add(CANARY_PADDING)?;
    if total_len > MAX_WITNESS_BYTES {
        return None;
    }

    // Pattern: "AAAA...AAAA"  -  printable, displays cleanly in repro
    // output. A real attacker would use a controlled payload (e.g.
    // a ROP chain), but for the proof-of-overflow witness this
    // suffices.
    let bytes = vec![b'A'; total_len];

    Some(ConcreteInput {
        bytes,
        render: format!(
            "python3 -c 'import sys; sys.stdout.buffer.write(b\"A\" * {total_len})' | ./victim"
        ),
        entrypoint_kind: EntrypointKind::Cli,
    })
}

/// Witness for `stack_overflow_gets`  -  same shape as
/// `stack_overflow_witness`, with a stdin-shaped repro command.
///
/// Returns `None` under the same capacity guard as `stack_overflow_witness`.
#[must_use]
pub fn gets_witness(dst_capacity: usize) -> Option<ConcreteInput> {
    let mut w = stack_overflow_witness(dst_capacity)?;
    // Reuse the shared CANARY_PADDING owner so the render's byte count can
    // never disagree with the bytes produced by stack_overflow_witness.
    let total_len = dst_capacity + CANARY_PADDING;
    w.render = format!(
        "python3 -c 'import sys; sys.stdout.write(\"A\" * {total_len})' | ./victim"
    );
    Some(w)
}

/// Witness for `format_string_user_controlled`  -  a chain of `%n`
/// specifiers triggers an arbitrary-write primitive. The classical
/// witness is `"%n"` repeated enough times that the format-arg
/// register file is consumed and a stack-resident pointer is
/// dereferenced for write. Eight `%n`s cover both x86-64 SysV
/// (6 integer-arg registers  -  RDI/RSI/RDX/RCX/R8/R9) and ARM64
/// AAPCS64 (8 integer-arg registers  -  X0..X7); two extra %n's at
/// the end ensure the stack is reliably reached on both ABIs.
#[must_use]
pub fn format_string_witness() -> ConcreteInput {
    let payload = "%n%n%n%n%n%n%n%n";
    ConcreteInput {
        bytes: payload.as_bytes().to_vec(),
        render: format!(
            "./victim '{payload}'   # %n chain triggers arbitrary write via va_arg overrun"
        ),
        entrypoint_kind: EntrypointKind::Cli,
    }
}

/// Witness for `stack_overflow_sprintf_unbounded`  -  sprintf into a
/// fixed-capacity destination has no length argument, so any format
/// expansion that exceeds `dst_capacity` overflows the next stack
/// frame's saved RBP / saved RIP on x86-64 SysV.
///
/// The witness is a single `%s` argument whose bytes  -  when expanded
/// by sprintf  -  write `dst_capacity + canary_padding` bytes. Same
/// canary-padding contract as `stack_overflow_witness`. The render
/// hint shows how to invoke the victim with the payload via argv.
///
/// For wide-char sprintf (`vswprintf` etc.) the same shape applies
/// at the bytewise level; the canary-padding in element units stays
/// the same.
///
/// Returns `None` under the same capacity guard as `stack_overflow_witness`.
#[must_use]
pub fn sprintf_unbounded_witness(dst_capacity: usize) -> Option<ConcreteInput> {
    let total_len = dst_capacity.checked_add(CANARY_PADDING)?;
    if total_len > MAX_WITNESS_BYTES {
        return None;
    }

    // sprintf expands `%s` byte-for-byte; the supplied argv string
    // becomes the in-buffer overflow payload. Any printable byte
    // works for the proof; `A` keeps repro output legible.
    let bytes = vec![b'A'; total_len];

    Some(ConcreteInput {
        bytes,
        render: format!(
            "python3 -c 'import sys; sys.stdout.write(\"A\" * {total_len})' \
             | xargs -I{{}} ./victim '{{}}'   # sprintf(dst, \"%s\", argv) overflows dst"
        ),
        entrypoint_kind: EntrypointKind::Cli,
    })
}

/// Witness for `heap_overflow_alloc_arith_wraps`  -  emit the smallest
/// attacker-controlled count `n` such that `n * struct_size` wraps the
/// pointer-width unsigned-multiplication and produces an undersized
/// allocation. The follow-up payload then writes past the wrapped
/// allocation into the next heap chunk header.
///
/// The math: on a `ptr_width_bytes = W`-byte pointer arch the size
/// argument is `usize` (`u32` for W=4, `u64` for W=8). Multiplication
/// wraps modulo `2^(8*W)`. The smallest `n` that wraps for a given
/// `struct_size` is `ceil(2^(8*W) / struct_size)`; the resulting
/// allocation is `(n * struct_size) mod 2^(8*W)`, which is small.
///
/// Returns `None` for `struct_size == 1` (no wrap possible  -  `n * 1`
/// equals `n`) or unsupported pointer widths (anything other than 4
/// or 8 bytes). `struct_size == 0` is admitted as a valid Class-1
/// witness: every allocator call `malloc(n * 0)` yields a 0-byte
/// allocation regardless of `n`, so writing any payload past the
/// allocation is unconditionally a heap overflow (audit
/// 2026-04-27 finding 6).
#[must_use]
pub fn alloc_arith_wraps_witness(
    struct_size: usize,
    ptr_width_bytes: usize,
) -> Option<ConcreteInput> {
    let limit: u128 = match ptr_width_bytes {
        4 => 1u128 << 32,
        8 => 1u128 << 64,
        _ => return None,
    };

    const CANARY_PADDING: u128 = 32;

    // ZST / flexible-array-member shape: the multiplication produces
    // 0 regardless of the count, so `malloc(n * 0)` is a 0-byte
    // allocation. Any subsequent write is an overflow. Witness uses
    // count = 1 (the shortest counter the bug needs) and emits the
    // canary-padding payload directly.
    if struct_size == 0 {
        let n: u128 = 1;
        let wrapped_size: u128 = 0;
        let write_len = wrapped_size + CANARY_PADDING;
        if write_len > usize::MAX as u128 {
            return None;
        }
        let n_le_bytes = n.to_le_bytes();
        if ptr_width_bytes > n_le_bytes.len() {
            return None;
        }
        let mut bytes = Vec::with_capacity(ptr_width_bytes + write_len as usize);
        bytes.extend_from_slice(&n_le_bytes[..ptr_width_bytes]);
        bytes.extend(std::iter::repeat_n(b'B', write_len as usize));
        return Some(ConcreteInput {
            bytes,
            render: format!(
                "# zero-size struct: alloc(n * 0) yields a 0-byte buffer; any write overflows\n\
                 python3 -c 'import struct,sys; sys.stdout.buffer.write(struct.pack(\"<{pack}\", 1) + b\"B\" * {write_len})' | ./victim",
                pack = if ptr_width_bytes == 4 { "I" } else { "Q" }
            ),
            entrypoint_kind: EntrypointKind::Cli,
        });
    }

    if struct_size == 1 {
        // `n * 1 == n` for every `n`; no integer wrap is possible.
        return None;
    }

    // Smallest n such that n * struct_size >= 2^(8*W). Using u128
    // keeps the math overflow-safe across both 32-bit and 64-bit.
    let struct_size_u128 = struct_size as u128;
    let n: u128 = limit.div_ceil(struct_size_u128);
    let wrapped_size: u128 = (n * struct_size_u128) % limit;

    // Payload is the count `n` encoded as a little-endian
    // ptr-width integer, then a write past the wrapped capacity.
    // The write payload is `wrapped_size + canary_padding` bytes
    // of `B` so the corrupted chunk header pattern is unambiguous
    // in heap dumps.
    let write_len = wrapped_size + CANARY_PADDING;

    // Refuse to materialize a payload that exceeds host address
    // space (32-bit hosts encoding a witness for a 64-bit target
    // would otherwise truncate `write_len as usize` and emit a
    // smaller payload than intended, OR attempt a multi-gigabyte
    // Vec allocation and OOM).
    if write_len > usize::MAX as u128 {
        return None;
    }
    let write_len_usize = write_len as usize;

    // Serialize `n` directly from `u128` so we never narrow through
    // `usize` (a 32-bit host emitting a 64-bit count would otherwise
    // truncate the upper 4 bytes of a count > u32::MAX).
    let n_le_bytes = n.to_le_bytes();
    if ptr_width_bytes > n_le_bytes.len() {
        return None;
    }

    let mut bytes = Vec::with_capacity(ptr_width_bytes + write_len_usize);
    bytes.extend_from_slice(&n_le_bytes[..ptr_width_bytes]);
    bytes.extend(std::iter::repeat_n(b'B', write_len_usize));

    Some(ConcreteInput {
        bytes,
        render: format!(
            "# alloc count = {n} (struct_size={struct_size}, wraps to {wrapped_size} bytes on \
             {bits}-bit ptr arch); subsequent write of {write_len} bytes overruns the wrapped chunk\n\
             python3 -c 'import struct,sys; sys.stdout.buffer.write(struct.pack(\"<{pack}\", {n}) + b\"B\" * {write_len})' | ./victim",
            bits = ptr_width_bytes * 8,
            pack = if ptr_width_bytes == 4 { "I" } else { "Q" }
        ),
        entrypoint_kind: EntrypointKind::Cli,
    })
}

/// Witness for `oob_write_unbounded_index`  -  emit the smallest
/// attacker-controlled index that overflows `capacity` for an
/// indexed store with no bounds check. The two variants:
///
/// - `signed = false`: positive overflow. Witness emits `capacity`
///   itself (the first illegal index) plus a single payload byte.
///   The store at `buf[capacity]` writes one slot past the array
///   into the next heap-chunk metadata (or saved frame on stack).
///
/// - `signed = true`: negative-index variant. A signed index
///   variable can be `-1`, which after pointer arithmetic indexes
///   `buf[-1]`  -  the byte preceding the buffer. On the stack this
///   is typically the saved RBP. Witness emits `-1` as a signed
///   integer of the appropriate width plus a payload byte.
///
/// `index_width_bytes` matches the sink ABI: 4 for `int` /
/// `unsigned`, 8 for `size_t` / `ssize_t` on 64-bit Linux. Returns
/// `None` for an unsupported width or if `capacity` does not fit
/// in the chosen width (audit 2026-04-27 finding 8  -  silent
/// `u32::MAX` clamping was a real bug).
#[must_use]
pub fn oob_write_index_witness(
    capacity: usize,
    signed: bool,
    index_width_bytes: usize,
) -> Option<ConcreteInput> {
    const PAYLOAD_BYTE: u8 = b'C';
    if !matches!(index_width_bytes, 4 | 8) {
        return None;
    }

    if signed {
        // Negative-index payload: `-1` encoded as a signed integer
        // of the requested width. The two's-complement
        // representation of -1 is all-ones in any width.
        let mut bytes = Vec::with_capacity(index_width_bytes + 1);
        let pack = if index_width_bytes == 4 { "i" } else { "q" };
        if index_width_bytes == 4 {
            bytes.extend_from_slice(&(-1_i32).to_le_bytes());
        } else {
            bytes.extend_from_slice(&(-1_i64).to_le_bytes());
        }
        bytes.push(PAYLOAD_BYTE);
        Some(ConcreteInput {
            bytes,
            render: format!(
                "# index = -1 (negative-index variant; capacity={capacity} ignored); \
                 buf[-1] writes one byte before the array, clobbering saved RBP / chunk header\n\
                 python3 -c 'import struct,sys; sys.stdout.buffer.write(struct.pack(\"<{pack}\", -1) + b\"C\")' | ./victim"
            ),
            entrypoint_kind: EntrypointKind::Cli,
        })
    } else {
        // Positive overflow: smallest illegal index = `capacity`.
        // Refuse to silently clamp when capacity exceeds the index
        // width; the caller must use a wider sink or accept that
        // no Class 1 witness is possible here.
        let pack = if index_width_bytes == 4 { "I" } else { "Q" };
        let illegal_index_bytes: Vec<u8> = if index_width_bytes == 4 {
            let v: u32 = capacity.try_into().ok()?;
            v.to_le_bytes().to_vec()
        } else {
            let v: u64 = capacity as u64;
            v.to_le_bytes().to_vec()
        };
        let mut bytes = Vec::with_capacity(index_width_bytes + 1);
        bytes.extend_from_slice(&illegal_index_bytes);
        bytes.push(PAYLOAD_BYTE);
        Some(ConcreteInput {
            bytes,
            render: format!(
                "# index = {capacity} (first illegal slot); \
                 buf[{capacity}] writes one byte past the array end\n\
                 python3 -c 'import struct,sys; sys.stdout.buffer.write(struct.pack(\"<{pack}\", {capacity}) + b\"C\")' | ./victim"
            ),
            entrypoint_kind: EntrypointKind::Cli,
        })
    }
}


/// Default shape parameters used when a rule primitive maps to a
/// closed-form witness and the rule did not thread per-finding shape
/// parameters. These match the 64-bit Linux launch corpus defaults.
const DEFAULT_CAPACITY: usize = 64;
const DEFAULT_PTR_WIDTH: usize = 8;
const DEFAULT_STRUCT_SIZE: usize = 8;
const DEFAULT_INDEX_WIDTH: usize = 4;

/// Map a rule primitive string to the canonical sink kind and default
/// shape parameters for closed-form witnesses.
///
/// This is the ONE-PLACE owner of the dispatch table; surgec consumes
/// it instead of re-implementing the same `if/else` chain.
#[must_use]
pub fn classify_primitive(primitive: &str) -> Option<(crate::SinkConstraintKind, Option<crate::ShapeParams>)> {
    // Order matters: more-specific patterns first so "sprintf" and
    // "strcpy" route to their dedicated witnesses rather than the
    // generic stack-overflow shape.
    if primitive.contains("sprintf") || primitive.contains("snprintf_unbounded") {
        Some((
            crate::SinkConstraintKind::StackOverflowSprintfUnbounded,
            Some(crate::ShapeParams::StackOverflow { dst_capacity: DEFAULT_CAPACITY }),
        ))
    } else if primitive.contains("strcpy") || primitive.contains("strcat") {
        Some((
            crate::SinkConstraintKind::StackOverflowStrcpy,
            Some(crate::ShapeParams::StackOverflow { dst_capacity: DEFAULT_CAPACITY }),
        ))
    } else if primitive == "gets" || primitive.contains("_gets") {
        Some((
            crate::SinkConstraintKind::StackOverflowGets,
            Some(crate::ShapeParams::StackOverflow { dst_capacity: DEFAULT_CAPACITY }),
        ))
    } else if primitive.contains("format_string") || primitive.contains("printf") {
        Some((crate::SinkConstraintKind::FormatStringUserControlled, None))
    } else if primitive.contains("alloc_arith") || primitive.contains("alloc_wraps") {
        Some((
            crate::SinkConstraintKind::HeapOverflowAllocArithWraps,
            Some(crate::ShapeParams::AllocArithWraps {
                struct_size: DEFAULT_STRUCT_SIZE,
                ptr_width_bytes: DEFAULT_PTR_WIDTH,
            }),
        ))
    } else if primitive.contains("off_by_one") {
        // Off-by-one: allocation = capacity, write = capacity+1.
        // Use struct_size=2 so wrapped_size is small but still wraps.
        Some((
            crate::SinkConstraintKind::HeapOverflowAllocArithWraps,
            Some(crate::ShapeParams::AllocArithWraps {
                struct_size: 2,
                ptr_width_bytes: DEFAULT_PTR_WIDTH,
            }),
        ))
    } else if primitive.contains("oob_write") || primitive.contains("oob_index_write") {
        Some((
            crate::SinkConstraintKind::OobWriteUnboundedIndex,
            Some(crate::ShapeParams::OobWriteIndex {
                capacity: DEFAULT_CAPACITY,
                signed: false,
                index_width_bytes: DEFAULT_INDEX_WIDTH,
            }),
        ))
    } else if primitive.contains("oob_read") || primitive.contains("oob_index_read") {
        // Read shape uses the same illegal-index payload as write.
        Some((
            crate::SinkConstraintKind::OobWriteUnboundedIndex,
            Some(crate::ShapeParams::OobWriteIndex {
                capacity: DEFAULT_CAPACITY,
                signed: false,
                index_width_bytes: DEFAULT_INDEX_WIDTH,
            }),
        ))
    } else {
        None
    }
}

#[cfg(test)]
#[path = "sink_payloads_tests.rs"]
mod tests;