Implement compression with Tower layer

Author: daxpedda

Cargo.toml

@@ -16,9 +16,8 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 
 wasm-bindgen-cli-support = "0.2"
 
-axum = { version = "0.5", default-features = false, features = ["http1", "headers"] }
+axum = { version = "0.5", default-features = false, features = ["http1"] }
 tokio = { version = "1.11", default-features = false, features = ["rt-multi-thread"] }
-tower-http = { version = "0.2", features = ["fs", "trace"] }
+tower-http = { version = "0.2", features = ["compression-full", "fs", "trace"] }
 tower = "0.4"
 fastrand = "1.5"
-flate2 = "1.0"

src/main.rs

@@ -1,7 +1,6 @@
 use std::path::PathBuf;
 
 use anyhow::{anyhow, ensure};
-use tracing::info;
 use tracing_subscriber::EnvFilter;
 
 use crate::server::Options;
@@ -32,20 +31,8 @@ fn main() -> Result<(), anyhow::Error> {
 
     let output = wasm_bindgen::generate(&wasm_file)?;
 
-    info!("compressed wasm output is {} large", pretty_size(output.compressed_wasm.len()));
-
     let rt = tokio::runtime::Runtime::new()?;
     rt.block_on(server::run_server(options, output))?;
 
     Ok(())
 }
-
-fn pretty_size(size_in_bytes: usize) -> String {
-    let size_in_kb = size_in_bytes as f32 / 1024.0;
-    if size_in_kb < 1024.0 {
-        return format!("{:.2}kb", size_in_kb);
-    }
-
-    let size_in_mb = size_in_kb / 1024.0;
-    return format!("{:.2}mb", size_in_mb);
-}

src/server.rs

@@ -1,11 +1,11 @@
 use std::net::SocketAddr;
 
-use axum::headers::ContentEncoding;
 use axum::http::{HeaderValue, StatusCode};
 use axum::response::{Html, IntoResponse, Response};
 use axum::routing::{get, get_service};
-use axum::{Router, TypedHeader};
+use axum::Router;
 use tower::ServiceBuilder;
+use tower_http::compression::CompressionLayer;
 use tower_http::services::ServeDir;
 
 use crate::wasm_bindgen::WasmBindgenOutput;
@@ -21,24 +21,20 @@ pub struct Options {
 }
 
 pub async fn run_server(options: Options, output: WasmBindgenOutput) -> Result<()> {
-    let WasmBindgenOutput { js, compressed_wasm } = output;
+    let WasmBindgenOutput { js, wasm } = output;
 
-    let middleware_stack = ServiceBuilder::new().into_inner();
+    let middleware_stack = ServiceBuilder::new().layer(CompressionLayer::new()).into_inner();
 
     let version = generate_version();
 
     let html = include_str!("../static/index.html").replace("{{ TITLE }}", &options.title);
 
     let serve_dir = get_service(ServeDir::new(".")).handle_error(internal_server_error);
 
-    let serve_wasm = || async move {
-        (TypedHeader(ContentEncoding::gzip()), WithContentType("application/wasm", compressed_wasm))
-    };
-
     let app = Router::new()
         .route("/", get(move || async { Html(html) }))
         .route("/api/wasm.js", get(|| async { WithContentType("application/javascript", js) }))
-        .route("/api/wasm.wasm", get(serve_wasm))
+        .route("/api/wasm.wasm", get(|| async { WithContentType("application/wasm", wasm) }))
         .route("/api/version", get(move || async { version }))
         .fallback(serve_dir)
         .layer(middleware_stack);

src/wasm_bindgen.rs

@@ -1,13 +1,10 @@
 use crate::Result;
-use anyhow::Context;
 use std::path::Path;
 use tracing::debug;
 
-const COMPRESSION_LEVEL: u32 = 2;
-
 pub struct WasmBindgenOutput {
     pub js: String,
-    pub compressed_wasm: Vec<u8>,
+    pub wasm: Vec<u8>,
 }
 pub fn generate(wasm_file: &Path) -> Result<WasmBindgenOutput> {
     debug!("running wasm-bindgen...");
@@ -26,23 +23,5 @@ pub fn generate(wasm_file: &Path) -> Result<WasmBindgenOutput> {
     let wasm = output.wasm_mut().emit_wasm();
     debug!("emitting wasm took {:?}", start.elapsed());
 
-    debug!("compressing wasm...");
-    let start = std::time::Instant::now();
-    let compressed_wasm = compress(&wasm).context("failed to compress wasm file")?;
-    debug!("compressing took {:?}", start.elapsed());
-
-    Ok(WasmBindgenOutput { js, compressed_wasm })
-}
-
-fn compress(bytes: &[u8]) -> Result<Vec<u8>, std::io::Error> {
-    use flate2::write::GzEncoder;
-    use flate2::Compression;
-    use std::io::prelude::*;
-
-    let mut encoder = GzEncoder::new(Vec::new(), Compression::new(COMPRESSION_LEVEL));
-
-    encoder.write_all(bytes)?;
-    let compressed = encoder.finish()?;
-
-    Ok(compressed)
+    Ok(WasmBindgenOutput { js, wasm })
 }