💣

「ゼロから学ぶRust 」正規表現エンジンにReDoS緩和策を実装する

2023/01/29に公開

はじめに

ゼロから学ぶRust システムプログラミングの基礎から線形型システム
https://www.kspub.co.jp/book/detail/5301951.html
では、6章で正規表現エンジンを実装します。

この記事はこの正規表現エンジンに
https://techlife.cookpad.com/entry/2022/12/12/162023
の手法を適用する話です。

ReDoS

「ゼロから学ぶRust」で実装する正規表現エンジンでは^などはないですが、以下の記事
https://qiita.com/prograti/items/9b54cf82a08302a5d2c7
で挙げられている

q(i+|t)+a

でReDoSが起こります。実際qiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiteとかをマッチングしてみるとかなり時間がかかるのが分かります。ベンチマークは後述します。
そこでキャッシュを用いて高速化を試みます。

キャッシュによるマッチングの高速化

方法は簡単で、「キャッシュによるRubyの正規表現のマッチングの高速化の紹介」にある通り、命令コードと位置をキャッシュします。またJump命令のみキャッシュしています。キャッシュにはHashSetを用います。

pub fn eval_with_cache(
    inst: &[Instruction],
    line: &[char],
    is_depth: bool,
) -> Result<bool, EvalError> {
    let mut cache = HashSet::new();
    if is_depth {
        eval_depth_with_cache(inst, line, 0, 0, &mut cache)
    } else {
        ...
    }
}

// 「ゼロから学ぶRust」のeval_depth関数が基本。差分だけ記載。
fn eval_depth_with_cache(
    inst: &[Instruction],
    line: &[char],
    mut pc: usize,
    mut sp: usize,
    cache: &mut HashSet<(usize, usize)>,
) -> Result<bool, EvalError> {
    ...
            Instruction::Jump(addr) => {
                if cache.contains(&(*addr, sp)) {
                    return Ok(false);
                }
                cache.insert((*addr, sp));
                pc = *addr;
            }
    ...

ベンチマーク

それではベンチマークをとってみます。ただしdo_matching_with_cachedo_matchingにおけるevaleval_with_cacheにしたものです。

use criterion::{criterion_group, criterion_main, Criterion};
use regex_zero::{do_matching, do_matching_with_cache};
use std::time::Duration;

const REDOS_INPUTS: &[(&str, &str, &str)] = &[
    ("i=22", "q(i+|t)+a", "qiiiiiiiiiiiiiiiiiiiiiite"),
    ("i=23", "q(i+|t)+a", "qiiiiiiiiiiiiiiiiiiiiiiite"),
    ("i=24", "q(i+|t)+a", "qiiiiiiiiiiiiiiiiiiiiiiiite"),
];

fn without_cache(c: &mut Criterion) {
    let mut g = c.benchmark_group("Without Cache");
    g.measurement_time(Duration::from_secs(80));

    for i in REDOS_INPUTS {
        g.bench_with_input(i.0, &(i.1, i.2), |b, args| {
            b.iter(|| do_matching(args.0, args.1, true))
        });
    }
}

fn with_cache(c: &mut Criterion) {
    let mut g = c.benchmark_group("With Cache");
    g.measurement_time(Duration::from_secs(80));

    for i in REDOS_INPUTS {
        g.bench_with_input(i.0, &(i.1, i.2), |b, args| {
            b.iter(|| do_matching_with_cache(args.0, args.1, true))
        });
    }
}

criterion_group!(benches, without_cache, with_cache);
criterion_main!(benches);

結果、自分の環境では

Without Cache/i=22      time:   [159.40 ms 159.94 ms 160.47 ms]
Without Cache/i=23      time:   [318.18 ms 319.59 ms 321.04 ms]
Without Cache/i=24      time:   [628.42 ms 632.70 ms 637.39 ms]
									   

With Cache/i=22         time:   [10.865 µs 10.939 µs 11.019 µs]
With Cache/i=23         time:   [11.749 µs 11.858 µs 11.980 µs]
With Cache/i=24         time:   [12.627 µs 12.809 µs 13.061 µs]

となり効果が確認できました。

おわりに

コードは
https://github.com/Catminusminus/regex-zero/tree/redos
にあります。

Discussion