Site: add brotli compression to POST request responses

Author: miwig

Cargo.lock

@@ -20,6 +20,7 @@ http = "0.2"
 home = "0.5"
 chrono = "0.4"
 rmp-serde = "0.15"
+brotli = "3.3.3"
 semver = "1.0"
 ring = "0.16.10"
 hex = "0.4.2"

site/Cargo.toml

@@ -1,3 +1,5 @@
+use brotli::enc::BrotliEncoderParams;
+use brotli::BrotliCompress;
 use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::path::Path;
@@ -313,6 +315,19 @@ async fn serve_req(server: Server, req: Request) -> Result<Response, ServerError
     let path = req.uri().path().to_owned();
     let path = path.as_str();
 
+    let allow_compression = req
+        .headers()
+        .get(hyper::header::ACCEPT_ENCODING)
+        .map_or(false, |e| e.to_str().unwrap().contains("br"));
+
+    let compression = if allow_compression {
+        let mut brotli = BrotliEncoderParams::default();
+        brotli.quality = 4;
+        Some(brotli)
+    } else {
+        None
+    };
+
     if let Some(response) = handle_fs_path(path) {
         return Ok(response);
     }
@@ -404,6 +419,7 @@ async fn serve_req(server: Server, req: Request) -> Result<Response, ServerError
             crate::comparison::handle_compare(check!(parse_body(&body)), &ctxt)
                 .await
                 .map_err(|e| e.to_string()),
+            &compression,
         )),
         "/perf/collected" => {
             if !server.check_auth(&req) {
@@ -412,7 +428,10 @@ async fn serve_req(server: Server, req: Request) -> Result<Response, ServerError
                     .body(hyper::Body::empty())
                     .unwrap());
             }
-            Ok(to_response(request_handlers::handle_collected().await))
+            Ok(to_response(
+                request_handlers::handle_collected().await,
+                &compression,
+            ))
         }
         "/perf/github-hook" => {
             if !verify_gh(&ctxt.config, &req, &body) {
@@ -435,6 +454,7 @@ async fn serve_req(server: Server, req: Request) -> Result<Response, ServerError
             match event.as_str() {
                 "issue_comment" => Ok(to_response(
                     request_handlers::handle_github(check!(parse_body(&body)), ctxt.clone()).await,
+                    &compression,
                 )),
                 _ => Ok(http::Response::builder()
                     .status(StatusCode::OK)
@@ -444,9 +464,11 @@ async fn serve_req(server: Server, req: Request) -> Result<Response, ServerError
         }
         "/perf/self-profile" => Ok(to_response(
             request_handlers::handle_self_profile(check!(parse_body(&body)), &ctxt).await,
+            &compression,
         )),
         "/perf/self-profile-raw" => Ok(to_response(
             request_handlers::handle_self_profile_raw(check!(parse_body(&body)), &ctxt).await,
+            &compression,
         )),
         "/perf/bootstrap" => Ok(
             match request_handlers::handle_bootstrap(check!(parse_body(&body)), &ctxt).await {
@@ -594,7 +616,7 @@ fn verify_gh_sig(cfg: &Config, header: &str, body: &[u8]) -> Option<bool> {
     Some(false)
 }
 
-fn to_response<S>(result: ServerResult<S>) -> Response
+fn to_response<S>(result: ServerResult<S>, compression: &Option<BrotliEncoderParams>) -> Response
 where
     S: Serialize,
 {
@@ -604,7 +626,7 @@ where
                 .header_typed(ContentType::octet_stream())
                 .header_typed(CacheControl::new().with_no_cache().with_no_store());
             let body = rmp_serde::to_vec_named(&result).unwrap();
-            response.body(hyper::Body::from(body)).unwrap()
+            maybe_compressed_response(response, body, compression)
         }
         Err(err) => http::Response::builder()
             .status(StatusCode::INTERNAL_SERVER_ERROR)
@@ -615,6 +637,30 @@ where
     }
 }
 
+fn maybe_compressed_response(
+    response: http::response::Builder,
+    body: Vec<u8>,
+    compression: &Option<BrotliEncoderParams>,
+) -> Response {
+    match compression {
+        None => response.body(hyper::Body::from(body)).unwrap(),
+        Some(brotli_params) => {
+            let compressed = compress_bytes(&body, brotli_params);
+            let response = response.header(
+                hyper::header::CONTENT_ENCODING,
+                hyper::header::HeaderValue::from_static("br"),
+            );
+            response.body(hyper::Body::from(compressed)).unwrap()
+        }
+    }
+}
+
+fn compress_bytes(mut bytes: &[u8], brotli_params: &BrotliEncoderParams) -> Vec<u8> {
+    let mut compressed = Vec::with_capacity(bytes.len());
+    BrotliCompress(&mut bytes, &mut compressed, brotli_params).unwrap();
+    compressed
+}
+
 fn to_triage_response(result: ServerResult<api::triage::Response>) -> Response {
     match result {
         Ok(result) => {