キャッシュによるRubyの正規表現のマッチングの高速化の紹介

9月からRuby開発チームにインターンシップとして参加している@makenowjustです。 総合研究大学院大学の学生で、普段は情報セキュリティに関する研究をしています。

インターンシップでは、キャッシュ (メモ化) を利用したRubyの正規表現の高速化を行いました。 ReDoSと呼ばれる、バックトラックが爆発することでマッチング時間が膨大になる脆弱性があります (ReDoSについては、拙作ですがWEB+DB PRESSに掲載された記事があります)。 近年、ReDoSは多く報告されており、Rubyもその例外ではありません (参考1、参考2)。 今回実装した最適化は、ReDoSを防ぐことを目的としたもので、多くの正規表現のマッチング時間が文字列の長さに対して線形となります。

ReDoSが起こる正規表現の例として、/^(a|a)*$/が挙げられます。 今回の修正の前後での実行時間を比較すると、次のようになります。

$ ruby --version
ruby 3.2.0preview2 (2022-09-09 master 35cfc9a3bb) [arm64-darwin21]

$ time ruby -e '/^(a|a)*$/ =~ "a" * 28 + "b"'
ruby -e '/^(a|a)*$/ =~ "a" * 28 + "b"'  8.49s user 0.04s system 98% cpu 8.663 total

$ ./miniruby --version
ruby 3.2.0dev (2022-10-27T13:39:56Z recache bc59b7cc1e) [arm64-darwin21]

$ time ./miniruby -e '/^(a|a)*$/ =~ "a" * 28 + "b"'
./miniruby -e '/^(a|a)*$/ =~ "a" * 28 + "b"'  0.01s user 0.01s system 8% cpu 0.310 total

修正前は実行に8秒かかっていた処理が、修正後は0.01秒で終了しており、800倍高速になっていることが分かります。

この最適化を含むパッチは既にマージされ、Ruby 3.2.0 Preview 3に含まれています。 そのため、すぐに手元で試すことができます。

本記事では、今回実装した最適化の実装や注意点について解説をしたいと思います。

正規表現とNFA

情報科学において、正規表現はNFA (非決定性有限状態オートマトン) と等価な表現力を持っていることが知られています。 そして、Rubyなどの正規表現マッチングは、バックトラックを用いてNFAの動作を模倣することで実装されています。

例えば、正規表現 /(a|a)*/ に対応するNFAは次の図のようになります。

/(a|a)*/ に対応するNFA

このNFAは状態が初期状態かつ受理状態で、文字で自身に戻ってくる遷移が2つあります。 そのため、aaaaab のような文字 a のあとに a 以外の文字に対してマッチングを行うと、バックトラックによって状態に到達する回数がaの個数の指数回となります。 このようにしてバックトラック回数が爆発することが、ReDoSの原因となっています。

このことをバックトラックによるマッチングの実装を通して確かめてみましょう。 実装は次のようになります。

# nfa.rb

# デバッグ出力を有効にするかどうか
$debug = false

# NFA (非決定性有限状態オートマトン) を表現するクラス。
class NFA
  # @param [Symbol] init 初期状態
  # @param [Array] accepts 受理状態
  # @param [Hash] transition 遷移関数
  def initialize(init, accepts, transition)
    @init = init
    @accepts = accepts
    @transition = transition
  end

  attr_reader :init, :accepts, :transition

  # バックトラックを用いてNFAのマッチングを行う。
  #
  # @param [String] input 入力文字列
  # @param [Symbol] state 現在の状態
  # @param [Integer] n 何文字目か
  def run_backtrack(input, state = init, n = 0)
    if $debug
      puts "run_backtrack(#{input.inspect}, #{state.inspect}, #{n.inspect})"
    end

    if input.size == n
      return accepts.include?(state)
    end

    c = input[n]
    (transition.dig(state, c) || []).each do |next_state|
      if run_backtrack(input, next_state, n + 1)
        return true
      end
    end

    false
  end
end

この実装でマッチングを行ってみます。

irb(main):001:0> require_relative './nfa.rb'
=> true
irb(main):002:0> # 正規表現 `(a|a)*` に対応するNFA
irb(main):003:0> nfa = NFA.new(
irb(main):004:1*   :q1,
irb(main):005:1*   [:q1],
irb(main):006:1*   {
irb(main):007:2*     q1: { 'a' => %i[q1 q1] }
irb(main):008:2>   }
irb(main):009:1> )
=> #<NFA:0x00000001520cac08 @init=:q1, @accepts=[:q1], @transition={:q1=>{"a"=>[:q1, :q1]}}>
irb(main):010:0> $debug = true
=> true
irb(main):011:0> nfa.run_backtrack("a" * 2 + "b")
run_backtrack("aab", :q1, 0)
run_backtrack("aab", :q1, 1)
run_backtrack("aab", :q1, 2)
run_backtrack("aab", :q1, 2)
run_backtrack("aab", :q1, 1)
run_backtrack("aab", :q1, 2)
run_backtrack("aab", :q1, 2)
=> false
irb(main):012:0> nfa.run_backtrack("a" * 3 + "b")
run_backtrack("aaab", :q1, 0)
run_backtrack("aaab", :q1, 1)
run_backtrack("aaab", :q1, 2)
run_backtrack("aaab", :q1, 3)
run_backtrack("aaab", :q1, 3)
run_backtrack("aaab", :q1, 2)
run_backtrack("aaab", :q1, 3)
run_backtrack("aaab", :q1, 3)
run_backtrack("aaab", :q1, 1)
run_backtrack("aaab", :q1, 2)
run_backtrack("aaab", :q1, 3)
run_backtrack("aaab", :q1, 3)
run_backtrack("aaab", :q1, 2)
run_backtrack("aaab", :q1, 3)
run_backtrack("aaab", :q1, 3)
=> false

このように、"a" * 2 + "b"に対するマッチでは出力が7行、"a" * 3 + "b"では15行と、倍々に増えていきます。 "a" * 30 + "b" のような入力にすると、延々と出力が続いて終了しなくなるので注意してください。

また、バックトラックによって同じ状態の同じ位置に何度も訪れていることにも注目してください。 バックトラックの爆発が起こる場合、このような現象が見られます。

キャッシュ (メモ化) の導入

先程の実装で、バックトラックの爆発が起こる際には、同じ状態で同じ位置に何度も訪れていることが分かりました。 しかし、同じ状態で同じ位置に到達した場合のNFAの動作は全く同じになるので、すでにマッチに失敗すると分かっていれば、わざわざ実行する必要はありません。 そこで、状態と位置をキーにして、そこに到達したかどうかを記録しておくことで、余計なバックトラックを防ぐことができます。 これが、キャッシュ (メモ化) による正規表現マッチングの最適化の原理となります。

キャッシュによる最適化の実装は次のようになります。

# cache.rb

class NFA
  # キャッシュによって最適化されたNFAのマッチングを行う。
  #
  # @param [String] input 入力文字列
  # @param [Symbol] state 現在の状態
  # @param [Integer] n 何文字目か
  # @param [Hash] cache キャッシュ
  def run_cached(input, state = init, n = 0, cache = {})
    # ここでキャッシュの確認と追加を行う。
    return false if cache[[state, n]]
    cache[[state, n]] = true

    if $debug
      puts "run_cached(#{input.inspect}, #{state.inspect}, #{n.inspect})"
    end

    if input.size == n
      return accepts.include?(state)
    end

    c = input[n]
    (transition.dig(state, c) || []).each do |next_state|
      if run_cached(input, next_state, n + 1, cache)
        return true
      end
    end

    false
  end
end

この実装を試してみます。 (先程のirbセッションの続きになっています)

irb(main):013:0> require_relative "./cache.rb"
=> true
irb(main):014:0> nfa.run_cached("a" * 2 + "b")
run_cached("aab", :q1, 0)
run_cached("aab", :q1, 1)
run_cached("aab", :q1, 2)
=> false
irb(main):015:0> nfa.run_cached("a" * 3 + "b")
run_cached("aaab", :q1, 0)
run_cached("aaab", :q1, 1)
run_cached("aaab", :q1, 2)
run_cached("aaab", :q1, 3)
=> false

今度は、"a" * 2 + "b"で3行、"a" * 3 + "b"で4行と、1行しか増えていません。 さらに"a"の個数を増やしても、4個だと5行、5個だと6行と、1行ずつしか増えていかず、先程のような指数的な爆発が起こらないことが確認できます。

たった2行の追加で、ここまで劇的に高速化できるのは驚きなのではないでしょうか？

実際の実装

実際のRubyの正規表現実装 (Onigmo) への修正は、次のPull Requestで行いました。

Implement cache optimization for regexp matching by makenowjust · Pull Request #6486 · ruby/ruby · GitHub

また、次のチケットで進捗を管理しています。

Feature #19104: Introduce the cache-based optimization for Regexp matching - Ruby master - Ruby Issue Tracking System

OnigmoはVM型の正規表現実装のため、状態と位置の組でキャッシュするのではなく、命令コードと位置でキャッシュを行います。

Pull Requestを見ると分かりますが、差分は800行弱と、想像よりも少ない差分になっています。 このように、キャッシュによる最適化は、元のコードを大きく変更することなく、正規表現のマッチングを高速化できることが、利点の一つだと考えています。

実装の工夫として、すべての命令コードに対してキャッシュをするのではなく、分岐の命令コードでのみキャッシュをするようにしています。 複数回訪れることがあるのは分岐の合流地点なので、本当はそちらでキャッシュした方がいいのですが、合流地点を求めるには分岐命令のジャンプ先などを全て記録する必要があり大変なため、現在は分岐命令でキャッシュを行っています。

これらの工夫や、キャッシュによる最適化というアイディアは、次の論文を参考にしました。

Davis, James C., Francisco Servant, and Dongyoon Lee. "Using selective memoization to defeat regular expression denial of service (ReDoS)." 2021 IEEE symposium on security and privacy (SP). IEEE, 2021. https://www.computer.org/csdl/proceedings-article/sp/2021/893400a543/1oak988ThvO

キャッシュ (メモ化) の注意点

今回のキャッシュ (メモ化) による最適化によって、マッチングの時間が線形になることが保証できました。 しかし、この実装にはいくつか注意するべき点があるので、それらについて説明します。

いくつかの正規表現は最適化できない

実装を見て気付いた方もいるかもしれませんが、キャッシュによる最適化はいくつかの正規表現の拡張や場合に対応していません。 具体的には、次のときに最適化が無効になるか、エラーになります。

1. 後方参照 \1、部分式の呼び出し \g<foo>、先読みや後読み (?<=...)、アトミックグループ (?>)、非包含オペレータ (?~...) を使っている正規表現の場合
2. 回数指定の繰り返しが繰り返しの中にある場合 (例 (fo{2})+)

2の場合に最適化できないのは完全に実装の都合なのですが、Onigmoの実装に多くの修正が必要になり、パッチの妥当性の検証が困難になることが予想されるため、実装を見送りました。

メモリを多く使う

キャッシュにはビット配列を用いているので、((キャッシュを行う命令コードの数) * (文字列の長さ)) / 8 バイト程度のメモリをマッチの際に追加で使用します。 そのため、極端に大きな文字列でマッチングを行なう場合や、極端にたくさんの分岐のある正規表現を利用している場合に、メモリ使用量が急に増加して、問題になるかもしれません。 とくに、大きな回数指定の繰り返しを使っている場合は、その繰り返し回数だけ分岐の数が増えることになるので、問題になりやすいです。

ただし、(現在の実装では) バックトラックの回数が文字列の長さを越えた場合にのみ、キャッシュによる最適化が有効になります。 キャッシュ用のメモリの確保もこのときに行うため、元々の正規表現に大きな問題が無ければ、メモリ使用量で問題になることも少ないのではないかと考えています。

また、チケットの方にも書いた通り、実際に使われている正規表現でも、分岐の数は多くても80個ほどで、文字列の長さに対して10倍メモリが使用される程度に収まることを確認しています。

まとめ

この記事では、Rubyに新しく導入される、キャッシュによる正規表現のマッチングの最適化について紹介しました。 最後に、キャッシュによる最適化の特徴や注意点を整理します。

• この最適化によって、正規表現のマッチングの時間が、文字列の長さに比例する程度に抑えられるようになる。
• バックトラックの回数が文字列の長さを越えたときに、最適化が有効になる。
• 一部の正規表現の拡張を使っている場合や、回数指定の繰り返しがネストしている場合に最適化が有効にならないことがある。
• マッチングに使うメモリの量がこれまでよりも増える可能性がある。

この最適化がReDoSを根絶することができるわけではありませんが、多くの場合に有効な対策となっていることを期待しています。