js_deobfuscator/format/
brace.rs

1//! Brace wrapping: ensure if/else/for/while bodies have `{}` blocks.
2//!
3//! `if (x) return 1;` → `if (x) { return 1; }`
4
5use oxc::allocator::Allocator;
6use oxc::ast::ast::{Program, Statement};
7use oxc::semantic::Scoping;
8use oxc::span::SPAN;
9
10use oxc_traverse::{Traverse, TraverseCtx, traverse_mut};
11
12use crate::engine::error::Result;
13use crate::engine::module::{Module, TransformResult};
14
15/// Brace wrapping module.
16pub struct BraceWrapper;
17
18impl Module for BraceWrapper {
19    fn name(&self) -> &'static str {
20        "BraceWrapper"
21    }
22
23    fn transform<'a>(
24        &mut self,
25        allocator: &'a Allocator,
26        program: &mut Program<'a>,
27        scoping: Scoping,
28    ) -> Result<TransformResult> {
29        let mut visitor = BraceVisitor { modifications: 0 };
30        let scoping = traverse_mut(&mut visitor, allocator, program, scoping, ());
31        Ok(TransformResult { modifications: visitor.modifications, scoping })
32    }
33}
34
35struct BraceVisitor {
36    modifications: usize,
37}
38
39impl<'a> Traverse<'a, ()> for BraceVisitor {
40    fn exit_statement(
41        &mut self,
42        stmt: &mut Statement<'a>,
43        ctx: &mut TraverseCtx<'a, ()>,
44    ) {
45        match stmt {
46            Statement::IfStatement(if_stmt) => {
47                if wrap_if_needed(&mut if_stmt.consequent, ctx) {
48                    self.modifications += 1;
49                }
50                if let Some(alt) = &mut if_stmt.alternate {
51                    // Don't wrap `else if`
52                    if !matches!(alt, Statement::IfStatement(_)) && wrap_if_needed(alt, ctx) {
53                        self.modifications += 1;
54                    }
55                }
56            }
57            Statement::WhileStatement(w) => {
58                if wrap_if_needed(&mut w.body, ctx) {
59                    self.modifications += 1;
60                }
61            }
62            Statement::ForStatement(f) => {
63                if wrap_if_needed(&mut f.body, ctx) {
64                    self.modifications += 1;
65                }
66            }
67            Statement::ForInStatement(f) => {
68                if wrap_if_needed(&mut f.body, ctx) {
69                    self.modifications += 1;
70                }
71            }
72            Statement::ForOfStatement(f) => {
73                if wrap_if_needed(&mut f.body, ctx) {
74                    self.modifications += 1;
75                }
76            }
77            _ => {}
78        }
79    }
80}
81
82/// Wrap a statement in a block if it isn't one already. Returns true if wrapped.
83fn wrap_if_needed<'a>(stmt: &mut Statement<'a>, ctx: &mut TraverseCtx<'a, ()>) -> bool {
84    if matches!(stmt, Statement::BlockStatement(_)) {
85        return false;
86    }
87
88    let inner = std::mem::replace(stmt, ctx.ast.statement_empty(SPAN));
89    let mut stmts = ctx.ast.vec();
90    stmts.push(inner);
91    *stmt = ctx.ast.statement_block(SPAN, stmts);
92    true
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use oxc::codegen::Codegen;
99    use oxc::parser::Parser;
100    use oxc::semantic::SemanticBuilder;
101    use oxc::span::SourceType;
102
103    fn wrap(source: &str) -> String {
104        let allocator = Allocator::default();
105        let mut program = Parser::new(&allocator, source, SourceType::mjs()).parse().program;
106        let scoping = SemanticBuilder::new().build(&program).semantic.into_scoping();
107        let mut module = BraceWrapper;
108        module.transform(&allocator, &mut program, scoping).unwrap();
109        Codegen::new().build(&program).code
110    }
111
112    #[test]
113    fn test_if_body() {
114        let code = wrap("if (x) return 1;");
115        assert!(code.contains("{"), "should add braces: {code}");
116    }
117
118    #[test]
119    fn test_already_braced() {
120        let code = wrap("if (x) { return 1; }");
121        assert!(code.contains("{"), "should keep braces: {code}");
122    }
123
124    #[test]
125    fn test_while_body() {
126        let code = wrap("while (x) x--;");
127        assert!(code.contains("{"), "should wrap while body: {code}");
128    }
129
130    #[test]
131    fn test_else_if_not_wrapped() {
132        let code = wrap("if (a) {} else if (b) {}");
133        // else if should not become else { if }
134        assert!(!code.contains("else {\n\tif") && !code.contains("else {\n  if"), "got: {code}");
135    }
136}
js_deobfuscator/format/brace.rs

js_deobfuscator/format/
brace.rs
