diff --git a/src/route/speed_test.rs b/src/route/speed_test.rs index ff5ffa1..d0852c9 100644 --- a/src/route/speed_test.rs +++ b/src/route/speed_test.rs @@ -7,11 +7,13 @@ use axum::{ http::header::CONTENT_LENGTH, response::Response, }; -use bytes::Bytes; +use bytes::{Bytes, BytesMut}; use rand::{prelude::SmallRng, RngCore, SeedableRng}; use crate::{route::forbidden, util::string_to_hash, AppState, MAX_KEY_TIME_DRIFT}; +const BUFFER_SIZE: usize = 16384; + // example: /t/5242880/1645930666/bce541b2a97788319e53a754b47e1801204ae7bf/43432228 pub(super) async fn speedtest(Path((size, time, hash)): Path<(u64, i64, String)>, data: State>) -> Response { // Check time & hash @@ -27,14 +29,17 @@ pub(super) fn random_response(size: u64) -> Response { Response::builder() .header(CONTENT_LENGTH, size) .body(Body::from_stream(stream! { - let mut buffer = [0; 8192]; + let mut buffer = BytesMut::zeroed(BUFFER_SIZE); SmallRng::from_entropy().fill_bytes(&mut buffer); - let buffer = Bytes::copy_from_slice(&buffer); + let buffer = buffer.freeze(); let mut filled = 0; while filled < size { - let size = cmp::min(size - filled, 8192) as usize; - yield Result::Ok::(buffer.slice(0..size)); + let size = cmp::min(size - filled, BUFFER_SIZE as u64) as usize; + if size == BUFFER_SIZE { + yield Ok(buffer.clone()); + } + yield Ok::(buffer.slice(0..size)); filled += size as u64; } }))