LitJWTに見るモダンなC#のbyte[]とSpan操作法

2019-05-27

LitJWT、という超高速な認証ライブラリを作りました。

github.com/Cysharp/LitJWT

なんと今回はUnity用が、ない！どころか.NET Standardですら、ない！.NET Core専用になってます（今のとこ）。理由はパフォーマンス都合で現状.NET CoreにしかないAPIを使いすぎたので修正が面倒ということなので、そのうちなんとかするかもしれませんかもしれません。

5倍高速

そもそも認証ライブラリでパフォーマンス追求しているものなんてない！ので！まぁそりゃそうだという感じではある。実際、そこまで認証で必要か？というと疑問符が付くところなので、ただのオーバーエンジニアリングなのですが、とはいえ速いというのは良いことです。シンプルに。

JWTに関しては特に説明することもないので（セッションにでも何にでも使えばいいんじゃないかしら、実際、私はMagicOnionのステートレスセッションのために必要なので用意しました）、ここから先は実装の話をします。

モダンBase64(Url)

JWTは大雑把にはJSONをBase64Urlでエンコードして署名を引っ付けたもの、です。Base64UrlというのはBase64の亜種で、URLセーフになるように、使う文字列が少し異なります。性質上、GETでURLにトークンが引っ付いたりするかもですしね。なるほど。

さて、しかしそんなマイナーなBase64Urlをエンコードするメソッドは用意されていないので、普通はこんな風に書いてます。

Convert.ToBase64String(input)
    .TrimEnd('=')      // 新しいstringを作る
    .Replace('+', '-') // 新しいstringを作る
    .Replace('/', '_') // 新しいstringを作る

改めてBase64Urlは、ようするにパディング（4の倍数に収まらない場合に末尾につく)の=が不要で、+が-、/が_なBase64なので、置換！ただたんに置換！する、すなわち新規文字列を無駄に作成！無駄に検索して無駄に作成！なわけです。

実際、別にこのBase64の変換表の一部を差し替えるだけの話なのに。

無駄すぎて発狂しちゃうので、ここは普通に自前でBase64を実装することで大解決しましょう。実際それしか方法はない、しょうがない。

せっかく作るので、今風のAPIにしましょう。例えばデコードのAPIはこんな感じに。

public static bool TryFromBase64UrlString(string s, Span<byte> bytes, out int bytesWritten)
public static bool TryFromBase64UrlChars(ReadOnlySpan<char> chars, Span<byte> bytes, out int bytesWritten)
public static bool TryFromBase64UrlUtf8(ReadOnlySpan<byte> utf8, Span<byte> bytes, out int bytesWritten)

stringだけ受け入れるのではなくて、ReadOnlySpan<char>と、UTF8を直接受け入れられるようにReadOnlySpan<byte>のオーバーロードを用意しましょう（面倒くせえ……）。中身の実装はcharとbyteで似てるようで若干違うので今回は雑にコピペコードで済ませてます。コピペ最強。

ともあれこれでゼロアロケーションなデコードです。

ちなみにSystem.Security.Cryptographyも、こうしたSpan対応のAPIが(.NET Core 2.1なら)あります。.NET Standard 2.0にはありません。2.1から、なのでまだ先です。

bool TryComputeHash(ReadOnlySpan<byte> source, Span<byte> destination, out int bytesWritten)
bool TrySignData(ReadOnlySpan<byte> data, Span<byte> destination, HashAlgorithmName hashAlgorithm, RSASignaturePadding padding, out int bytesWritten)
bool VerifyData(ReadOnlySpan<byte> data, ReadOnlySpan<byte> signature, HashAlgorithmName hashAlgorithm, RSASignaturePadding padding)

今回の最初のリリースが.NETCore Appのみなのは、主にこの辺が理由です。迂回できないこともないんですけどねえ。

stackallocとArrayPoolをめっちゃ使う

先のBase64のデコード繋がりで説明すると、デコード先のbyte[]をどう用意するか、という話であり。headerのBase64とかsignatureのBase64とか、あまり大きくないのが確定しているので、stackallocをSpanで受けて、デコード先を作ります。

Span<byte> bytes = stackalloc byte[Base64.GetMaxBase64UrlDecodeLength(header.Length)];
if (!Base64.TryFromBase64UrlUtf8(header, bytes, out var bytesWritten))

Payloadは長さがわからない（そこそこ大きい可能性もある）ので、stackallocで受けるのは不安があるので、ArrayPoolを使いましょう。

var rentBytes = ArrayPool<byte>.Shared.Rent(Base64.GetMaxBase64UrlDecodeLength(payload.Length));
try
{
    Span<byte> bytes = rentBytes.AsSpan();
    if (!Base64.TryFromBase64UrlUtf8(payload, bytes, out var bytesWritten))
    {
        return DecodeResult.InvalidBase64UrlPayload;
    }
    bytes = bytes.Slice(0, bytesWritten);

    // ....
}
finally
{
    ArrayPool<byte>.Shared.Return(rentBytes);
}

ようするに、今どきnew byte[]なんてしたら殺されるぞ！

ReadOnlySpanの辞書を作る

ReadOnlySpan<byte>はref struct！つまりDictionaryのKeyにはできない！けどルックアップはしたい！

どーいうことかというと、例えば

HeaderのJSONを舐めて、デコードに使うアルゴリズムが何であるかあるかチェックしたいわけですが、まず、今どきはJSONをstringで検索したりはしません。UTF8のままやります(System.Text.Json(preview)やUtf8Jsonを使いましょう)。特に、今回はBase64Urlからデコードしてきたバイナリなので、更にstringにデコードしてしまうのは無駄の極みなので、絶対避けたいわけです。

そうして、algのvalue部分に相当するReadOnlySpanが切り出せたとしましょう。さて、じゃあこれが何であるか。HS256なのかRS512なのか、そして、それをキーにしてIJwtAlgorithmを取り出したいわけです。必要なデータ構造はDictionary<ReadOnlySpan<byte>, IJwtAlgorithm>>なわけです。が、それは無理。C#の言語仕様がそれを許してくれないのです。困ったねえ。

もちろん、答えは自作しましょう。今回はReadOnlyUtf8StringDictionaryというものを用意しました。Dictionary内部で持っておくキーは別にSpanである必要はないので、普通にbyte[]で確保しておきます。ルックアップだけ

public bool TryGetValue(ReadOnlySpan<byte> key, out TValue value)

というAPIを用意すればOKという寸法です。

実装において、byte[]の一致比較はSpanのSequenceEqualを使えば良いんですが、GetHashCodeの実装だけはどうにもなりません（Utf8Stringも控えてることだし、標準でいい感じのがそろそろ入るといいんですけどねえ）。私は延々と使いまわせいているFarmHashの実装をコピペで用意していますが、適当にxxHashを実装したり何かするといいと思います。適当に拾ってきたものを使うとパフォーマンス的に意味のないクソ実装の可能性もあるので、その辺は適当に気をつけましょう。

最後まで配列の切り出しをしない実装を作る

jwtEncoderのEncodeメソッドは、3つのオーバーロード(名前違い含む)を持ってます。

string Encode<T>(...)
byte[] EncodeAsUtf8Bytes<T>(...)
void Encode<T>(IBufferWriter<byte> writer, ...)

一番使うのは、stringだとは思います。Httpのヘッダーとかに埋めたりするケースが多いと思うので、stringが要求されるのでしょーがない。でも、byte[]を返すもののほうが高速です。内部的には全てUtf8 byte[]で処理しているので、stringへのエンコード処理をバイパスできるからです。例えばgRPCは(MagicOnionも)、バイナリヘッダーを許容しているので、stringヘッダーよりも高速に処理できます。

// gRPC Header
var metadata = new Metadata();
metadata.Add("auth-token-bin", encoder.EncodeAsUtf8Bytes());

さて、じゃあ最後の IBufferWriter<byte> はなにかというと、直接これに書き込みます。まぁ、Span<byte>,int bytesWrittenみたいなものですが、Span<byte>を渡すのが使えるのって、処理後の長さが概ね分かっているときで、JwtのエンコードはPayloadの処理とかあるので、基本的には処理が完了するまで分かりません。ので、bytesWritten形式のAPIは向いてません。

IBufferWriterはStreamみたいなもので、これに直接書き込みます。新しいI/O APIである System.IO.Pipelines で使われているAPIで、つまりは、一応それに対応しているということで。MessagePack-CSharpのv2(現在絶賛制作中)も、IBufferWriterが主役になっています。時代はダイレクトライト。

System.IdentityModel.Tokens.Jwtは最低

JWTの話は特にするつもりはなかったんですが、とにかくSystem.IdentityModel.Tokens.Jwtが最低だということは言っておきたい！とにかくAPIがヤバい！まぁ、これ、他の認証系も統合された抽象化の上に乗っているので、JWT的に不要で意味不明なものがいっぱいついているうえに、その抽象化がエンタープライズグレード（笑）の重厚長大な酷いもので、Microsoftの認証が難しいと感じるとしたら（実際難しい）、ただたんにライブラリのAPIが腐ってるから難しいだけですからね。

何かのフレームワークと統合されてて、ワンポチで導入される、とかだったらまだいいんですが、直接は触りたくないですねえ。誰が作ってんだかって感じですが(お、公開されてる先はAzure配下かな……)

まとめ

MagicOnionで――というのもありますが、認証系はJWT中心に、ちょっと色々考えてます。あとまぁ、さすがにパフォーマンスだけが差別化要因というのはしょっぱいので、Unity対応しよ。