pyrograph 0.1.0

GPU-accelerated taint analysis for supply chain malware detection
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

pub const SHADER_SRC: &str = r#"
@group(0) @binding(0) var<storage, read> node_labels: array<u32>;
@group(0) @binding(1) var<storage, read> edge_offsets: array<u32>;
@group(0) @binding(2) var<storage, read> edge_targets: array<u32>;
@group(0) @binding(3) var<storage, read> source_nodes: array<u32>;
@group(0) @binding(4) var<storage, read_write> findings: array<vec4<u32>>;
@group(0) @binding(5) var<storage, read_write> finding_count: array<atomic<u32>>;
@group(0) @binding(6) var<uniform> params: Params;
@group(0) @binding(7) var<storage, read_write> visited_set: array<u32>;

struct Params {
    num_sources: u32,
    num_nodes: u32,
    max_findings: u32,
    max_depth: u32,
    words_per_source: u32,
    _pad0: u32,
    _pad1: u32,
    _pad2: u32,
}

const LABEL_SOURCE: u32 = 1u;
const LABEL_SINK: u32 = 2u;
const MAX_BFS_QUEUE: u32 = 4096u;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let source_idx = gid.x;
    if source_idx >= params.num_sources { return; }
    let start_node = source_nodes[source_idx];
    
    // BFS from start_node
    var queue: array<u32, 4096>;
    var queue_head: u32 = 0u;
    var queue_tail: u32 = 1u;
    queue[0] = start_node;
    
    let base = source_idx * params.words_per_source;
    
    // Mark start as visited
    let start_word = start_node / 32u;
    let start_bit = 1u << (start_node % 32u);
    visited_set[base + start_word] = visited_set[base + start_word] | start_bit;

    var depth: u32 = 0u;

    while queue_head < queue_tail && depth < params.max_depth {
        let level_end = queue_tail;
        while queue_head < level_end {
            let node = queue[queue_head];
            queue_head += 1u;

            // node_labels packed format: (kind << 24) | (label_type << 16) | (label_id & 0xFFFF)
            let label_type = (node_labels[node] >> 16u) & 0xFFu;

            // Stop traversal at sanitizers (label_type 4)
            if label_type == 4u {
                continue;
            }

            // Check if this node is a sink (skip self-taint on source==sink nodes)
            if (label_type == LABEL_SINK || label_type == 3u) && node != start_node { // 3 is BOTH
                let idx = atomicAdd(&finding_count[0], 1u);
                if idx < params.max_findings {
                    findings[idx] = vec4<u32>(start_node, node, depth, source_idx);
                }
            }

            // Explore neighbors
            let edge_start = edge_offsets[node];
            let edge_end = edge_offsets[node + 1u];
            for (var e = edge_start; e < edge_end; e++) {
                let tgt = edge_targets[e];
                let word = tgt / 32u;
                let bit = 1u << (tgt % 32u);
                let visited_word = visited_set[base + word];
                if (visited_word & bit) == 0u {
                    visited_set[base + word] = visited_word | bit;
                    if queue_tail < MAX_BFS_QUEUE {
                        queue[queue_tail] = tgt;
                        queue_tail += 1u;
                    }
                }
            }
        }
        depth += 1u;
    }
}
"#;