diff --git a/crates/biliup/src/downloader/flv_writer.rs b/crates/biliup/src/downloader/flv_writer.rs index fd3adaf..042630f 100644 --- a/crates/biliup/src/downloader/flv_writer.rs +++ b/crates/biliup/src/downloader/flv_writer.rs @@ -68,7 +68,9 @@ impl<'a> FlvFile<'a> { ) -> std::io::Result { self.write_tag_header(tag_header)?; self.buf_writer.write_all(body)?; - self.buf_writer.write(previous_tag_size) + // write 允许部分写入,短写会静默丢字节并破坏 FLV 结构,必须用 write_all + self.buf_writer.write_all(previous_tag_size)?; + Ok(previous_tag_size.len()) } pub fn write_tag_header(&mut self, tag_header: &TagHeader) -> std::io::Result<()> { @@ -86,7 +88,9 @@ impl<'a> FlvFile<'a> { writer: &mut impl Write, previous_tag_size: u32, ) -> std::io::Result { - writer.write(&previous_tag_size.to_be_bytes()) + let bytes = previous_tag_size.to_be_bytes(); + writer.write_all(&bytes)?; + Ok(bytes.len()) } } @@ -104,7 +108,8 @@ pub struct FlvTag<'a> { pub fn to_json(mut writer: impl Write, t: &T) -> std::io::Result { serde_json::to_writer(&mut writer, t)?; - writer.write("\n".as_ref()) + writer.write_all(b"\n")?; + Ok(1) } #[derive(Debug, PartialEq, Serialize)] @@ -124,3 +129,39 @@ pub enum TagDataHeader<'a> { }, Script(ScriptData<'a>), } + +#[cfg(test)] +mod tests { + use super::*; + + /// 模拟每次调用最多接受 1 字节的 Writer。 + /// 对这类 Writer,`write` 只写入部分数据也返回 Ok,必须用 `write_all` 才能保证完整写入。 + struct ShortWriter(Vec); + + impl Write for ShortWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let take = buf.len().min(1); + self.0.extend_from_slice(&buf[..take]); + Ok(take) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + #[test] + fn previous_tag_size_is_fully_written_even_on_short_writes() { + let mut writer = ShortWriter(Vec::new()); + let written = FlvFile::write_previous_tag_size(&mut writer, 0x0102_0304).unwrap(); + assert_eq!(written, 4); + assert_eq!(writer.0, [0x01, 0x02, 0x03, 0x04]); + } + + #[test] + fn to_json_writes_trailing_newline_even_on_short_writes() { + let mut writer = ShortWriter(Vec::new()); + to_json(&mut writer, &serde_json::json!({"k": "v"})).unwrap(); + assert!(writer.0.ends_with(b"\n")); + } +}