From 6ae87364b4a67fd5f3725585f77c8230337f7acc Mon Sep 17 00:00:00 2001
From: amtoaer <amtoaer@gmail.com>
Date: Wed, 22 Jan 2025 00:18:04 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E4=B8=BA=E4=B8=8B=E8=BD=BD=E5=8A=A0?=
 =?UTF-8?q?=E5=85=A5=20flush=20=E4=B8=8E=20content-length=20=E6=A3=80?=
 =?UTF-8?q?=E6=9F=A5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 crates/bili_sync/src/downloader.rs | 31 +++++++++++++++++++++---------
 1 file changed, 22 insertions(+), 9 deletions(-)

diff --git a/crates/bili_sync/src/downloader.rs b/crates/bili_sync/src/downloader.rs
index 22e718c..34e721d 100644
--- a/crates/bili_sync/src/downloader.rs
+++ b/crates/bili_sync/src/downloader.rs
@@ -1,10 +1,11 @@
+use core::str;
 use std::path::Path;
 
-use anyhow::{anyhow, Result};
+use anyhow::{bail, ensure, Result};
 use futures::StreamExt;
 use reqwest::Method;
 use tokio::fs::{self, File};
-use tokio::io;
+use tokio::io::{self, AsyncWriteExt};
 
 use crate::bilibili::Client;
 pub struct Downloader {
@@ -24,10 +25,25 @@ impl Downloader {
             fs::create_dir_all(parent).await?;
         }
         let mut file = File::create(path).await?;
-        let mut res = self.client.request(Method::GET, url, None).send().await?.bytes_stream();
-        while let Some(item) = res.next().await {
-            io::copy(&mut item?.as_ref(), &mut file).await?;
+        let resp = self.client.request(Method::GET, url, None).send().await?;
+        let expected = resp.content_length().unwrap_or_else(|| {
+            warn!("content length is missing, fallback to 0");
+            0
+        });
+        let mut received = 0u64;
+        let mut stream = resp.bytes_stream();
+        while let Some(bytes) = stream.next().await {
+            let bytes = bytes?;
+            received += bytes.len() as u64;
+            io::copy(&mut bytes.as_ref(), &mut file).await?;
         }
+        file.flush().await?;
+        ensure!(
+            received >= expected,
+            "received {} bytes, expected {} bytes",
+            received,
+            expected
+        );
         Ok(())
     }
 
@@ -46,10 +62,7 @@ impl Downloader {
             .output()
             .await?;
         if !output.status.success() {
-            return match String::from_utf8(output.stderr) {
-                Ok(err) => Err(anyhow!(err)),
-                _ => Err(anyhow!("ffmpeg error")),
-            };
+            bail!("ffmpeg error: {}", str::from_utf8(&output.stderr).unwrap_or("unknown"));
         }
         Ok(())
     }