FFRIエンジニアブログ

株式会社FFRIセキュリティのエンジニアが執筆する技術者向けブログです

LLVM バックエンド書き換えによる独自命令の挿入

はじめに

FFRIリサーチエンジニアの中川です。

以前投稿した記事「Shadow Stack を使った Stack Buffer Overflow 検知機構の実装」で

今回は LLVM バックエンドの書き換えにより、sspush sspop が関数のエピローグとプロローグに自動的に挿入されるようにしました。

LLVM バックエンドのカスタマイズ方法に関しては、記事が長くなることから、次の記事で紹介することにします。

のように書いていました。本記事では、以前の記事で後回しにしていた LLVM バックエンドのカスタマイズによる sspush sspop 命令の挿入方法について説明します。Shadow Stack を使った Stack Buffer Overflow 検知機構の実装の記事を未読の方はまずそちらをご参照ください。

LLVMでアセンブリが生成されるまでの流れ

最初にソースコードがコンパイルされターゲットアーキテクチャのアセンブリが生成されるまでの流れについて説明します。大雑把に以下の過程を経てソースコードがアセンブリへと変換されます。

  • 字句解析: ソースコードをトークンの並びに変換
  • 構文解析: 字句解析で作成したトークンの並びから抽象構文木を生成
  • 意味解析: 言語の記述仕様を満たしているのかの確認を実施し、中間表現・シンボルテーブルを生成
  • 最適化: 中間表現に対し、デッドコード削除・ピープホール最適化などの最適化を適用
  • コード生成: 中間表現からターゲットアーキテクチャ向けのアセンブリを出力

アセンブリが生成された後、アセンブラがアセンブリを機械語に変換し、リンカーがランタイムとのリンクを行い、最終的な実行ファイルが生成されます。

上に述べた 5 つの段階のうち、意味解析までをフロントエンドと呼び、最適化の部分をミドルエンド、コード生成の部分をバックエンドと呼びます。

LLVM はミドルエンドとバックエンドの部分の役割を担っています。フロントエンドの役割を担っているのは別のプログラムで、C/C++ 向けですと clang が使われています。

フロントエンドでは入力ソースコードの LLVM Intermidiate Representation (IR) への変換が実行されます。LLVM はこの IR に対して最適化を施し、ターゲットアーキテクチャ向けのアセンブリを出力します。 この言語やアーキテクチャから独立した IR に対して最適化を行うことが LLVM の大きな特徴の一つです。この特徴により、新しい言語向けのコンパイラを設計する際の労力を大幅に削減できます。ソースコードを IR に変換するフロントエンドさえ書いてしまえば、既存の最適化、ターゲットアーキテクチャ向けのバックエンドを使い回すことができるからです。

本記事では、コード生成に当たるバックエンド部分の書き換えを行います。

準備

まず、LLVM をビルドする環境を整えます。clang の公式サイトを参考にしながら進めていきます。

LLVM のソースコードをダウンロードします。

$ git clone https://github.com/llvm/llvm-project.git

また、crosstool-ng 経由で risc-v をターゲットとした gcc も導入しておきます (これはあとでリンカーとして利用)。その手順についてはこちらのブログ記事をご参照ください。

次に、LLVM をビルドします。マシンスペックにもよりますが、数時間程度かかると思ったほうが良いでしょう。また、並列ビルドする際に使うコア数を多くしすぎると、メモリが枯渇し死んでしまうことがあるので、そちらも注意です。以下にビルド方法の一例を示します。ビルドの高速化を図るため、ninja を使っています。

$ cd llvm-project
$ mkdir build
$ cd build
$ cmake -DLLVM_ENABLE_PROJECTS=clang -G Ninja ../llvm -DCMAKE_BUILD_TYPE="Release" \
  -DBUILD_SHARED_LIBS=True -DLLVM_USE_SPLIT_DWARF=True \
  -DLLVM_OPTIMIZED_TABLEGEN=True -DLLVM_BUILD_TESTS=False \
  -DDEFAULT_SYSROOT="/path/to/riscv64-unknown-elf" \
  -DLLVM_DEFAULT_TARGET_TRIPLE="riscv64-unknown-elf" \
  -DLLVM_TARGETS_TO_BUILD="RISCV"
$ cmake --build .

ひとまず、何もカスタマイズを加えていない状態で正しく動作することを確認しておきましょう。

$ cd ~
$ echo -e "#include <stdio.h>\nint main(){ puts(\"hoge\"); }" > test.c
$ ./llvm-project/build/bin/clang --target=riscv64 -c --sysroot=/etc/opt/riscv-toolchain/riscv64-unknown-elf/ test.c -march=rv64gc -mabi=lp64d
$ riscv64-unknown-elf-gcc -march=rv64gc -mabi=lp64d test.o -o a.out

コンパイルする際の注意点として marchmabi の 2 つを明示的に指定することが挙げられます。これを指定し忘れると ld: test.o: can't link soft-float modules with double-float modules というメッセージが表示されリンク時に失敗するので、必ず指定しましょう。

RISC-V ISA シミュレーター Spike で実行し、正しく実行できることを確認します。

$ spike pk a.out
bbl loader
hoge

コンパイルが正常に行われていることを確認できました。

sspush sspop 命令の挿入

次にバックエンド書き換え作業に移ります。

以前の記事で書いた、今回のバックエンド書き換えで実現したいことについて確認しておきます。

  1. 関数プロローグ直前に sspush 命令を挿入
  2. 関数エピローグ直後、return する直前に以下の 2 つの命令を挿入
    • sspop 命令
    • sspop 命令で取り出したリターンアドレスの値と、ra レジスタ (RISC-V では関数リターン直前に、リターンアドレスの値が ra レジスタに格納されている ) に格納されたリターンアドレスの値とを比較し、等しくない場合プログラムを終了する関数へジャンプする命令 (RISC-V だと bne 命令がこれに該当)

これからやることは以下の 2 つです。

  • 新規に定義した命令用の TableGen の作成
  • コード生成部分のうち、関数エピローグ生成部分 (RISCVFrameLowering::emitEpilogue) と関数プロローグ生成部分 (RISCVFrameLowering::emitPrologue) に、挿入したい命令を挿入する処理を記述

TableGen の作成

最初に TableGen について少しだけ補足説明しておきます。TableGen はターゲットアーキテクチャの詳細を記述するために用いられる Domain Specific Language (DSL) で、ターゲットアーキテクチャの命令・レジスタなどの情報を記述するのに用います。LLVM のビルドの過程で C++ のソースコードへと変換され、他の C++ のソースコードから、その情報を参照できるようになります。

今回、sspush sspop のエンコーディングは以下の表 1 のように設定しました。表 1 の情報に従って、sspush sspop 命令の情報を TableGen を使って記述していきます。

表 1 sspush sspop のエンコーディング。rs1 rs2 はソースレジスタ、rd はデスティネーションレジスタを表す。

まず sspush sspop 用に専用のオペコードを追加しておきます。llvm-project/llvm/lib/Target/RISCV/RISCVInstrFormats.td を開き、以下の行を追加します。

 def OPC_JALR      : RISCVOpcode<0b1100111>;
 def OPC_JAL       : RISCVOpcode<0b1101111>;
 def OPC_SYSTEM    : RISCVOpcode<0b1110011>;
 def OPC_CUSTOM0   : RISCVOpcode<0b0001011>; // <-- custom0 のオペコードを追加

次に sspush sspop 命令を追加します。llvm-project/llvm/lib/Target/RISCV/RISCVInstrInfo.td を開き、以下のように編集します。

 // SSPUSH 命令の定義を追加
 let hasSideEffects = 0, mayLoad = 0, mayStore = 0 in
 def SSPUSH : RVInstR<0b0000000, 0b111, OPC_CUSTOM0, (outs GPR:$rd), (ins GPR:$rs1, GPR:$rs2),
                    "sspush", "$rd, $rs1, $rs2">;

 // SSPOP 命令の定義を追加
 let hasSideEffects = 0, mayLoad = 0, mayStore = 0 in
 def SSPOP  : RVInstR<0b0000001, 0b111, OPC_CUSTOM0, (outs GPR:$rd), (ins GPR:$rs1, GPR:$rs2),
                    "sspop", "$rd, $rs1, $rs2">;

ここで、RVInstR というのを使っていますが、これは RISC-V の命令フォーマットのうち R 形式でエンコードすることを意味します。以下の TableGen のソースコードを見て想像のつく方もいるかと思いますが、funct7 (命令のうち 31--25 のビットのこと) の値, funct3 (命令のうち 14--12 のビットのこと) の値, opcode, destination register, source register1, source register2, 命令の文字列, オペランド の順に引数を設定して、命令オペコード、アセンブリ出力時の命令名を定義することができます。

class RVInstR<bits<7> funct7, bits<3> funct3, RISCVOpcode opcode, dag outs,
              dag ins, string opcodestr, string argstr>
    : RVInst<outs, ins, opcodestr, argstr, [], InstFormatR> {
  bits<5> rs2;
  bits<5> rs1;
  bits<5> rd;

  let Inst{31-25} = funct7;
  let Inst{24-20} = rs2;
  let Inst{19-15} = rs1;
  let Inst{14-12} = funct3;
  let Inst{11-7} = rd;
  let Opcode = opcode.Value;
}

関数エピローグとプロローグに命令の挿入

RISCVFrameLowering::emitEpilogueRISCVFrameLowering::emitPrologue の 2 つを書き換え、関数エピローグと関数プロローグに命令を挿入します。

関数プロローグ書き換えにあたっては、スタックフレーム生成の直前に sspush 命令を追加し、リターンアドレスの値を Shadow Stack にプッシュします。

void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
                                      MachineBasicBlock &MBB) const {

  ...

  // Early exit if there is no need to allocate on the stack
  if (StackSize == 0 && !MFI.adjustsStack())
    return;

  // 今回の追加部分 (スタックフレーム生成処理の直前に sspush x0, x1, x0 命令を追加)
  BuildMI(MBB, MBBI, DL, TII->get(RISCV::SSPUSH), RISCV::X0).addReg(RISCV::X1).addReg(RISCV::X0);

  // Allocate space on the stack if necessary.
  adjustReg(MBB, MBBI, DL, SPReg, SPReg, -StackSize, MachineInstr::FrameSetup);

  // Emit ".cfi_def_cfa_offset StackSize"
  unsigned CFIIndex = MF.addFrameInst(
      MCCFIInstruction::createDefCfaOffset(nullptr, -StackSize));
  BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
      .addCFIIndex(CFIIndex);

BuildMI という関数呼び出しを追加していますが、これでターゲットアーキテクチャ向けの命令を生成することができます。上記関数のコメントにも書いていますが、x1 レジスタ (ra レジスタの別名) の内容を Shadow Stack にプッシュする sspush x0, x1, x0 をスタックフレーム生成処理の直前に追加しています。

同じように関数エピローグには、Shadow Stack の内容をポップして、呼び出し元アドレスが格納されている ra の内容との比較処理を追加します。

void RISCVFrameLowering::emitEpilogue(MachineFunction &MF,
                                      MachineBasicBlock &MBB) const {
  ...

  // Deallocate stack
  adjustReg(MBB, MBBI, DL, SPReg, SPReg, StackSize, MachineInstr::FrameDestroy);

  // 今回の追加部分 (スタックフレームを破棄した後で sspop x5, x0, x0 命令の追加)
  BuildMI(MBB, MBBI, DL, TII->get(RISCV::SSPOP), RISCV::X5).addReg(RISCV::X0).addReg(RISCV::X0);
  // 今回の追加部分 (リターンアドレスと Shadow Stack からポップした内容の比較処理を追加 命令としては、bne x1, x5, stack_failure を追加している)
  BuildMI(MBB, MBBI, DL, TII->get(RISCV::BNE))
      .addReg(RISCV::X1)
      .addReg(RISCV::X5)
      .addExternalSymbol("stack_failure"); // stack_failure

  // After restoring $sp, we need to adjust CFA to $(sp + 0)
  // Emit ".cfi_def_cfa_offset 0"
  unsigned CFIIndex =
      MF.addFrameInst(MCCFIInstruction::createDefCfaOffset(nullptr, 0));
  BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
      .addCFIIndex(CFIIndex);
}

以上です。あとは、プロセスを終了させる stack_failure 関数が定義されたオブジェクトファイルを用意し、リンクさせれば完了です。

LLVM のビルドとバックエンド書き換え結果の確認

さて、LLVM をリビルドし、命令の挿入が意図した通りに実現できていることを確認します。

LLVM を再度ビルドします。

$ cd llvm-project/build
$ cmake --build .

以下のソースコードをコンパイルし、アセンブリを出力します。

C のソースコード

// test.c
#include <stdio.h>

__attribute__((noinline))
void func(const char* input) {
    puts(input);
}

int main() {
    func("Hello World");
}

コンパイル時のコマンド

$ ./llvm-project/build/bin/clang --target=riscv64 -c --sysroot=/etc/opt/riscv/riscv64-ctng/riscv64-unknown-elf test.c -march=rv64gc -mabi=lp64d -S

逆アセンブル結果

   .text
    .file    "test.c"
    .globl   func                    # -- Begin function func
    .p2align 1
    .type    func,@function
func:                                   # @func
# %bb.0:
    sspush   zero, ra, zero ; <-- sspush をスタックフレーム生成の直前に実行
    addi sp, sp, -32
    sd   ra, 24(sp)
    sd   s0, 16(sp)
    addi s0, sp, 32
    sd   a0, -24(s0)
    ld   a0, -24(s0)
    call  puts
    ld   s0, 16(sp)
    ld   ra, 24(sp)
    addi sp, sp, 32
    sspop    t0, zero, zero ; <-- 関数 return 直前に sspop を実行し、Shadow Stack のトップの値を t0 に代入
    bne  ra, t0, stack_failure ; <-- t0 の値と ra の値を比較し、違っていたら stack_failure へジャンプ
    ret
.Lfunc_end0:
    .size    func, .Lfunc_end0-func
                                        # -- End function
    .globl   main                    # -- Begin function main
    .p2align 1
    .type    main,@function
main:                                   # @main
# %bb.0:
    sspush   zero, ra, zero ; <-- sspush をスタックフレーム生成の直前に実行
    addi sp, sp, -16
    sd   ra, 8(sp)
    sd   s0, 0(sp)
    addi s0, sp, 16
    lui  a0, %hi(.L.str)
    addi a0, a0, %lo(.L.str)
    call  func
    mv   a0, zero
    ld   s0, 0(sp)
    ld   ra, 8(sp)
    addi sp, sp, 16
    sspop    t0, zero, zero ; <-- 関数 return 直前に sspop を実行し、Shadow Stack のトップの値を t0 に代入
    bne  ra, t0, stack_failure ; <-- t0 の値と ra の値を比較し、違っていたら stack_failure へジャンプ
    ret
.Lfunc_end1:
    .size    main, .Lfunc_end1-main
                                        # -- End function

上記のアセンブリにコメントとして挿入された命令について説明を加えています。

func 関数、main 関数それぞれについて、関数エピローグとプロローグに所望の命令が挿入されていることを確認できました。

まとめ

LLVM バックエンドの書き換えにより、関数エピローグとプロローグに独自の命令を挿入する方法について紹介しました。

RISC-V の大きな特徴として、CPU 設計者が独自の命令を定義できるというところがあります。そうして定義した命令を呼び出すにあたり、コンパイラの書き換えが必要になる場面があると思います。この記事が RISC-V やコンパイラをカスタマイズして使うことを考えている人の助けとなれば幸いです。

今回紹介しなかったですが、LLVM のカスタマイズとして、Pass を作成する方法もあります。Pass を使う方法であれば、LLVM 本体を書き換えず、LLVM IR に独自の処理を追加・あるいは編集することが可能になります。Pass を作成する方法についてはこちらの記事などをご参照ください。

追記

ちょうど、この記事を執筆している最中に lowRISC の twitter で、ターゲットアーキテクチャとしての RISC-V の指定が experimental ではなく official となったことがアナウンスされました。これにより、RISC-V ターゲットとした LLVM のビルドがデフォルトで行われるようになりました。嬉しいニュースです。

https://twitter.com/lowRISC/status/1175073534001012736