Zenn
🔫

clang で「安全なループを今、書き換えて！」みる

2022/08/28に公開
C++
AST
C
clang
tech

TL;DR

https://lycoris-recoil.com/

あるいは、clang プラグインを使って C言語の AST を弄り倒す話。

これはなに

リコリス・リコイル、みなさん好きですよね。私も大好きです。頼むからちさたき幸せになってほしい

そんなリコリコの ED テーマ、さユりさんの歌う「花の塔」

https://www.youtube.com/watch?v=H2r25lVcIHw

のラスサビ前の歌詞は以下の通りです。

僕は選んでみたいの 高鳴る心 謎だらけの空を
安全なループを今、書き換えて！ [1]

書き換えてやろうじゃありませんか。

安全なループを書き換えるには

安全なループ

前提として、あんまり困難な言語を選んでしまうと困難になってしまうので、ここは無難に C 言語を選びます。

さて、プログラムにおいて「安全」という言葉はとても難しく、特にC言語においては不安全な処理を容易に惹起できてしまうのですが[2]、安全なループについて真面目に考えるのではなく、ひとまず無限ループから脱出することを考えてみたいと思います。

無限ループというのは、たとえば以下の通り。

loop.c
#include <stdio.h>

int main(void) {
   for (;;) {
       printf("窓に飾った絵画をなぞる\n");
   }
   printf("謎だらけの空\n");
   return 0;
}

みなさんお手持ちの clang でコンパイルしてやると、安全なループの中でひたすら窓に飾った絵画をなぞることができます。

$ clang loop.c -o loop.o
$ ./loop.o
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
窓に飾った絵画をなぞる
^C⏎

書き換える

我々の手元にはソースコードがあるので、書き換える手段は実に様々あります。たとえば、vim を使ってただ編集することも「書き換える」ことの範疇ですし、sed で正規表現を掛けたり、はたまたあなたのお好きなプログラム言語で処理して結果を上書きしてもよいでしょう。
が、今回はせっかくのラスサビ前の盛り上げたいところですので、もうちょっとドラスティックにいきましょう。

C言語の処理系といえば clang ですけれども[3]、大変ありがたいことに clang は プラグインをサポートしているので、これを使って書き換えていくのが今回の趣旨です。

書き換えるための基礎知識

以下の環境です。

Ubuntu 20.04.4 LTS (focal) on WSL2, Windows 11 Pro, 22H2 (22621.382)

とりあえず clang をビルドする

LLVM clang をビルドします。LLVM の README に書いてある通りです。
どこのご家庭にもある cmake と Ninja を使います。ご利用の環境に合わせて -jN は適当に調節してください[4]。

$ git clone https://github.com/llvm/llvm-project.git
$ cd llvm-project
$ git switch release/14.x
$ cmake -S llvm -B build -G "Ninja" \
    -DLLVM_ENABLE_PROJECTS="clang;clang-tools-extra" \
    -DCLANG_BUILD_EXAMPLES=1 \
    -DCMAKE_BUILD_TYPE=Debug \
    -DBUILD_SHARED_LIBS=ON
... (略) ...
$ cmake --build build -j8
... (略) ...

コーヒーでも飲んでビルドを待てば、やがて出来上がりです。lib にサンプルのプラグイン (PrintFunctionNames.so) が入っているはずなので確認してみましょう。

$ ./build/bin/clang -v
clang version 14.0.6 (https://github.com/llvm/llvm-project.git f28c006a5895fc0e329fe15fead81e37457cb1d1)
Target: x86_64-unknown-linux-gnu
Thread model: posix
InstalledDir: /home/karno/lycoreco/llvm-project/./build/bin
Found candidate GCC installation: /usr/lib/gcc/x86_64-linux-gnu/9
Selected GCC installation: /usr/lib/gcc/x86_64-linux-gnu/9
Candidate multilib: .;@m64
Selected multilib: .;@m64
$ ls ./build/lib/ | grep "PrintFunctionNames"
PrintFunctionNames.so

サンプルプラグイン PrintFunctionNames を調べる

さて、プラグインの例として ./build/lib に PrintFunctionNames.so が入っているので、ひとまず適当に動かしてみましょう。

$ ./build/bin/clang -Xclang -load -Xclang ./build/lib/PrintFunctionNames.so -Xclang -add-plugin -Xclang print-fns ../loop.c
top-level-decl: "size_t"
top-level-decl: "va_list"
top-level-decl: "__gnuc_va_list"
top-level-decl: "__u_char"
top-level-decl: "__u_short"
top-level-decl: "__u_int"
top-level-decl: "__u_long"
top-level-decl: "__int8_t"
... (略) ...
top-level-decl: "ftrylockfile"
top-level-decl: "funlockfile"
top-level-decl: "__uflow"
top-level-decl: "__overflow"
top-level-decl: "main"

延々と続く top-level-decl: hogehoge の最後には top-level-decl: "main" と出力されているので、どうやら動いていそうです。

そもそも何が書かれているのか確認するために PrintFunctionNames を見に行くと、以下のような感じです。

https://github.com/llvm/llvm-project/blob/release/14.x/clang/examples/PrintFunctionNames/PrintFunctionNames.cpp

AST をなめ回しているっぽいですね。そして FrontendPluginRegistry::Add してる名前で呼び出せる、ふむふむ。なるほど、よく分かりました。さっさと次に行きましょう。

ところで AST どうなってんの

clang は特に何も無くとも AST を dump できます。

$ ./build/bin/clang -Xclang -ast-dump -fsyntax-only ../loop.c
TranslationUnitDecl 0x56149ab0d8d8 <<invalid sloc>> <invalid sloc>
|-TypedefDecl 0x56149ab0e100 <<invalid sloc>> <invalid sloc> implicit __int128_t '__int128'
| `-BuiltinType 0x56149ab0dea0 '__int128'
... (略) ...
`-FunctionDecl 0x55805ed95250 <../loop.c:3:1, line:9:1> line:3:5 main 'int (void)'
  `-CompoundStmt 0x55805ed95590 <col:16, line:9:1>
    |-ForStmt 0x55805ed95450 <line:4:4, line:6:4>
    | |-<<<NULL>>>
    | |-<<<NULL>>>
    | |-<<<NULL>>>
    | |-<<<NULL>>>
    | `-CompoundStmt 0x55805ed95438 <line:4:13, line:6:4>
    |   `-CallExpr 0x55805ed953e0 <line:5:8, col:52> 'int'
    |     |-ImplicitCastExpr 0x55805ed953c8 <col:8> 'int (*)(const char *, ...)' <FunctionToPointerDecay>
    |     | `-DeclRefExpr 0x55805ed952f0 <col:8> 'int (const char *, ...)' Function 0x55805ed84cc8 'printf' 'int (const char *, ...)'
    |     `-ImplicitCastExpr 0x55805ed95420 <col:15> 'const char *' <NoOp>
    |       `-ImplicitCastExpr 0x55805ed95408 <col:15> 'char *' <ArrayToPointerDecay>
    |         `-StringLiteral 0x55805ed95348 <col:15> 'char[35]' lvalue "\347\252\223\343\201\253\351\243\276\343\201\243\343\201\237\347\265\265\347\224\273\343\202\222\343\201\252\343\201\236\343\202\213\n"
    |-CallExpr 0x55805ed95508 <line:7:4, col:33> 'int'
    | |-ImplicitCastExpr 0x55805ed954f0 <col:4> 'int (*)(const char *, ...)' <FunctionToPointerDecay>
    | | `-DeclRefExpr 0x55805ed95488 <col:4> 'int (const char *, ...)' Function 0x55805ed84cc8 'printf' 'int (const char *, ...)'
    | `-ImplicitCastExpr 0x55805ed95548 <col:11> 'const char *' <NoOp>
    |   `-ImplicitCastExpr 0x55805ed95530 <col:11> 'char *' <ArrayToPointerDecay>
    |     `-StringLiteral 0x55805ed954a8 <col:11> 'char[20]' lvalue "\350\254\216\343\201\240\343\202\211\343\201\221\343\201\256\347\251\272\n"
    `-ReturnStmt 0x55805ed95580 <line:8:4, col:11>
      `-IntegerLiteral 0x55805ed95560 <col:11> 'int' 0

ふむふむ。構造がわかりますね。もういっちょいってみましょう。

loop2.c
int main(void) {
    for (int i = 0; i < 10; i++) {
        break;
    }
    return 0;
}

こんどは stdio.h を省いたので省略なしでお見せできます。

./build/bin/clang -Xclang -ast-dump -fsyntax-only ../loop2.c                                                                                Sun Aug 21 11:10:02 2022
TranslationUnitDecl 0x5586363555b8 <<invalid sloc>> <invalid sloc>
|-TypedefDecl 0x558636355de0 <<invalid sloc>> <invalid sloc> implicit __int128_t '__int128'
| `-BuiltinType 0x558636355b80 '__int128'
|-TypedefDecl 0x558636355e50 <<invalid sloc>> <invalid sloc> implicit __uint128_t 'unsigned __int128'
| `-BuiltinType 0x558636355ba0 'unsigned __int128'
|-TypedefDecl 0x558636356158 <<invalid sloc>> <invalid sloc> implicit __NSConstantString 'struct __NSConstantString_tag'
| `-RecordType 0x558636355f30 'struct __NSConstantString_tag'
|   `-Record 0x558636355ea8 '__NSConstantString_tag'
|-TypedefDecl 0x5586363561f0 <<invalid sloc>> <invalid sloc> implicit __builtin_ms_va_list 'char *'
| `-PointerType 0x5586363561b0 'char *'
|   `-BuiltinType 0x558636355660 'char'
|-TypedefDecl 0x5586363564e8 <<invalid sloc>> <invalid sloc> implicit __builtin_va_list 'struct __va_list_tag[1]'
| `-ConstantArrayType 0x558636356490 'struct __va_list_tag[1]' 1
|   `-RecordType 0x5586363562d0 'struct __va_list_tag'
|     `-Record 0x558636356248 '__va_list_tag'
`-FunctionDecl 0x5586363beb30 <../loop2.c:1:1, line:6:1> line:1:5 main 'int (void)'
  `-CompoundStmt 0x5586363bee30 <col:16, line:6:1>
    |-ForStmt 0x5586363bedc8 <line:2:5, line:4:5>
    | |-DeclStmt 0x5586363bece0 <line:2:10, col:19>
    | | `-VarDecl 0x5586363bec58 <col:10, col:18> col:14 used i 'int' cinit
    | |   `-IntegerLiteral 0x5586363becc0 <col:18> 'int' 0
    | |-<<<NULL>>>
    | |-BinaryOperator 0x5586363bed50 <col:21, col:25> 'int' '<'
    | | |-ImplicitCastExpr 0x5586363bed38 <col:21> 'int' <LValueToRValue>
    | | | `-DeclRefExpr 0x5586363becf8 <col:21> 'int' lvalue Var 0x5586363bec58 'i' 'int'
    | | `-IntegerLiteral 0x5586363bed18 <col:25> 'int' 10
    | |-UnaryOperator 0x5586363bed90 <col:29, col:30> 'int' postfix '++'
    | | `-DeclRefExpr 0x5586363bed70 <col:29> 'int' lvalue Var 0x5586363bec58 'i' 'int'
    | `-CompoundStmt 0x5586363bedb0 <col:34, line:4:5>
    |   `-BreakStmt 0x5586363beda8 <line:3:9>
    `-ReturnStmt 0x5586363bee20 <line:5:5, col:12>
      `-IntegerLiteral 0x5586363bee00 <col:12> 'int' 0

いろいろ細かいことはすっ飛ばすとして、以下のような構造を見かけたら:

  • FunctionDecl main
    • CompoundStmt
      • ...
      • ForStmt
        • CompoundStmt
          • ...

以下のように BreakStmt を付け足してやれば安全なループを書き換えられるのではないでしょうか？

  • FunctionDecl main
    • CompoundStmt
      • ...
      • ForStmt
        • CompoundStmt
          • ...
          • BreakStmt

完全理解。

プラグインを作る準備をする

今更諸々準備するのも面倒なので、llvm-project/clang/examples に相乗りすることにします。

llvm_project/clang/examples/LoopRewriter という名前でもろもろ作っていきます。

llvm-project/clang/examples/CMakeLists.txt
if(NOT CLANG_BUILD_EXAMPLES)
  set_property(DIRECTORY PROPERTY EXCLUDE_FROM_ALL ON)
  set(EXCLUDE_FROM_ALL ON)
endif()

if(CLANG_PLUGIN_SUPPORT)
  add_subdirectory(PrintFunctionNames)
  add_subdirectory(AnnotateFunctions)
  add_subdirectory(Attribute)
  add_subdirectory(CallSuperAttribute)
  add_subdirectory(PluginsOrder)
+  add_subdirectory(LoopRewriter)
endif()

基本的な構造は前節までで見てきた PrintFunctionNames を流用します。

AST を読む

さて、PrintFunctionNames をよくよく見ると、だいたい以下の流れで動いていることがわかります。

名前からも分かるとおり、RecursiveASTVisitor をなんとかしてやればよさそうです。
このあたりの処理には幸いにもドキュメントがあり、これを読みつつ多少試行錯誤してやると、TraverseDecl とか TraverseStmt を呼び出すと VisitXXX に流れていき、また、Decl や Stmt は dyn_cast でダウンキャストしたりしてなんやかんやできるということがわかります。

さて、今回書き換えたいのは ForStmt の中身なので、適当な RecursiveASTVisitor で VisitForStmt を実装してやります (Context はあとで必要になるので持ってきます)。

class ForStmtVisitor : public RecursiveASTVisitor<ForStmtVisitor> {
private:
  ASTContext &Context;

public:
  ForStmtVisitor(ASTContext &Context) : Context(Context) {}

  bool VisitForStmt(ForStmt *FS) {
    FS->dump();
    return true;
  }
};

あとはこいつを ASTConsumer で呼んでやれば ForStmt が dump できることになります。

class RewriteLoopConsumer : public ASTConsumer {
  CompilerInstance &Instance;
  ForStmtVisitor *visitor;

public:
  RewriteLoopConsumer(CompilerInstance &Instance)
      : Instance(Instance),
        visitor(new ForStmtVisitor(&Instance.getASTContext())) {}

  bool HandleTopLevelDecl(DeclGroupRef DG) override {
    for (DeclGroupRef::iterator i = DG.begin(), e = DG.end(); i != e; ++i) {
      Decl *D = *i;
      if (NamedDecl *ND = dyn_cast<NamedDecl>(D)) {
        if (ND->getNameAsString() == std::string("main")) {
          llvm::errs() << "inside main()\n";
          visitor->TraverseDecl(ND);
        }
      }
    }

    return true;
  }
};

ここまで来れば、あとはさっき出力した AST とにらめっこしながら気合いで読み解けます。
たぶん以下のような感じ:

  bool VisitForStmt(ForStmt *FS) {
    if (CompoundStmt *CS = dyn_cast<CompoundStmt>(FS->getBody())) {
      llvm::errs() << "inside for loop\n";
      for (CompoundStmt::body_iterator i = CS->body_begin(), e = CS->body_end();
           i != e; ++i) {
	Stmt *S = *i;
	S->dump();
      }
    }
    return true;
  }

cmake して loop.c を喰わせてやると、以下のような形で出力できます。

$ ./build/bin/clang -Xclang -load -Xclang ./build/lib/LoopRewriter.so -Xclang -add-plugin -Xclang loop-rewriter ../loop
.c -o loop.o
inside main()
inside for loop
CallExpr 0x55e4120bd680 'int'
|-ImplicitCastExpr 0x55e4120bd668 'int (*)(const char *, ...)' <FunctionToPointerDecay>
| `-DeclRefExpr 0x55e4120bd590 'int (const char *, ...)' Function 0x55e4120acaf8 'printf' 'int (const char *, ...)'
`-ImplicitCastExpr 0x55e4120bd6c0 'const char *' <NoOp>
  `-ImplicitCastExpr 0x55e4120bd6a8 'char *' <ArrayToPointerDecay>
    `-StringLiteral 0x55e4120bd5e8 'char[35]' lvalue "\347\252\223\343\201\253\351\243\276\343\201\243\343\201\237\347\265\265\347\224\273\343\202\222\343\201\252\343\201\236\343\202\213\n"

AST を (実力行使で) 書き換える

ここまで来たら、ForStmt の中の CompoundStmt までもう分かっている状態になりますが、今回我々は安全なループを、いや AST を書き換えたいわけなので、もう一手間加えましょう。

構造としては ForStmt -> CompoundStmt -> For の中身 となっているのですが、CompoundStmt には中身を付け足すためのインタフェースを持っていません。
仕方が無いので、以前の CompoundStmt の末尾に BreakStmt を追加した新しい CompoundStmt を作り、これを ForStmt の Body にセットするようにしてみましょう。幸い、ForStmt は setBody メソッドを持っています。

https://github.com/karno/llvm-project/blob/release/14.x/clang/include/clang/AST/Stmt.h#L2583

新しい CompoundStmt を作る

ところで、 CompoundStmt::Create のシグネチャは以下の通りです。
https://github.com/karno/llvm-project/blob/release/14.x/clang/include/clang/AST/Stmt.h#L1418-L1419

このうち、第1引数の const ASTContext &C は先ほど「あとで使う」と言っていたやつで、CompilerInstance::getASTContext() から取得できます。また、第3・4引数はこの CompoundStmt のソースコード上の位置を示すもので、これは元の CompoundStmt::getLBracLoc()、CompoundStmt::getRBracLoc() を拝借しましょう。

CompoundStmt *CSNew = CompoundStmt::Create(
    Context, /* Stmts: TBD */,
    CS->getLBracLoc(), CS->getRBracLoc());

新しい CompoundStmt の中身を作る

第二引数の ArrayRef<Stmt *> は割と何からでも作れるので、今回は単純に配列を作ることにします。

Stmt **stashedBody = new Stmt *[CS->size() + 1];
int stashIndex = 0;

// ... copy CompoundStmt body into stashedBody ...

// then, create new CompoundStmt()
CompoundStmt *CSNew = CompoundStmt::Create(
    Context, ArrayRef<Stmt *>(stashedBody, stashIndex),
    CS->getLBracLoc(), CS->getRBracLoc());

新しい CompoundStmt の中身に BreakStmt を付け足す

さて、残るは stashedBody に BreakStmt を付け足してやるだけです。BreakStmt 自体のシグネチャは大したことなく、ソースコード上の位置はデバッグ用途にしか使われないので適当な位置を補ってやれば OK です。

https://github.com/karno/llvm-project/blob/release/14.x/clang/include/clang/AST/Stmt.h#L2730-L2731

ところで、これをどうインスタンス化させるかですが、どうやら Placement New の引数として Context を渡してやる必要がありそう で、以下のようにすると上手く動きました。

new (Context) BreakStmt (CS->getEndLoc())

全部まとめるとこうなる

というわけで、 苦労して作ってやった CompoundStmt を ForStmt::setBody に渡してやることで無事 AST を書き換えることができます。

  bool VisitForStmt(ForStmt *FS) {
    // add statement into CompoundStmt
    if (CompoundStmt *CS = dyn_cast<CompoundStmt>(FS->getBody())) {
      llvm::errs() << "inside for loop\n";
      // create stash array
      Stmt **stashedBody = new Stmt *[CS->size() + 1];
      int stashIndex = 0;

      // iterate body and pack them into stash array
      for (CompoundStmt::body_iterator i = CS->body_begin(), e = CS->body_end();
           i != e; ++i) {
        Stmt *S = *i;
        stashedBody[stashIndex++] = S;
      }

      // and add BreakStmt
      stashedBody[stashIndex++] = new (Context) BreakStmt(CS->getEndLoc());

      // then, create new CompoundStmt()
      CompoundStmt *CSNew = CompoundStmt::Create(
          Context, ArrayRef<Stmt *>(stashedBody, stashIndex), CS->getLBracLoc(),
          CS->getRBracLoc());
      // and set for content
      FS->setBody(CSNew);

      // cleanup
      delete stashedBody;
    }
    return true;
  }

最後の一手間

さて、ここまで来たら書き換えられたも同然ですが、最後に大事な設定があります。

PluginASTAction には getActionType() というメソッドがあるのですが、これはデフォルトでは CmdlineAfterMainAction になっています。ここで言う MainAction はすなわちコード生成のことであり、コード生成の後にプラグインが実行されることとなります。

しかしながら、今回 AST をいじっているので、MainAction の前にプラグインを実行してくれない限り変更が反映されません。ということで、以下の通り設定します。

  ActionType getActionType() override { return AddBeforeMainAction; }

また、こうすることでこれまで長ったらしく (-Xclang -load -Xclang ./build/lib/LoopRewriter.so -Xclang -add-plugin -Xclang loop-rewriter) 設定していた引数がシンプルになり、-fplugin=./build/lib/LoopRewriter.so だけで済むようになります。

実際に書き換えてみよう

ということで、作ったプラグインの全体像です。

https://github.com/karno/llvm-project/blob/lrlr/clang/examples/LoopRewriter/LoopRewriter.cpp

このプラグインを使って、冒頭に掲出した loop.c をコンパイルしてみましょう。

loop.c
#include <stdio.h>

int main(void) {
   for (;;) {
       printf("窓に飾った絵画をなぞる\n");
   }
   printf("謎だらけの空\n");
   return 0;
}

$ cmake --build build -j8
$ ./build/bin/clang -fplugin=./build/lib/LoopRewriter.so ../loop.c -o loop.o
inside main()
inside for loop
$ ./loop.o
窓に飾った絵画をなぞる
謎だらけの空
$

無事に安全なループを書き換えられたようです。お疲れ様でした。

おわりに

思い立ってからなんとか形になるまで休日3日間を潰しましたが、なんとか書き換えることができてよかったです。こんなことをしているうちにリコリコ本編は大変なことになっていますが、気を強く持って生きていきたいと思います。いのちだいじに。

今回の話題は Web 上のドキュメントが本当に少なく、Stack Overflow で「modify AST clang plugin」で調べても「無理だよ」という回答を見かけることが多かったです。私自身 C++ は雰囲気でしか書けないので大変苦労しました。
しかしながら着実に知識を蓄積している先人も多く、たとえば以下のリソースは穴が空くほど見ました。ありがとうございました。

https://github.com/tpugh/mutation-lab-class-1/blob/a094cb8f654ff5cb01da0948e19b1a2da8e04a47/programs/vim/mull-cxx-frontend/src/ASTMutator.h

https://www.youtube.com/watch?v=_rUwW8Awc5s

脚注

  1. さユり 「花の塔」 より引用, https://www.uta-net.com/song/321038/ ↩︎

  2. C makes it easy to shoot yourself in the foot という言葉があるとおり ↩︎

  3. 諸説あり ↩︎

  4. 筆者は盛りすぎて BSOD したりした ↩︎

  5. 詳しい話はこのあたりにあります ↩︎

Discussion