static-keys

Static key是Linux内核中的一个底层机制，用于加速对很少改变的特性的条件判断检查。我们将这一特性迁移到了用户态Rust程序中，适用于Linux、macOS和Windows。（目前需要Nightly版本的Rust，原因请参见FAQ）

目前在CI中经过测试的支持平台包括：

• Linux

  • x86_64-unknown-linux-gnu
  • x86_64-unknown-linux-musl
  • i686-unknown-linux-gnu
  • aarch64-unknown-linux-gnu
  • riscv64gc-unknown-linux-gnu
  • loongarch64-unknown-linux-gnu

• macOS

  • aarch64-apple-darwin

• Windows

  • x86_64-pc-windows-msvc
  • i686-pc-windows-msvc

需要注意，如果使用cross-rs交叉编译loongarch64-unknown-linux-gnu平台，需要使用GitHub上的最新版cross-rs。更多细节可参见Evian-Zhang/static-keys#4。

更详细的解释和FAQ可参见GitHub Pages(English version).

出发点

现代程序往往可以通过命令行选项或配置文件进行配置。在配置中开启或关闭的选项往往在程序启动后是不会被改变的，但由在整个程序中被频繁访问。

let flag = CommandlineArgs::parse();
loop {
    if flag {
        do_something();
    }
    do_common_routines();
}

尽管flag在程序初始化后不会被修改，在每次循环中仍然会执行if判断。在x86-64架构中，这一判断一般会被编译为test-jnz指令

    test    eax, eax           ; 检查eax寄存器是否为0
    jnz     do_something       ; 如果不为0，则跳转do_something
do_common_routines:
    ; 执行相关指令
    ret
do_something:
    ; 执行相关指令
    jmp     do_common_routines ; 跳转至do_common_routines

尽管if判断只是test-jnz指令，这还可以进一步加速。我们是不是可以把这个判断优化为jmp指令（不执行do_something分支）或nop指令（总是执行do_something分支）？这就是static keys的原理。简单来说，我们需要动态修改指令。在从命令行获得flag之后，我们根据flag的值将if flag {}动态修改为jmp或nop指令。

例如，如果用户给出的flag为false，那么生成的指令将被动态修改为下面的nop指令

    nop     DWORD PTR [rax+rax*1+0x0]
do_common_routines:
    ; 执行相关指令
    ret
do_something:
    ; 执行相关指令
    jmp     do_common_routines

如果flag为true，那么生成的指令将被修改为下面的无条件跳转指令

    jmp     do_something
do_common_routines:
    ; 执行相关指令
    ret
do_something:
    ; 执行相关指令
    jmp     do_common_routines

这里的指令就不会再有test和条件跳转指令了，而仅仅是一个nop（即不做任何事）或jmp指令。

尽管将test-jnz替换为nop提升的性能可能微乎其微，但是Linux内核文档中描述了，如果这样的配置检查涉及全局变量，那么这样的替换可以减小缓存压力。并且在服务端程序中，这些配置有可能会通过Arc在多线程中共享，那么用nop或jmp就可以有更大的提升空间。

使用方式

目前需要nightly版本的Rust来使用这个库。在使用者的代码中，需要声明对unstable特性asm_goto的使用。

#![allow(unused)]
#![feature(asm_goto)]
fn main() {
}

首先，在Cargo.toml中加入相应依赖：

[dependencies]
static-keys = "0.6"

在main函数开头，需要调用static_keys::global_init进行初始化。

fn main() {
    static_keys::global_init();
    // 执行其他指令
}

随后需要定义一个static key来记录用户传入的flag，并根据用户传入的值来控制这个static key的值。

#![allow(unused)]
fn main() {
use static_keys::define_static_key_false;
struct CommandlineArgs {}
impl CommandlineArgs { fn parse() -> bool { true } }
// FLAG_STATIC_KEY初始值为`false`
define_static_key_false!(FLAG_STATIC_KEY);

fn application_initialize() {
    let flag = CommandlineArgs::parse();
    if flag {
        unsafe {
            FLAG_STATIC_KEY.enable();
        }
    }
}
}

同一个static key可以在任意时刻多次更改值。但需要注意的是，在多线程环境中修改static key是非常危险的。因此，如果需要使用多线程，请在完成对static key的修改后再创建新线程。不过，在多线程环境中使用static key是绝对安全的。此外，对static key的修改相对比较慢，但由于static key一般用于控制很少被修改的特性，所以这样的修改相对比较少，因此慢点也没有太大影响。请参见FAQ了解更多。

在定义static key之后，就可以像平常一样用if语句来使用这个static key了（这个和这个介绍了likely和unlikely API的语义）。同一个static key可以在多个if语句中被使用。当这个static key被修改时，所有使用这个static key的if语句都将被统一修改为jmp或nop。

#![allow(unused)]
fn main() {
#![feature(asm_goto)]
use static_keys::{define_static_key_false, static_branch_unlikely};
struct CommandlineArgs {}
impl CommandlineArgs { fn parse() -> bool { true } }
fn do_something() {}
fn do_common_routines() {}
define_static_key_false!(FLAG_STATIC_KEY);
fn run() {
    loop {
        if static_branch_unlikely!(FLAG_STATIC_KEY) {
            do_something();
        }
        do_common_routines();
    }
}
}

参考链接

• Linux内核官方文档：static-keys
• Linux static_key internals
• Rust for Linux项目也实现了static key，请查看Rust-for-Linux/linux#1084。我们基于break的内联汇编实现是受到这一实现的启发的。

底层实现

正如在简介中所说，static key的使用流程如下：

1. 全局初始化相应结构。
2. 定义一个static key。
3. 根据用户传入的值修改static key。
4. 在if判断处使用static key。

在本节中，我们使用如下术语：

• Static key

  静态变量，用于存储相关信息来控制static branch的选取。

• Jump entry

  静态变量，用于存储static branch的信息。这个变量用于定位static branch。

• Static branch

  使用static key的if判断。

简化逻辑

在简化的逻辑中，我们可以把static key和jump entry理解为以下结构体：

struct StaticKey {
    enabled: bool,
    jump_entries: Vec<JumpEntry>,
}

struct JumpEntry {
    code: &'static Instruction,
}

当我们修改static key时，会进行以下的步骤：

1. 修改static key的enable字段
2. 根据jump_entries字段去找所有与这个static key相关联的jump entry
3. 对于每个jump entry，根据其static_branch字段来定位static branch
4. 根据enable字段的值来修改static branch为nop或jmp

我们在if判断处使用static key，就是在static key的entries字段增加一个元素，记录当前if判断的位置。

在理解了简化逻辑之后，我们还需要以下补充：

• jump entry的位置
• static branch的修改内容
• static branch的修改准则
• static branch的修改方式

jump entry的位置

根据之前介绍的使用方式，我们可以在多处if判断中使用同一个static key。因此，一个static key可能会与多个jump entry相关联。但是，我们不能分布式地创建一个编译期的vector：我们无法定义一个静态vector之后，在各处代码中在编译期往这个vector里加入元素。因此，我们必须将jump entry存储在生成的二进制文件中，在运行时把jump entry与static key相关联。

具体来说，我们将jump entry存储在生成的二进制文件的特定节中。这个节的名称在不同操作系统中不同。例如，在Linux ELF中，我们将这个节称为__static_keys。

在运行时的初始化阶段，我们会收集每个static key关联的jump entry。但是，由于jump entry都位于一个已经载入内存的节中，所以如果再把这个jump entry加入static key的vector中，那么内存占用就会翻倍，这是我们不想看到的。

为了解决这个问题，我们把jump_entries字段定义为一个指针而非一个vector。这个指针可以直接指向相应节中的jump entry，因此可以减少内存占用。为了做到这样，我们需要对这个节中的jump entry进行排序，来确保相同的static key的jump entry应该相邻，这样的话jump_entries字段可以指向static key关联的第一个jump entry。

为了能够进行排序，我们需要在JumpEntry中再加入一个字段：static key的地址。这样的话，我们就可以根据这个地址对jump entry进行排序。需要注意到的一点是，在实现中，考虑到ASLR，这些地址都是相对地址。

因此，相应的结构体需要修改为

struct StaticKey {
    enabled: bool,
    jump_entries: *const JumpEntry,
}

struct JumpEntry {
    code: &'static Instruction,
    /// static key的相对地址
    key: usize,
}

jump_entries字段在一开始是null，然后在初始化时，这个字段被更新为指向其关联的第一个jump entry的指针。

static branch的修改内容

当修改static branch时，我们需要把nop修改为jmp或把jmp修改为nop。在大多数指令集架构中，nop指令的长度可以有多种。例如，在x86-64架构中，由于jmp一般5字节，所以我们选择了5字节长度的nop来替换相应指令。这样可以保证我们不会污染相邻的指令。

但是，在修改nop为jmp时，我们该如何填jmp的目的地址？这并不能被很简单地计算出来。因此，我们需要在JumpEntry结构体中增加另一个字段来记录跳转的目的地址：

#![allow(unused)]
fn main() {
struct JumpEntry {
    code: &'static Instruction,
    /// 跳转目标的相对地址
    target: usize,
    key: usize,
}
}

为了在static branch处生成相应的jump entry，我们使用了如下的内联汇编（以x86-64为例）。当使用static_branch_likely!和static_branch_unlikely!宏时，会展开为如下代码片段（具体细节可能会有所差别）：

'my_label {
    ::core::arch::asm!(
        r#"
        2:
        .byte 0x0f,0x1f,0x44,0x00,0x00
        .pushsection __static_keys, "awR"
        .balign 8
        .quad 2b - .
        .quad {0} - .
        .quad {1} + {2} - .
        .popsection
        "#
        label {
            break 'my_label false;
        },
        sym MY_STATIC_KEY,
        const true as usize,
    );
    break 'my_label true;
}

这看上去非常复杂，我们来分段讲解。

汇编片段

第一行2:代表一个汇编标签，用于表示当前0x0f, 0x1f, 0x44, 0x00, 0x00这个数据的地址。这5个字节构成了一个nop指令。

然后我们使用一对.pushsection和.popsection来切换至另一个节（当前节为.text，用于记录指令），用于记录jump entry。

在新的节中，我们使用三个.quad，定义了三个8字节的值。这三个值分别对应为JumpEntry结构体的三个字段。第一个8字节值是2b - .，这里面2b代表与当前最近的2标签，也就是刚刚定义的nop指令的地址。而.代表当前的位置，也就是现在的8字节的值的地址。因此，2b - .就代表了一个与nop的相对地址，也就是JumpEntry的code字段。

第二个8字节值是{0} - .。这里{0}代表内联汇编的第一个参数，也就是label { break 'my_label false; }。这就是jmp指令的目的地址，也就对应于JumpEntry的target字段。这将在后面更详细地解释。

第三个8字节值是{1} + {2} - .，存储了static key的相应信息以及其初始值（需要注意的是，由于static key总是8字节对齐，因此其地址的最后一个字节总是0x00，所以我们就可以用这个字节去记录额外信息）。这个初始值我们也将在之后详细解释。

通过执行这个内联汇编，在"__static_keys"节就可以在编译期生成一个jump entry。

跳转标签部分

由于这里的内联汇编并不影响控制流，所以我们把上面的代码片段简化一下，只看其跳转标签部分：

'my_label {
    // 内联汇编
    break 'my_label true;
}

这段代码会被Rust编译器理解为一个true值。由于我们是在if判断处使用这些宏，所以if判断就变成了

if true {
    do_a();
} else {
    do_b();
}
do_c();

因此，Rust编译器就会将这些指令优化为

nop        ; 0x0f,0x1f,0x44,0x00,0x00
call do_a  ; do_a()

但是，do_b()并不会被优化掉：内联汇编的参数中用到了这个分支---label { break 'my_label false; }。正如前面所说，这个参数代表了break 'label false;语句的地址。当把这个语句放在if判断中时，就变成了一个false条件。因此，这个语句就会被编译为一个对do_b()的调用，而这个调用在静态控制流中是永远不会被执行的。为了理解得更清晰，我们来看看生成的汇编：

    nop           ; 0x0f,0x1f,0x44,0x00,0x00
    call    do_a  ; do_a()
DO_C:
    call    do_c  ; do_c()
    ret           ; 函数结尾
DO_B:
    call    do_b  ; do_b()
    jmp     DO_C  ; goto DO_C

DO_B处的基本块在静态控制流中永远不会被执行，而我们把它的地址存储在了jump entry中。

当我们把这个static branch修改为jmp，汇编代码就变成了

    jmp     DO_B  ; 此处被修改
    call    do_a  ; do_a()
DO_C:
    call    do_c  ; do_c()
    ret           ; 函数结尾
DO_B:
    call    do_b  ; do_b()
    jmp     DO_C  ; goto DO_C

一切符合预期。

static branch的修改准则

分支布局

正如之前所介绍的，有两个分支会被执行：一个分支在nop后被执行，与主要的部分相邻。另一个分支需要通过两个额外的jmp执行，它的位置一般在函数的结尾。一般来说，不太可能被执行到的分支应当被放在后者，而前者则应当是更有可能被执行到的分支。这样的分支布局是通过static_branch_likely!和static_branch_unlikely!来控制的。

在使用static_branch_likely!时，更有可能被执行到的分支会放在true分支，也就是紧邻着主要部分，在执行完nop后就被执行。而false分支则被放在了其他位置，需要通过两个jmp来执行。

在内联汇编中，这两种布局的差别是通过结尾的break 'my_label true和break 'my_label false来体现的。

初始指令

在得到正确的分支布局之后，下一个问题就是，我们在生成可执行程序时，对应的默认的初始指令该是nop还是jmp呢？如果在程序启动后，我们不修改static key，那么这个默认的初始指令就会被执行。所以，正确的方案是：

• 对于 static_branch_likely!

  • 如果static key的初始值为true，则生成nop
  • 如果static key的初始值为false，则生成jmp

• 对于static_branch_unlikely!

  • 如果static key的初始值为false，则生成nop
  • 如果static key的初始值为true，则生成jmp

修改的方向

另一个问题是，当修改static key时，我们应该更新至哪个指令？我们应当把jmp更新为nop，还是把nop更新为jmp？为了解决这个问题，我们应当使用JumpEntry结构体key字段的最后一字节记录的初始状态。

初始状态是布尔值，表示这个if判断中，更有可能的分支是否是true分支。正如上面所说，更有可能执行到的分支应当总是与主要部分相邻。而决定这个布尔值的是我们调用的是static_branch_likely!还是static_branch_unlikely!，以及static key的初始值/

当我们修改static branch时，我们可以通过将static key的新值与jump entry中记录的初始状态异或得到。例如，如果更有可能执行到的分支是true分支，并且static key的新值是true，那么我们应当将jmp更新为nop。这是因为我们需要执行紧挨着static branch判断的分支。

static branch的修改方式

最后一个需要被解决的问题是如何修改static branch。

众所周知，程序的指令都位于text节。在大多数平台上，text节拥有执行权限而没有写权限，从而阻止攻击者修改程序指令执行恶意逻辑。这种保护一般被称为DEP或者W^X。

为了修改static branch处的指令，我们就需要暂时绕过DEP一会儿。这种绕过虽然是危险的，但是只会发生在static key的修改过程中。一旦结束修改，我们就会恢复相应的保护。因此，在修改static key时需要额外注意。

FAQs

为什么static key应当被用在较少改变的特性上？

两点原因：

• 对static key的修改需要绕过DEP，会带来潜在的安全风险。不过DEP会在修改结束后重新生效。
• 对static key的修改比较慢，因为涉及到了许多系统调用。

为什么static key必须在单线程环境下修改？

在用户态，如果要修改别的线程可能会执行到的指令会非常复杂。Linux内核社区曾经提出过text_poke系统调用，但是如今仍不可用。顺带一提，Linus好像不太喜欢这个，并且他说的很有道理。

另一个原因是我们需要操作内存保护权限来绕过DEP，但是在多线程环境下，这会引发保护权限本身的race condition。尽管可以用mutex来解决数据竞争的问题，但是如果cargo解析出多版本的static-keys依赖，那么每个版本中都会有一个全局mutex实例，这种方法就失效了。这个可以被RFC 1977: public & private dependencies解决。rust-lang/cargo#2363亦可供参考。

为什么需要nightly Rust？

我们在内部使用了内联汇编，并且使用了asm_goto和asm_const这两个特性。只要这两个特性稳定了，我们就能使用stable Rust了。

为什么static_branch_likely!和static_branch_unlikely!是宏？

因为内联汇编的sym参数需要是静态路径，这在函数里是做不到的。

如果要扩展到新的操作系统，需要实现哪些操作系统特性？

• 标志一个自定义节的开始和结束的符号

  这是用来对static key排序，以及标志循环结束

• 保证自定义节不会被链接器的GC回收的方案

• 绕过DEP的方案

  用于更新static branch的指令

如果要扩展到新的指令集架构，需要实现哪些架构特性？

• 与jmp等长的nop指令（或者可以整除，如2字节nop与4字节jmp）
• 在Linux上清除指令缓存的方式（需加入到Evian-Zhang/clear-cache）
• Rust支持的内联汇编

我可以在no_std环境中使用吗？

可以