diff --git a/Cargo.lock b/Cargo.lock index 2a4e2db..da87dd3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -399,6 +399,7 @@ dependencies = [ "axum", "bytes", "dotenvy", + "futures-util", "hex", "hmac", "openrouter-rs", @@ -409,6 +410,8 @@ dependencies = [ "subtle", "thiserror 2.0.18", "tokio", + "tokio-stream", + "tokio-util", ] [[package]] @@ -1379,6 +1382,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.18" diff --git a/Cargo.toml b/Cargo.toml index 2dba50b..d036f9a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,9 @@ edition = "2024" [dependencies] reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } tokio = { version = "1.52", features = ["full"] } +tokio-stream = "0.1" +tokio-util = "0.7" +futures-util = "0.3" serde_json = "1.0" serde = { version = "1.0", features = ["derive"] } openrouter-rs = "0.10" diff --git a/src/bot.rs b/src/bot.rs index 72b80ab..d03d2fd 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -1,9 +1,11 @@ -use std::time::Duration; - +use futures_util::stream::TryStreamExt; use serde::Deserialize; +use std::time::Duration; +use tokio::io::AsyncReadExt; +use tokio_util::io::StreamReader; use crate::{ - consts::{BOT_PROCESS_MSG, REVIEW_PROMPT}, + consts::{BOT_PROCESS_MSG, MAX_DIFF_SIZE, REVIEW_PROMPT}, env::EnvConfig, errors::AppError, gitea::{GiteaAPI, ReviewPayload, WebhookType}, @@ -197,7 +199,18 @@ impl Bot { async fn download_git_diff(&self, url: &str) -> anyhow::Result { let response = self.http_client.get(url).send().await?; - let body = response.text().await?; - Ok(body) + let stream = response.bytes_stream().map_err(std::io::Error::other); + + let mut buf = Vec::with_capacity(MAX_DIFF_SIZE); + StreamReader::new(stream) + .take((MAX_DIFF_SIZE + 1) as u64) + .read_to_end(&mut buf) + .await?; + + if buf.len() > MAX_DIFF_SIZE { + anyhow::bail!("Git diff exceeds the maximum allowed size of 1 Mo"); + } + + Ok(String::from_utf8_lossy(&buf).into_owned()) } } diff --git a/src/consts.rs b/src/consts.rs index ded647d..f3ae2dd 100644 --- a/src/consts.rs +++ b/src/consts.rs @@ -1,6 +1,8 @@ pub const GITEA_SIG_HEADER_NAME: &str = "x-gitea-signature"; pub const GITEA_EVENT_TYPE_HEADER_NAME: &str = "x-gitea-event-type"; + pub const MAX_WEBHOOK_BODY_SIZE: usize = 1024 * 1024; // 1 MiB +pub const MAX_DIFF_SIZE: usize = 1024 * 1024; // 1 MiB pub const BOT_PROCESS_MSG: &str = " Review in progress with the model \"{model}\"...