JS `Control Flow Graph` (CFG) 重建与程序流分析

各位听众，早上好！今天咱们聊聊一个听起来高大上，实际上也挺高大上的话题：JavaScript的控制流图（CFG）重建与程序流分析。别怕，我会尽量把它讲得接地气一点，争取让各位听完之后，觉得自己也能撸一个出来。

一、什么是控制流图（CFG）？

想象一下你写了一段JavaScript代码，这代码就像一条蜿蜒的小路，里面充满了各种岔路口和选择。控制流图呢，就是把这条小路画成一张地图，清清楚楚地标明每个路口，以及从一个路口到另一个路口的方向。

更正式地说，控制流图是一个有向图，其中：

• 节点（Nodes）： 代表代码中的基本块（Basic Blocks）。基本块是指一段顺序执行的代码，没有跳转语句，只有一个入口和一个出口。
• 边（Edges）： 代表控制流的转移。例如，if语句的两个分支，for循环的循环体和循环退出等等。

用人话说，节点就是一段挨在一起的代码，中间没有if、for、return之类的，边就是代码执行的顺序。

二、为什么要重建CFG？

有了CFG，我们就能做很多有意思的事情，比如：

• 代码优化： 分析CFG可以帮助我们找到冗余的代码，进行死代码消除、常量传播等优化。
• 漏洞检测： 分析CFG可以帮助我们发现潜在的安全漏洞，例如，变量未初始化就使用、数据流污染等。
• 代码覆盖率测试： CFG可以帮助我们确定代码覆盖率，确保测试用例覆盖了所有可能的执行路径。
• 程序理解： 反编译/混淆代码分析，CFG可以帮助安全人员理解混淆的代码逻辑。

总之，CFG就像一个放大镜，让我们能够更清晰地观察代码的内部运作。

三、CFG重建的步骤

重建CFG的过程，大致可以分为以下几个步骤：

1. 词法分析（Lexical Analysis）： 将源代码分解成一个个的token，比如关键字、变量名、运算符等等。
2. 语法分析（Syntactic Analysis）： 将token流转换成抽象语法树（AST）。AST是代码的结构化表示，更容易进行分析。
3. 基本块划分（Basic Block Identification）： 将AST划分成基本块。
4. 控制流连接（Control Flow Linking）： 确定基本块之间的控制流关系，建立CFG的边。

咱们一步一步来，用代码说话。

3.1 词法分析和语法分析（AST生成）

这部分通常依赖现成的工具，比如acorn、esprima、babel等等。这里以acorn为例：

const acorn = require("acorn");

const code = `
  function add(a, b) {
    if (a > 0) {
      return a + b;
    } else {
      return a - b;
    }
  }
`;

const ast = acorn.parse(code, {
  ecmaVersion: 2020,
  sourceType: "script", // or "module"
});

console.log(JSON.stringify(ast, null, 2)); // 打印AST

这段代码会把code变量中的JavaScript代码解析成一个AST对象，然后打印出来。 AST 看起来比较复杂， 我们可以用可视化工具观察，如 AST Explorer。

3.2 基本块划分

基本块划分的原则是：

• 基本块的第一个语句是入口语句。
• 基本块的最后一个语句是出口语句。
• 基本块内部没有跳转语句。

这意味着，if、for、while、return、throw等语句都是基本块的边界。

下面是一个简单的基本块划分的示例代码：

function splitIntoBasicBlocks(ast) {
  const basicBlocks = [];
  let currentBlock = [];

  function addStatementToBlock(statement) {
    currentBlock.push(statement);
  }

  function createNewBlock() {
    if (currentBlock.length > 0) {
      basicBlocks.push(currentBlock);
    }
    currentBlock = [];
  }

  function walk(node) {
    if (!node || typeof node !== 'object') {
      return;
    }

    switch (node.type) {
      case 'Program':
      case 'BlockStatement':
        node.body.forEach(walk);
        break;

      case 'FunctionDeclaration':
        addStatementToBlock(node); // 整个函数声明作为一个基本块的开始
        walk(node.body);
        break;

      case 'IfStatement':
        // 在if语句前创建一个新块
        createNewBlock();
        addStatementToBlock(node); // if 语句本身
        // 处理consequent (then) block
        walk(node.consequent);
        // 处理alternate (else) block, 如果存在
        if (node.alternate) {
          walk(node.alternate);
        }
        // 在if语句后创建一个新块
        createNewBlock();
        break;

      case 'ForStatement':
      case 'WhileStatement':
        // 在循环语句前创建一个新块
        createNewBlock();
        addStatementToBlock(node);
        walk(node.body);
        createNewBlock();
        break;

      case 'ReturnStatement':
      case 'ThrowStatement':
        addStatementToBlock(node);
        createNewBlock(); // return/throw 之后创建一个新块
        break;

      default:
        addStatementToBlock(node);
    }
  }

  walk(ast);
  createNewBlock(); // 确保最后一个块被添加

  return basicBlocks;
}

const basicBlocks = splitIntoBasicBlocks(ast);

basicBlocks.forEach((block, index) => {
  console.log(`Basic Block ${index + 1}:`);
  block.forEach(statement => {
    console.log(statement.type); // 打印语句类型
  });
  console.log('---');
});

这段代码遍历AST，根据上面提到的原则，将代码划分成基本块。 splitIntoBasicBlocks 函数接收AST作为输入，返回一个二维数组，其中每个子数组代表一个基本块，包含了该块中的AST节点。

3.3 控制流连接

控制流连接的目标是确定基本块之间的控制流关系。这需要分析每个基本块的最后一个语句，以及if、for、while等语句的条件。

function connectBasicBlocks(basicBlocks) {
    const cfg = {};

    for (let i = 0; i < basicBlocks.length; i++) {
        cfg[i] = {
            block: basicBlocks[i],
            next: [],
            type: 'normal' // 默认是普通块
        };
    }

    for (let i = 0; i < basicBlocks.length; i++) {
        const currentBlock = cfg[i];
        const lastStatement = currentBlock.block[currentBlock.block.length - 1];

        if (!lastStatement) {
            // 空块，连接到下一个块
            if (i + 1 < basicBlocks.length) {
                currentBlock.next.push(i + 1);
            }
            continue;
        }

        switch (lastStatement.type) {
            case 'IfStatement':
                // if 语句，连接到 then 和 else 块
                if (lastStatement.consequent) {
                    // 找到consequent block对应的 index
                    const consequentIndex = findBlockIndex(cfg, lastStatement.consequent);
                    if (consequentIndex !== -1) {
                        currentBlock.next.push(consequentIndex);
                    }
                }
                if (lastStatement.alternate) {
                    // 找到alternate block对应的 index
                    const alternateIndex = findBlockIndex(cfg, lastStatement.alternate);
                     if (alternateIndex !== -1) {
                        currentBlock.next.push(alternateIndex);
                    }
                } else {
                    // 如果没有 else 块，连接到 if 语句后的块
                    if (i + 1 < basicBlocks.length) {
                        currentBlock.next.push(i + 1);
                    }
                }
                currentBlock.type = 'conditional';
                break;

            case 'ForStatement':
            case 'WhileStatement':
                // 循环语句，连接到循环体
                const loopBodyIndex = findBlockIndex(cfg, lastStatement.body);
                if (loopBodyIndex !== -1) {
                    currentBlock.next.push(loopBodyIndex);
                }
                // 同时，循环体也连接到循环语句本身（用于循环）
                cfg[loopBodyIndex].next.push(i);
                // 循环结束后，连接到下一个块
                if (i + 1 < basicBlocks.length) {
                    currentBlock.next.push(i + 1);
                }
                currentBlock.type = 'loop';
                break;

            case 'ReturnStatement':
            case 'ThrowStatement':
                // return 或 throw 语句，没有后续块
                currentBlock.type = 'exit';
                break;

            default:
                // 其他情况，连接到下一个块
                if (i + 1 < basicBlocks.length) {
                    currentBlock.next.push(i + 1);
                }
        }
    }

    return cfg;
}

function findBlockIndex(cfg, astNode) {
    for (const index in cfg) {
        const block = cfg[index].block;
        for(const node of block){
            if(node === astNode){
                return parseInt(index); // 返回块的索引
            }
        }
    }
    return -1; // 未找到
}

const cfg = connectBasicBlocks(basicBlocks);

// 打印CFG
for (const blockIndex in cfg) {
    const blockInfo = cfg[blockIndex];
    console.log(`Block ${blockIndex}:`);
    blockInfo.block.forEach(statement => console.log(`  ${statement.type}`));
    console.log(`  Next: ${blockInfo.next.join(', ')}`);
    console.log(`  Type: ${blockInfo.type}`);
}

这段代码遍历基本块，根据每个基本块的最后一个语句，确定控制流的转移方向。 connectBasicBlocks 函数接收基本块数组作为输入，返回一个CFG对象。 findBlockIndex 函数用于根据AST节点查找基本块的索引。

四、程序流分析

有了CFG，我们就可以进行各种程序流分析了。这里举几个简单的例子：

4.1 可达性分析（Reachability Analysis）

判断某个基本块是否可达，也就是从程序的入口点是否可以通过控制流到达该基本块。这可以用来检测死代码。

function reachabilityAnalysis(cfg) {
  const reachable = {};
  const entryPoint = 0; // 假设第一个块是入口点

  function dfs(blockIndex) {
    if (reachable[blockIndex]) {
      return; // 已经访问过
    }

    reachable[blockIndex] = true;

    cfg[blockIndex].next.forEach(nextIndex => {
      dfs(nextIndex);
    });
  }

  dfs(entryPoint);

  for (const blockIndex in cfg) {
    console.log(`Block ${blockIndex} is reachable: ${!!reachable[blockIndex]}`);
  }
}

reachabilityAnalysis(cfg);

4.2 数据流分析（Data Flow Analysis）

例如，活跃变量分析（Live Variable Analysis），判断某个变量在某个基本块的出口处是否仍然可能被使用。这可以用来进行寄存器分配和死代码消除。

这是一个简化的示例，没有考虑所有情况，只是为了说明概念：

function liveVariableAnalysis(cfg) {
  const liveVariables = {};

  // 初始化每个块的活跃变量集合为空
  for (const blockIndex in cfg) {
    liveVariables[blockIndex] = new Set();
  }

  // 迭代直到活跃变量集合不再改变
  let changed = true;
  while (changed) {
    changed = false;

    // 反向遍历CFG
    for (let blockIndex = Object.keys(cfg).length - 1; blockIndex >= 0; blockIndex--) {
      const block = cfg[blockIndex].block;
      const prevLiveVariables = new Set(liveVariables[blockIndex]); // 复制之前的集合

      // 模拟变量使用和定义
      for (let i = block.length - 1; i >= 0; i--) {
        const statement = block[i];

        // 简化的分析，只考虑Identifier类型的节点
        if (statement.type === 'Identifier') {
          // 假设这是一个变量使用
          liveVariables[blockIndex].add(statement.name);
        } else if (statement.type === 'VariableDeclaration') {
          // 假设这是一个变量定义，移除活跃变量集合中的变量
          statement.declarations.forEach(declaration => {
            if (declaration.id.type === 'Identifier') {
              liveVariables[blockIndex].delete(declaration.id.name);
            }
          });
        }
        //实际分析需要根据AST节点类型做更细致的分析，这里只做演示
      }

      // 从后继块传递活跃变量
      cfg[blockIndex].next.forEach(nextIndex => {
        for (const variable of liveVariables[nextIndex]) {
          liveVariables[blockIndex].add(variable);
        }
      });

      // 检查活跃变量集合是否发生改变
      if (liveVariables[blockIndex].size !== prevLiveVariables.size) {
        changed = true;
      } else {
        for (const variable of liveVariables[blockIndex]) {
          if (!prevLiveVariables.has(variable)) {
            changed = true;
            break;
          }
        }
      }
    }
  }

  // 打印结果
  for (const blockIndex in cfg) {
    console.log(`Block ${blockIndex} live variables: ${Array.from(liveVariables[blockIndex]).join(', ')}`);
  }
}

// 示例代码：
const code2 = `
  function example() {
    let x = 10;
    let y = x + 5;
    if (y > 15) {
      return y;
    } else {
      return x;
    }
  }
`;

const ast2 = acorn.parse(code2, { ecmaVersion: 2020, sourceType: 'script' });
const basicBlocks2 = splitIntoBasicBlocks(ast2);
const cfg2 = connectBasicBlocks(basicBlocks2);

liveVariableAnalysis(cfg2);

五、总结

CFG重建和程序流分析是一个相对复杂的过程，涉及到词法分析、语法分析、基本块划分、控制流连接等多个步骤。但是，有了CFG，我们就可以进行各种代码优化、漏洞检测、代码覆盖率测试等工作，从而提高代码的质量和安全性。

当然，上面的代码只是一个非常简化的示例，实际应用中需要考虑更多的情况，例如异常处理、函数调用等等。此外，还可以使用更高级的算法和数据结构来提高CFG重建和程序流分析的效率。

希望今天的讲座能够帮助大家对CFG重建和程序流分析有一个初步的了解。如果大家有任何问题，欢迎提问！