概要

Goでファイルを読み込んでいる時に、そのファイルのタイプを判別したいことがたまにあります。例えばGzipかどうか分からないけど、もしGzipならgzip.NewReader噛ませたい、みたいな場合です。雑にgzip.NewReader噛ませてerr返すかどうかで判定とかやってみたんですが、普通に10バイト読み進められちゃうのでerr返ったあとに別のファイルタイプとして処理しようとするとinvalidなヘッダーになって死にます。実は読み進められたバイトを戻す方法あるよ、という場合は教えて下さい。

そもそもGzip以外の判定をしたいときもあるので、NewReaderの方針も必ず使えるわけではありません。もしファイルがos.Fileとかbufio.Readerの形であればReadしてからSeekしたりPeekしたり出来るのですが、io.Readerの場合どうやるのか分からなかったので調べました。

ちなみに自分では全く思いつかなくて containers/image 読んでて見つけました。賢いなーと思ったので書いておきます。よくよく考えると以前も同じようなのを見て賢いなーと思った記憶があり、何度も感動できてお得なのですが時間の無駄ではあるので次の感動が来ないように書いておきます。

解説は良いからコード見せろという人のために先に置いておきます。
https://github.com/containers/skopeo/blob/8f24d281302cc6aaa7ba9301d4a87bfceb7141b6/vendor/github.com/containers/image/v5/pkg/compression/compression.go#L96

実装

今回の例ではGzipかどうかの判定をしようと思います。マジックバイトの判定なら何でも同じ方法で行けるはずです。

func isGzip(input io.Reader) (io.Reader, bool, error) {
    buf := [3]byte{}

    n, err := io.ReadAtLeast(input, buf[:], len(buf))
    if err != nil {
        return nil, false, err
    }

    isGzip := buf[0] == 0x1F && buf[1] == 0x8B && buf[2] == 0x8
    return io.MultiReader(bytes.NewReader(buf[:n]), input), isGzip, nil
}

まず、io.Readerから普通にReadします。ここでは参考にしたコードに合わせてio.ReadAtLeastを使ってますが、len(buf)にしてるのでio.ReadFullと同じじゃないかなーと思ってます。とりあえず3バイト読み込みます。

次に、その3バイトを使ってGzipかどうか判定します。もちろん他のファイルタイプの場合は違う処理になりますし、3バイトではないかもしれません。

そして最後がポイントですが、bytes.NewReader(buf[:n])を使って読み込んだ3バイトをio.Readerインタフェースに合わせます。そしてio.MultiReaderを使って残りの4バイト目以降と結合します。こうすることで、io.MultiReaderによって作られたio.Readerは最初のReaderから3バイト読み込み、次のReaderから4バイト目以降を読み込みます。結果として戻り値として返されるReaderは最初のinputと同じ値を読み込むことが可能になります。

下の例ではbytes.NewBufferなのであまり良い例ではないですがサンプルコードということで一応置いておきます。
https://play.golang.org/p/VH19FHQnqdr

追記（2019/01/12）
@yamasaki-masahide さんからのコメントにあるように、io.Copy使う場合はMultiReaderにWriteToが実装されていない関係でパフォーマンス的にあまり良くないようです。オシャレさに感動して紹介しましたが、大人しくbufio.Reader使う方が多くのケースでは良さそうです。

まとめ

io.MultiReaderいつ使うんじゃ！と思っていたのですが普通に便利でした。こうしたちょっとしたテクニックは忘れがちなのでちゃんと書いておきたい所存なのですが、その所存も忘れがち問題があります。

io.Readerのファイルタイプを判定する

概要

実装

まとめ