在现代网络应用中，处理二进制数据，例如文件上传和下载，是不可避免的任务。本文将深入探讨如何使用 Rust 构建强大的 API 服务器，以处理发送和接收字节数据的各种场景。我们将使用 axum 框架，这是一个非常适合路由请求和声明式解析请求的 Web 应用框架。

准备工作

首先，在 Cargo.toml 文件中添加必要的依赖项：

[dependencies]
tokio = { version = "1", features = ["macros"] }
lambda_http = "0.13.0"
axum = { version = "0.7.5", features = ["multipart"] }
tower-http = { version = "0.5.2", features = ["limit"] }
regex = "1.10.6"
urlencoding = "2.1.3"

这里我们添加了 axum 用于构建 Web 应用程序，tokio 用于异步运行时，tower-http 用于设置请求体大小限制，regex 用于解析文件名，urlencoding 用于处理文件名编码，以及 lambda_http 用于在 AWS Lambda 上部署（可选）。

接收二进制数据

直接接收字节数据

当直接接收请求体中的字节数据时，我们需要在处理函数中完成以下步骤：

1. 从 Content-Disposition 头部信息中提取 file_name。
2. 将 Bytes 类型转换为 Vec<u8> 类型，以便后续处理，例如保存文件。

use axum::{
    http::HeaderMap,
    Json,
    extract::Bytes,
    response::{IntoResponse, Response},
};
use serde_json::{json, Value};
use regex::Regex;
use urlencoding::decode;

async fn post_bytes(headers: HeaderMap, bytes: Bytes) -> Json<Value> {
    let Some(content_disposition) = headers.get("content-disposition") else {
        return Json(json!({
            "message": "content disposition not available"
        }));
    };

    let Ok(content_disposition_string) = content_disposition.to_str() else {
        return Json(json!({
            "message": "content disposition not available"
        }));
    };

    let Ok(re) = Regex::new(r#"((.|\\s\\S|\\r|\\n)*)filename\*?=utf-8''(?<name>((.|\\s\\S|\\r|\\n)*))"#) else {
        return Json(json!({
            "message": "Regex failed to create."
        }));
    };

    let Some(captures) = re.captures(content_disposition_string) else {
        return Json(json!({
            "message": "File name not found."
        }));
    };
    let file_name = decode(&captures["name"]).expect("UTF-8").to_string();
    let data: Vec<u8> = match bytes.bytes().collect() {
        Ok(data) => data,
        Err(err) => {
            return Json(json!({
                "message": format!("Bytes data not available. Error: {}", err.to_string())
            }));
        }
    };

    return Json(json!({
        "file_name": file_name,
        "data_length": data.len()
    }));
}

这段代码首先尝试从请求头中获取 Content-Disposition 字段，然后使用正则表达式提取文件名。最后，将字节数据收集到 Vec<u8> 中，并返回包含文件名和数据长度的 JSON 响应。

处理 Multipart/form-data 数据

对于 Multipart/form-data 类型的请求，我们需要使用 axum::extract::Multipart 来处理。

use axum::extract::Multipart;

async fn post_bytes_multiparts(mut multipart: Multipart) -> Json<Value> {
    while let Some(field) = multipart.next_field().await.unwrap() {
        let name = field.name().unwrap().to_string();
        let file_name = field.file_name().unwrap().to_string();
        let content_type = field.content_type().unwrap().to_string();
        let data = field.bytes().await.unwrap();

        if name == "file" {
            return Json(json!({
                "file_name": file_name,
                "content_type": content_type,
                "data_length": data.len()
            }));
        }
    }

    Json(json!({
        "message": "file field is missing in the form."
    }))
}

这段代码使用 Multipart 提取器迭代表单中的每个字段，并检查字段名称是否为 "file"。如果是，则提取文件名、内容类型和数据，并返回 JSON 响应。

设置请求体大小限制

为了防止恶意请求导致服务器资源耗尽，我们需要设置请求体大小限制。默认情况下，axum 会将请求体大小限制为 2MB。

use axum::body::Body;
use axum::http::Request;
use tower::ServiceBuilder;
use tower_http::limit::RequestBodyLimitLayer;
use std::convert::Infallible;

// ...

let app = Router::new()
    // ... 路由
    // 设置不同的限制：10GB
    .layer(
        ServiceBuilder::new()
            // Disable the default body limit.
            .layer(RequestBodyLimitLayer::disable())
            // Set a limit of 10 GiB.
            .layer(RequestBodyLimitLayer::new(10 * 1024 * 1024 * 1024))
    );

这段代码首先使用 RequestBodyLimitLayer::disable() 禁用默认的请求体大小限制，然后使用 RequestBodyLimitLayer::new() 设置新的限制为 10GB。

发送二进制数据

与接收二进制数据不同，发送二进制数据时，我们需要确保响应头中的 Content-Type 和 Content-Disposition 设置正确。

use axum::http::{StatusCode, HeaderValue};
use urlencoding::encode;

fn build_error_response(message: &str) -> Response {
    let mut json_header = HeaderMap::new();
    json_header.insert("content-type", HeaderValue::from_static("application/json"));

    let mut response = Response::new(json!({
        "message": message
    }).to_string());
    *response.status_mut() = StatusCode::BAD_REQUEST;
    return (json_header, response).into_response();
}

async fn echo_bytes(headers: HeaderMap, bytes: Bytes) -> Response {
    let Some(content_type) = headers.get("content-type") else {
        return build_error_response("content type not available");
    };

    let Some(content_disposition) = headers.get("content-disposition") else {
        return build_error_response("content disposition not available");
    };

    let Ok(content_disposition_string) = content_disposition.to_str() else {
        return build_error_response("content disposition not available");
    };

    let Ok(re) = Regex::new(r#"((.|\\s\\S|\\r|\\n)*)filename\*?=utf-8''(?<name>((.|\\s\\S|\\r|\\n)*))"#) else {
        return build_error_response("Regex failed to create");
    };

    let Some(captures) = re.captures(content_disposition_string) else {
        return build_error_response("File name not found in Content-Disposition");
    };

    let file_name = decode(&captures["name"]).expect("UTF-8").to_string();
    let data: Vec<u8> = match bytes.bytes().collect() {
        Ok(data) => data,
        Err(err) => {
            return build_error_response(&format!("Bytes data not available. Error: {}", err.to_string()));
        }
    };

    let file_name_encode = encode(&file_name).to_string();
    let mut headers = HeaderMap::new();
    headers.insert("content-type", content_type.to_owned());
    headers.insert("content-length", HeaderValue::from(bytes.len()));
    headers.insert("content-disposition", HeaderValue::from_str(&format!("attachment; filename*=utf-8''{}", file_name_encode)).unwrap());

    (headers, data).into_response()
}

这段代码首先从请求头中获取 Content-Type 和 Content-Disposition 字段，然后提取文件名。接着，将字节数据收集到 Vec<u8> 中。最后，构建新的响应头，并将数据作为响应体返回。

总结

本文详细介绍了如何使用 Rust 构建 API 服务器来处理发送和接收字节数据。我们学习了如何使用 axum 框架处理不同类型的请求，如何设置请求体大小限制，以及如何正确设置响应头以发送二进制数据。