Rust 代码挑战系列 3 - Cookie 的读取

这是一个学习 Rust 的小系列文章,通过完成来自 shuttle 平台所举办的 2023 Christmas Code Hunt 里的每个小挑战,来学习 rust web 框架的使用,此为第三篇文章。

本篇内 Part 7 涉及 Cookie 的获取。

前言

随着路由函数越来越多,全都写在 main.rs 里会显得很臃肿,是时候调整下目录结构了:

新建 src/handler 目录,我们的路由函数将都放到这里,然后在该目录下分别创建 d1.rs、d4.rs 等文件,我们将路由函数分别移动到各自的文件里,并添加 pub 前缀,然后再在该目录下新建 mod.rs 文件,其内容如下:

mod d1;
mod d4;
mod d5;
// ...

// 重新导出一遍(re-export),这样在外面使用时可以写成 handler::xx 而不是 handler::d1::xx
pub use d1::*;
pub use d4::*;
pub use d5::*;
// ...

接着更新 main.rs:

// ...

// 导入 handler 模块
mod handler;

#[tokio::main]
async fn main() {
    let router = Router::new()
        // 然后就可以使用 handler::xx 的方式调用路由函数
        .route("/1/*nums", get(handler::d1))
        .route("/4/strength", post(handler::d4_1))
        .route("/4/contest", post(handler::d4_2))
}

现在我们的目录结构变得很清爽了,开始进入正题吧。

Part 7

空气中弥漫着刚出炉的饼干(cookies)的香气,圣诞老人也要开始制作饼干了,他的配方(recipe)保存在了网上,但却忘了当时是用什么编码方式保存的。

Task 1

请帮助圣诞老人解开配方的内容!

# 示例输入输出
curl http://localhost:8000/7/decode \
  -H 'Cookie: recipe=eyJmbG91ciI6MTAwLCJjaG9jb2xhdGUgY2hpcHMiOjIwfQ=='

{"flour":100,"chocolate chips":20}

这个任务中涉及到了请求头中 Cookie 的处理,其中 recipe 内容的一眼鉴别为是 base64 编码,那么我们来看看如何处理他们,新增一个路由函数:

Router.route("/7/decode", get(handler::d7_1))

然后在 handler/d7.rs 里编写路由函数(记得要在 mod.rs 里先 mod d7 将其包含进来,然后 pub use d7::* 重新导出一下,后续不再强调此做法):

use axum_extra::extract::CookieJar;
use base64::{engine::general_purpose, Engine};

pub async fn d7_1(jar: CookieJar) -> String {
    let recipe = jar.get("recipe").unwrap();
    let recipe = general_purpose::STANDARD.decode(recipe.value()).unwrap();

    String::from_utf8(recipe).unwrap()
}

上述中使用到了两个新的 crate,因为在 axum 的 0.7+ 版本中不再内置 Cookie 操作,它将多数功能都移到了 axum-extra crate 中,base64 的处理则需要用到 base64 crate,现在我们来安装一下他们:

cargo add axum-extra --features cookie
cargo add base64

在处理 cookie 时,我们使用到了 CookieJar 提取器,调用 get() 来获取 cookie 里某个字段的内容,我们知道每个 cookie 字段除了包含 name 和 value,还包含了过期时间、HttpOnly、Secure 等信息,因此 get() 返回的是一个 struct 而非直接是字符串内容,因此需要进一步调用对应的 value() 可以获取到内容值。

对于 base64 字符串的解码,base64 包提供的 api 看着不是那么直观,读者们可阅读其文档做深入的了解。

最后我们返回了一个字符串内容,使用 cch23-validator 7 验证通过便可。

Task 2

圣诞老人知道了配方内容后,现在要开始做香喷喷的饼干啦!但他不确定做饼干所需的每种原料(ingredient)是否足够分给每个小精灵。现在需要我们帮他计算下他的食品柜(pantry)总共够做多少饼干。

# 示例输入输出
curl http://localhost:8000/7/bake \
  -H 'Cookie: recipe=eyJyZWNpcGUiOnsiZmxvdXIiOjk1LCJzdWdhciI6NTAsImJ1dHRlciI6MzAsImJha2luZyBwb3dkZXIiOjEwLCJjaG9jb2xhdGUgY2hpcHMiOjUwfSwicGFudHJ5Ijp7ImZsb3VyIjozODUsInN1Z2FyIjo1MDcsImJ1dHRlciI6MjEyMiwiYmFraW5nIHBvd2RlciI6ODY1LCJjaG9jb2xhdGUgY2hpcHMiOjQ1N319'

{
  "cookies": 4,
  "pantry": {
    "flour": 5,
    "sugar": 307,
    "butter": 2002,
    "baking powder": 825,
    "chocolate chips": 257
  }
}

这个任务中我们会从 Cookie 里获取到配方所需的原料和食品柜所含有的原料,然后返回能做多少个饼干,以及还剩下多少原料。新增一个路由:

Router.route("/7/bake", get(handler::d7_2))

路由函数:

/// Cookie 解码后的内容
#[derive(Deserialize)]
struct BakeQuantity {
    recipe: Ingredient,
    pantry: Ingredient,
}

/// 原料信息
#[derive(Deserialize, Serialize, Clone, Copy)]
pub struct Ingredient {
    flour: usize,
    sugar: usize,
    butter: usize,
    #[serde(rename = "baking powder")]
    baking_powder: usize,
    #[serde(rename = "chocolate chips")]
    chocolate_chips: usize,
}

/// 响应信息
#[derive(Serialize)]
pub struct BakeResult {
    cookies: usize,
    pantry: Ingredient,
}

/// 运算符重载:两个原料相除(多 / 少),得到最多可做的份数
impl Div for Ingredient {
    type Output = usize;

    fn div(self, rhs: Self) -> Self::Output {
        let flour = self.flour / rhs.flour;
        let sugar = self.sugar / rhs.sugar;
        let butter = self.butter / rhs.butter;
        let baking_powder = self.baking_powder / rhs.baking_powder;
        let chocolate_chips = self.chocolate_chips / rhs.chocolate_chips;
        return flour
            .min(sugar)
            .min(butter)
            .min(baking_powder)
            .min(chocolate_chips);
    }
}

/// 运算符重载:两个原料相减(多 - 少),得到剩余原料
impl Sub for Ingredient {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        let max = self / rhs;

        return Self {
            flour: self.flour - rhs.flour * max,
            sugar: self.sugar - rhs.sugar * max,
            butter: self.butter - rhs.butter * max,
            baking_powder: self.baking_powder - rhs.baking_powder * max,
            chocolate_chips: self.chocolate_chips - rhs.chocolate_chips * max,
        };
    }
}

pub async fn d7_2(jar: CookieJar) -> Json<BakeResult> {
    let recipe = jar.get("recipe").unwrap();
    let recipe = general_purpose::STANDARD.decode(recipe.value()).unwrap();
    let string = String::from_utf8(recipe).unwrap();
    let quantity = serde_json::from_str::<BakeQuantity>(&string).unwrap();

    let max_used_flour = quantity.pantry / quantity.recipe;
    let left = quantity.pantry - quantity.recipe;

    Json::from(BakeResult {
        cookies: max_used_flour,
        pantry: left,
    })
}

这个任务中,我们刻意采用了运算符重载的方式来实现计算逻辑,毕竟本系列以学习为主,多使用不同技巧可以更全面的了解这门语言嘛,其中在实现 trait 所需函数时,记得使用 IDE 的功能来帮助我们自动补全函数前面,而不需要自己一个字母一个字母的敲击。

上述我们给原料 Ingredient 添加了默认的 Clone 和 Copy 的实现,为何呢?

当你移除这两个 trait 时,会发现编译器报错 “move occurs because quantity.pantry has type Ingredient, which does not implement the Copy trait”。从 Div 和 Sub 所要实现的函数签名可以看出,其参数是 self 而非 &self,即会涉及值的移动(value move),所以值使用一次后便不能再使用了,除非新复制一份,可这么做会让我们的代码瞬间变得复杂,由于该任务中的 Ingredient 里的字段类型都是比较 cheap 的,复制一份新的并不会有太大开销,且 rust 已经默认为 usize 实现了 Copy trait,当 struct 的字段都实现了 Copy 那么这个 struct 也可以直接通过设置 derive 来实现 Copy,因此我们便根据报错提示为其添加 Copy trait,添加后又会报错提示还需要实现 Clone trait,故最终添加了这两个 trait。

对于 Clone 和 Copy 的更准确了解大家可查看该 StackOverflow 解答,笔者也会在后续分享个人的一些理解。

Task 3

一些爱玩的小精灵们发现了上述端口后,他们准备要尝试一些创新做法,现在食谱所需原料和食品柜里储藏的原料的种类和个数都不再是固定的了。

curl http://localhost:8000/7/bake \
  -H 'Cookie: recipe=eyJyZWNpcGUiOnsic2xpbWUiOjl9LCJwYW50cnkiOnsiY29iYmxlc3RvbmUiOjY0LCJzdGljayI6IDR9fQ=='

{
  "cookies": 0,
  "pantry": {
    "cobblestone": 64,
    "stick": 4
  }
}

路由仍不变:

Router.route("/7/bake", get(handler::d7_2))

该任务里的解法看着是能够兼容任务 2 的需求,我们注释掉任务 2 的相关内容,来重新实现一套逻辑:

#[derive(Deserialize, Serialize)]
struct Ingredient {
    // 原料内容不再固定,故此处我们使用 HashMap
    inner: HashMap<String, usize>,
}

/// 实现默认值
impl Default for Ingredient {
    fn default() -> Self {
        Self {
            inner: Default::default(),
        }
    }
}

impl Div for &Ingredient {
    type Output = usize;

    // self 是食材柜,rhs 是食谱
    fn div(self, rhs: Self) -> Self::Output {
        let mut res: HashMap<String, usize> = HashMap::new();
        // 根据食谱所需食材进行遍历
        for key in rhs.inner.keys() {
            // 判断食材柜里是否有食谱所需的食材
            match self.inner.get(key) {
                Some(val) => {
                    if *val == 0 {
                        continue;
                    }
                    res.insert(key.to_string(), val / rhs.inner.get(key).unwrap());
                }
                None => {}
            }
        }
        return *res.values().min().unwrap_or(&0);
    }
}

impl Sub for &Ingredient {
    type Output = Ingredient;

    // self 是食材柜,rhs 是食谱
    fn sub(self, rhs: Self) -> Self::Output {
        let max = self / rhs;

        // 一个都做不了,则原样返回食材柜的食材
        if max == 0 {
            return Ingredient {
                inner: self.inner.clone(),
            };
        }

        let mut res: HashMap<String, usize> = self.inner.clone();
        rhs.inner.keys().for_each(|key| match self.inner.get(key) {
            Some(val) => {
                res.insert(key.to_string(), val - rhs.inner.get(key).unwrap() * max);
            }
            None => {}
        });
        return Ingredient { inner: res };
    }
}

#[derive(Deserialize, Debug)]
struct BakeQuantity {
    recipe: HashMap<String, usize>,
    pantry: HashMap<String, usize>,
}

#[derive(Serialize)]
pub struct BakeResult {
    cookies: usize,
    pantry: HashMap<String, usize>,
}

pub async fn d7_3(jar: CookieJar) -> Json<BakeResult> {
    let recipe = jar.get("recipe").unwrap();
    let recipe = general_purpose::STANDARD.decode(recipe.value()).unwrap();
    let quantity = serde_json::from_slice::<BakeQuantity>(&recipe).unwrap();

    let pantry = Ingredient {
        inner: quantity.pantry,
    };
    let recipe = Ingredient {
        inner: quantity.recipe,
    };

    let max_used_flour = &pantry / &recipe;
    let left = &pantry - &recipe;

    Json::from(BakeResult {
        cookies: max_used_flour,
        pantry: left.inner,
    })
}

在本任务的解决过程中,业务复杂度主要在于处理食谱所需原料和储藏柜所含原料的不同场景,两者种类不会一一对应,且有的原料数量会是 0。

我们使用了较多的 HashMap 相关 api,且代码中涉及了 &、*、clone() 的操作,乍一看让人很头大,但在编写过程中通过查看函数签名以及 rust-analyzer 的报错提示,我们可以明确知道是需要引用,还是需要值。

上述的 Ingredient 移除了 Copy 和 Clone trait,且在实现运算符重载时,不像任务 2 中是给 Ingredient 实现,而是给 &Ingredient 实现,为什么呢?

因为 Ingredient 的字段值从简单值(usize)而变成了 HashMap,HashMap 的复制是一项开销较大的操作,它可能会包含很多键值对,rust 也默认没有为其实现 Copy trait,自然我们也没法直接给 Ingredient 添加 Copy,因为我们仅涉及数据的读取计算,并不会修改它们,自然使用引用更为合理。

测试用例编写

目前我们每次测试路由逻辑,都需要运行服务然后手动触发请求,我们来看看如何编写单元测试用例,让调试和验证都更加便捷些。axum 本身没有提供用于测试的相关能力,我们将使用一个非官方封装的 crate axum-test 来简化测试用例的编写:

cargo add axum-test --dev

我们将测试用例都写在各自的路由函数内,现在在 d7.rs 里编写用例:

#[cfg(test)]
mod tests {
    use super::*;
    use axum::{
        http::{header::COOKIE, HeaderValue},
        routing::get,
        Router,
    };
    use axum_test::TestServer;
    use serde_json::{json, Value};

    #[tokio::test]
    async fn d7_3_test() {
        let app = Router::new().route("/7/bake", get(d7_3));
        let server = TestServer::new(app).unwrap();

        let response = server.get("/7/bake").add_header(COOKIE, HeaderValue::from_static("recipe=eyJwYW50cnkiOnsiY29jb2EgYmVhbiI6NSwiY29ybiI6NSwiY3VjdW1iZXIiOjB9LCJyZWNpcGUiOnsiY2hpY2tlbiI6MCwiY29jb2EgYmVhbiI6MX19")).await;

        assert_eq!(
            response.json::<Value>(),
            json!({
              "cookies": 5,
              "pantry": {
                "cocoa bean": 0,
                "corn": 5,
                "cucumber": 0
              }
            })
        );
    }
}

上述使用 axum-test 的 TestServer 创建了一个 mock 服务器,然后调用其 get/post() 模拟请求,该服务器默认没有真正使用到 tcp 服务,因此也无需担心端口冲突问题,当然其提供了配置项来配置使用真实的 tcp 服务来发送和接收请求,读者们可阅读其文档做进一步了解。

结尾

继续冲鸭!