pyrograph 0.1.0

GPU-accelerated taint analysis for supply chain malware detection
Documentation
use std::collections::HashMap;

use super::*;
use super::expressions::prop_name;
use crate::ir::{EdgeKind, NodeKind};

pub struct Scope {
    bindings: HashMap<String, NodeId>,
}

impl Scope {
    pub fn new() -> Self {
        Self {
            bindings: HashMap::new(),
        }
    }

    pub(crate) fn define(&mut self, name: String, id: NodeId) {
        self.bindings.insert(name, id);
    }

    pub(crate) fn resolve(&self, name: &str) -> Option<NodeId> {
        self.bindings.get(name).copied()
    }
}

impl JsParser {
    pub(crate) fn cur(&mut self) -> &mut Scope {
        if self.scopes.is_empty() {
            self.scopes.push(Scope::new());
        }
        // SAFETY: the if-block above guarantees scopes is non-empty
        #[allow(clippy::unwrap_used)]
        self.scopes.last_mut().unwrap()
    }

    pub(crate) fn resolve(&self, name: &str) -> Option<NodeId> {
        self.scopes.iter().rev().find_map(|s| s.resolve(name))
    }

    pub(crate) fn var(&mut self, name: &str) -> NodeId {
        self.resolve(name).unwrap_or_else(|| {
            let id = self.graph.add_node(NodeKind::Variable, name.to_string(), None);
            self.cur().define(name.to_string(), id);
            id
        })
    }

    pub(crate) fn lit(&mut self) -> NodeId {
        self.graph.add_node(NodeKind::Literal, "<lit>".to_string(), None)
    }

    pub(crate) fn name(&self, id: NodeId) -> String {
        self.graph.node(id).map(|n| n.name.clone()).unwrap_or_default()
    }

    pub(crate) fn flow(&mut self, from: NodeId, to: NodeId) {
        self.graph.add_edge(from, to, EdgeKind::Assignment);
    }

    pub(crate) fn bind_pat(&mut self, pat: &Pat) {
        let _ = self.bind_pat_with_source(pat, None);
    }

    pub(crate) fn bind_pat_with_source(&mut self, pat: &Pat, source: Option<NodeId>) -> Vec<NodeId> {
        let mut ids = Vec::new();
        match pat {
            Pat::Ident(i) => {
                let id = self
                    .graph
                    .add_node(NodeKind::Variable, i.id.sym.to_string(), None);
                if let Some(src) = source {
                    self.flow(src, id);
                    if let Some(alias) = self.node_binding_alias(src) {
                        if let Some(node) = self.graph.node_mut(id) {
                            node.alias = Some(alias);
                        }
                    }
                    self.materialize_composite_bindings(&i.id.sym.to_string(), src);
                }
                self.cur().define(i.id.sym.to_string(), id);
                ids.push(id);
            }
            Pat::Array(a) => {
                for (index, elem) in a.elems.iter().enumerate() {
                    let Some(elem) = elem else {
                        continue;
                    };
                    let elem_source = match elem {
                        Pat::Rest(_) => source,
                        _ => source.map(|src| self.indexed_binding_source(src, index)),
                    };
                    ids.extend(self.bind_pat_with_source(elem, elem_source));
                }
            }
            Pat::Object(o) => {
                for p in &o.props {
                    match p {
                        ObjectPatProp::KeyValue(kv) => {
                            let prop_source = match (&kv.key, source) {
                                (PropName::Computed(_), src) => src,
                                (key, Some(src)) => prop_name(key)
                                    .map(|name| self.member_binding_source(src, &name))
                                    .or(Some(src)),
                                (_, None) => None,
                            };
                            ids.extend(self.bind_pat_with_source(&kv.value, prop_source));
                        }
                        ObjectPatProp::Assign(a) => {
                            let id = self
                                .graph
                                .add_node(NodeKind::Variable, a.key.sym.to_string(), None);
                            if let Some(src) =
                                source.map(|src| self.member_binding_source(src, &a.key.sym.to_string()))
                            {
                                self.flow(src, id);
                            }
                            if let Some(ref val_expr) = a.value {
                                let default_val = self.eval_expr(val_expr);
                                self.flow(default_val, id);
                            }
                            self.cur().define(a.key.sym.to_string(), id);
                            ids.push(id);
                        }
                        ObjectPatProp::Rest(r) => {
                            ids.extend(self.bind_pat_with_source(&r.arg, source))
                        }
                    }
                }
            }
            Pat::Rest(r) => ids.extend(self.bind_pat_with_source(&r.arg, source)),
            _ => {}
        }
        ids
    }
}