m13o.net

2020-12-07 Mon 22:48
Rustでバイナリを読む その6   AdventCalendar2020 Rust programming

前回の最後に, wavファイルを読み取るcrateを実装するだなんだとほざきましたが, binary_readerの話をします. といいますのも, wav編を始めるにあたって見直していたらどうしても気になる処がありました. なぜ今になって気付いてしまったのか……

以下のように, 特定バイト読み取った後, バッファのサイズを超えた読み取りをしようとした時の事を考えます.

#[test]
    fn buffer_overrun_after_read_buffer() {
        let buffer = vec![0x02, 0x34, 0x58, 0xFF, 0x80, 0x99, 0x00, 0x42];
        let mut reader = BinaryReader::new(&buffer);
        let (values, rest) = reader.read(5).unwrap();
        assert_eq!(values.len(), 5);
        assert_eq!(values[0], 0x02);
        assert_eq!(values[1], 0x34);
        assert_eq!(values[2], 0x58);
        assert_eq!(values[3], 0xFF);
        assert_eq!(values[4], 0x80);
        assert_eq!(rest.len(), 3);
        assert_eq!(rest[0], 0x99);
        assert_eq!(rest[1], 0x00);
        assert_eq!(rest[2], 0x42);

        let result = reader.read(4);
        assert_eq!(result,
                   Err(BufferOverRunError {
                       buffer_size: 8,
                       index: 5,
                       read_size: 4
                   }));
    }

これは現行のコードでも想定通りテストをパスしますが, この時のreaderの内部状態はどうなっているかというと……

pub fn read(&mut self, size: usize) -> Result<(&'a [u8], &'a [u8]), BufferOverRunError> {
    // ...

    self.current_index += size; // <- ここ!
    if self.current_index > max_buffer_size {
        return Err(BufferOverRunError {
            buffer_size: max_buffer_size,
            index,
            read_size: size,
        });
    }

    // ...
}

current_indexにsizeを加算してから比較しているので, Ok(T)を返そうがErr(E)を返そうが, 読取用のインデックスが進んでしまう事になります. 一度読み取りサイズを超えてしまったら永続的にエラーとする場合はこれでも良いかもしれませんが, current_indexの型はusizeなのでどこかでオーバーフローを発生させてしまいます. また, そもそも読み取っていないにもかかわらずそのインデックスがズレてしまうのはよくありません. ですので, これを修正します.

先程のテスト, buffer_overrun_after_read_buffer() を read_buffer_after_overrun()にリネームして, 以下のようにテストを追加します.

#[test]
fn read_buffer_after_overrun() { // <- rename!
    let buffer = vec![0x02, 0x34, 0x58, 0xFF, 0x80, 0x99, 0x00, 0x42];
    let mut reader = BinaryReader::new(&buffer);

    // ...

    let result = reader.read(4);
    assert_eq!(result,
               Err(BufferOverRunError {
                   buffer_size: 8,
                   index: 5,
                   read_size: 4
               }));

    let (values, rest) = reader.read(3).unwrap();
    assert_eq!(values.len(), 3);
    assert_eq!(values[0], 0x99);
    assert_eq!(values[1], 0x00);
    assert_eq!(values[2], 0x42);
    assert_eq!(rest.len(), 0);
}

ここまで記述した状態で一度 cargo testを実行してみましょう.

test binary_reader::tests::read_buffer_after_overrun ... FAILED


failures:

---- binary_reader::tests::read_buffer_after_overrun stdout ----
thread 'binary_reader::tests::read_buffer_after_overrun' panicked at 'called `Result::unwrap()` on an `Err` value: バッファオーバーラン バッファサイズ 8, 読み出し位置 9, 読み出しバイト数 3', binary_reader/src/binary_reader.rs:219:45
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace


failures:
    binary_reader::tests::read_buffer_after_overrun

このように読み出し位置がおかしくなっているので, 期待した動作を満たしません. このテストが通るようにread()メソッドを修正します.

pub fn read(&mut self, size: usize) -> Result<(&'a [u8], &'a [u8]), BufferOverRunError> {
    let index = self.current_index;

    // ....

    if (index + size) > max_buffer_size {
        return Err(BufferOverRunError {
            buffer_size: max_buffer_size,
            index,
            read_size: size,
        });
    }

    self.current_index += size;
    Ok((
        &self.buffer[index..self.current_index],
        &self.buffer[self.current_index..],
    ))
}

このように, current_indexにsizeを加算する処理をエラー条件判定の後ろに持ってきて, 実際の条件判定の処理では, indexとsizeを加算した物を比較に利用してcurrent_indexの状態が変わらないようにしています.

これで改めてcargo testを実行してみると……

running 28 tests
test binary_reader::tests::read_bytes ... ok
test binary_reader::tests::read_buffer_overrun ... ok
test binary_reader::tests::read_buffer_after_overrun ... ok
test binary_reader::tests::read_empty_buffer ... ok

# .......

test result: ok. 28 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out

無事テストをパスしました. これでまた一つBinaryReaderの品質が上がりました. めでたしめでたし.