tiktok_api/
media_transfer.rs

1use reqwest::{Body, Client, StatusCode};
2use url::Url;
3
4//
5pub const CHUNK_SIZE_MIN: usize = 1024 * 1024 * 5;
6pub const CHUNK_SIZE_MAX: usize = 1024 * 1024 * 64;
7pub const CHUNK_COUNT_MIN: usize = 1;
8pub const CHUNK_COUNT_MAX: usize = 1000;
9
10//
11fn get_chunk_size(chunk_size: usize) -> usize {
12    core::cmp::min(core::cmp::max(chunk_size, CHUNK_SIZE_MIN), CHUNK_SIZE_MAX)
13}
14
15pub fn get_chunk_size_and_total_chunk_count(
16    video_size: usize,
17    chunk_size: usize,
18) -> (usize, usize) {
19    let chunk_size = get_chunk_size(chunk_size);
20
21    if video_size <= chunk_size {
22        (video_size, 1)
23    } else {
24        (
25            chunk_size,
26            (video_size as f64 / chunk_size as f64).floor() as usize,
27        )
28    }
29}
30
31//
32//
33//
34pub async fn upload_part<T>(
35    client: Client,
36    upload_url: Url,
37    content_type: &str,
38    byte_range: core::ops::Range<usize>,
39    video_size: usize,
40    stream: T,
41) -> Result<StatusCode, UploadError>
42where
43    T: Into<Body>,
44{
45    let content_length = byte_range.end - byte_range.start;
46    let content_range = format!(
47        "bytes {}-{}/{}",
48        byte_range.start,
49        byte_range.end - 1,
50        video_size,
51    );
52
53    // println!("content_length:{content_length} content_range:{content_range}");
54
55    //
56    let response = client
57        .put(upload_url)
58        .header("Content-Type", content_type)
59        .header("Content-Length", content_length)
60        .header("Content-Range", content_range)
61        .body(stream)
62        .send()
63        .await
64        .map_err(UploadError::RespondFailed)?;
65
66    //
67    let response_status = response.status();
68
69    match response_status {
70        StatusCode::PARTIAL_CONTENT | StatusCode::CREATED => Ok(response_status),
71        status => {
72            let response_body = response
73                .bytes()
74                .await
75                .map_err(UploadError::ReadResponseBodyFailed)?;
76            let response_body = response_body.as_ref();
77
78            Err(UploadError::ResponseMismatch(
79                status,
80                response_body.to_vec(),
81            ))
82        }
83    }
84}
85
86#[cfg(feature = "with_tokio")]
87pub async fn upload_part_from_reader_stream<S>(
88    client: Client,
89    upload_url: Url,
90    content_type: &str,
91    byte_range: core::ops::Range<usize>,
92    video_size: usize,
93    stream: S,
94) -> Result<StatusCode, UploadError>
95where
96    S: tokio::io::AsyncRead + Send + Sync + 'static,
97{
98    use tokio_util::io::ReaderStream;
99
100    upload_part(
101        client,
102        upload_url,
103        content_type,
104        byte_range,
105        video_size,
106        Body::wrap_stream(ReaderStream::new(stream)),
107    )
108    .await
109}
110
111#[cfg(feature = "with_tokio_fs")]
112pub async fn upload_part_from_file(
113    client: Client,
114    upload_url: Url,
115    content_type: &str,
116    file_path: &std::path::PathBuf,
117    file_index: core::ops::Range<usize>,
118    file_size: usize,
119) -> Result<StatusCode, UploadError> {
120    use tokio::{
121        fs::File,
122        io::{AsyncReadExt as _, AsyncSeekExt as _, SeekFrom},
123    };
124
125    let file_index_start = file_index.start;
126    let file_index_end = file_index.end;
127
128    let file_take_size = file_index_end - file_index_start;
129
130    let mut file = File::open(&file_path)
131        .await
132        .map_err(UploadError::OpenFileFailed)?;
133    file.seek(SeekFrom::Start(file_index_start as u64))
134        .await
135        .map_err(UploadError::OpenFileFailed)?;
136    let file = file.take(file_take_size as u64);
137
138    upload_part_from_reader_stream(
139        client.to_owned(),
140        upload_url,
141        content_type,
142        file_index,
143        file_size,
144        file,
145    )
146    .await
147}
148
149#[cfg(feature = "with_tokio_fs")]
150pub async fn upload_from_file(
151    client: Client,
152    upload_url: Url,
153    content_type: &str,
154    file_path: &std::path::PathBuf,
155    chunk_size: Option<usize>,
156) -> Result<Vec<Result<StatusCode, UploadError>>, UploadError> {
157    let crate::tokio_fs_util::Info {
158        file_size,
159        file_name: _,
160    } = crate::tokio_fs_util::info(file_path)
161        .await
162        .map_err(UploadError::GetFileInfoFailed)?;
163
164    let video_size = file_size as usize;
165    let (chunk_size, total_chunk_count) =
166        get_chunk_size_and_total_chunk_count(video_size, chunk_size.unwrap_or(CHUNK_SIZE_MAX));
167
168    //
169    if total_chunk_count > CHUNK_COUNT_MAX {
170        return Err(UploadError::ChunkSizeTooSmaillOrFileTooLarge);
171    }
172
173    let mut ret_list = vec![];
174    for chunk_count in CHUNK_COUNT_MIN..=CHUNK_COUNT_MAX {
175        let chunk_index = chunk_count - 1;
176
177        let file_index_start = chunk_index * chunk_size;
178        let file_index_end = if total_chunk_count == chunk_count {
179            video_size
180        } else {
181            file_index_start + chunk_size
182        };
183
184        match upload_part_from_file(
185            client.to_owned(),
186            upload_url.to_owned(),
187            content_type,
188            file_path,
189            file_index_start..file_index_end,
190            file_size as usize,
191        )
192        .await
193        {
194            Ok(x) => ret_list.push(Ok(x)),
195            Err(err) => {
196                ret_list.push(Err(err));
197                break;
198            }
199        }
200
201        if file_index_end >= video_size {
202            break;
203        }
204    }
205
206    Ok(ret_list)
207}
208
209//
210//
211//
212#[derive(Debug)]
213pub enum UploadError {
214    RespondFailed(reqwest::Error),
215    ReadResponseBodyFailed(reqwest::Error),
216    ResponseMismatch(StatusCode, Vec<u8>),
217    #[cfg(feature = "with_tokio_fs")]
218    GetFileInfoFailed(std::io::Error),
219    #[cfg(feature = "with_tokio_fs")]
220    OpenFileFailed(std::io::Error),
221    ChunkSizeTooSmaillOrFileTooLarge,
222}
223impl core::fmt::Display for UploadError {
224    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
225        write!(f, "{self:?}")
226    }
227}
228impl std::error::Error for UploadError {}