做个简单的 reverse proxy

在 Tubi，我们会对客户端软件做 End-to-end testing（以下简称 E2ET）。E2ET 如果只测试 UI 和 UI 相关的逻辑，有足够多的开源工具和商业工具，不过，如果要同时要验证网络层的输入输出是否符合预期，比如在展示某个 UI 的同时，发送了相应的 analytics event，这就需要很多额外的工作。之前我们在客户端做网络层的 interceptor，虽然能工作，但它有一些局限：1) 测试脚本需要能访问到客户端拦截并记录下来的网络请求和响应，这意味着记录的内容需要上传至某个特定的，可以公开访问的位置，比如 s3。2) 所有的客户端都需要做类似的 interceptor 的工作。所以考虑再三，我们决定尝试做个独立的 proxy server，让客户端通过访问这个 proxy server，来访问实际的 API。这样，我们可以在这个 proxy server 里做很多事情，比如 traffic log，比如通过一系列规则来做 traffic rewrite，可以对 API 访问做缓存，甚至可以把某些请求写入 sinkhole。

做这样一个 proxy server，更确切地说，一个 reverse proxy，首先考虑的是，可不可以直接利用已有的 nginx server。nginx server 可以满足 proxy 的需求，但会带来两个问题：1) nginx 承载很多服务，我们不希望这个简单的 proxy 影响 prod/staging 的业务。2) 虽然我们可以在 nginx 中通过 lua 脚本拿到请求和响应并将其记录下来，但我们还需要启动一个 server 来对外提供记录的内容。因而，即便 nginx 替我们做了 proxy，我们还是需要单独做个服务，这无法避免。

思来想去，我决定做个简单的 POC，尝试用 Rust 来做个独立的工具，来满足 E2ET 的需求。同时，我希望这个工具还可以某种程度帮助客户端开发者更好地在本地开发。于是，我和同事一起大概构思了这样一个架构：

测试脚本可以驱动客户端运行，客户端在发送网络请求时，会附带一个额外的 session-id 的头。这个 session-id 作为唯一标识，可以用于之后获取数据。测试脚本可以监听某个 session-id 下的所有网络请求，也可以请求满足特定条件的网络请求。测试脚本还可以动态发送一些规则，来修改某个请求的响应，比如某些情况下需要测试用户 token 过期的逻辑，此时规则可以要求下一个请求，服务器返回 403。

对于 proxy server，它可以配置 proxy 运行在哪个端口，收到的请求该怎么处理：1) 如果 action_type 是 Forward，则做 proxy，2) 如果 action_type 是 Load，则从 本地加载返回：

server:
  port: 8080
  db_uri: "sqlite://./data/traffic.db"
proxies:
  - addr: "0.0.0.0:8081"
    cache: true
    actions:
      - action_type: Forward
        route: /*path
        dst: "https://api1.awesome-server.com"
  - addr: "0.0.0.0:8082"
    cache: true
    actions:
      - action_type: Forward
        route: /*path
        dst: "https://api2.awesome-server.com"
      - action_type: Load
        route: /newly/added/api
        dst: /path/to/api.json
        content_type: "application/json"

Load 可以用于当上游的 API 尚未实现的场景，此时，通过为 proxy server 提供样例数据，可以让 proxy server mock 服务端的行为，这样，客户端依旧可以测试未实现的 API 行为。

对于 Forward，proxy server 会把请求原封不动发送给目标服务器，同时把请求和响应记录下来，并且 1) publish 出去，如果有任何客户端 subscribe 对应的 session-id，客户端会收到这个记录；2) 把记录写入到 sqlite DB 中，以便日后查询之用。

好，说了这么多背景信息，我们进入正题，讲讲我在做这个 reverse proxy 的一点心得。

首先，这样一个 proxy，性能并不是最重要的，请求和响应数据的记录更为重要。所以我使用了 axum + reqwest + sqlite 的组合。axum 来提供 HTTP server，然后把收到的请求转给 reqwest 发送，reqwest 收到的响应，再转回给客户端。同时，我们可以把需要记录的请求和响应用 RequestInfo / ResponseInfo 封装起来，统一转换成一个 TrafficLog 结构。在 proxy 的上下文中，我们不直接写 sqlite，而是通过一个 mpsc channel 把数据转送出去，在另一个线程下接收并写入数据库：

因为我们不需要一个像 nginx 那样高性能的服务，所以在 proxy server 中，我们需要把请求和响应的 body 完全读取出来，便于记录。这里就遇到了第一个坑：因为我们的客户端 app 有可能在请求时允许 gzip 或者其他压缩方式，当我们把客户端的 headers 都透传给服务器时，服务器就有可能返回压缩过的数据，我们直接读到的 body 的内容就是压缩后的内容，这个内容发送给客户端没有问题，客户端会自动解压，但如果记录下来，需要先解压，否则记录的内容无法根据 content-type 来处理。本来我想通过客户端 app 发送的 request header 中的 “accept-encoding” 来确定如何解压，后来发现 reqwest 提供了自动解压的能力，我们可以在构建 HTTP Client 时，使用 Client::builder().gzip(true) 来允许自动解压，如下所示：

第二个坑是如何 decode request/response body。HTTP 协议设计的非常灵活，需要根据 request/response header 的 content-type 里的 charset 来决定如何处理对应的 body。在阅读了 reqwest 的源码后，我发现 Rust 有个 encoding_rs 库，可以帮忙处理这类问题：

第三个坑，或者说心得，是如何比较舒服地在若干种可能中找到最优匹配的 session-id。我们的 proxy 要尽量容错，比如允许客户端把 session-id 放在 header 或者 query 中。我定义的优先级是这样的：

1. 如果 header 中有 session-id，则使用之；
2. 否则，如果 header 中有 session_id，则使用之；
3. 否则，如果 query 中有 session_id，则使用之；
4. 否则，返回 None

这样的场景在开发中并不少见。如果我们使用 if-else 实现，非常啰嗦，且以后很难修改，所以，定义规则，解析并处理规则是更好的方式。在弱类型语言下，这样的规则定义起来很简单，放在一个数组或者列表中即可，然后使用类似 reduce_while 的语义统一处理即可。但在 Rust 下，我们无法很轻松地生成这样的规则，因为不同数据源的类型可能不同。比如 headers 是 HeaderMap 类型，而 query 是 HashMap 类型，所以我们需要先将它们统一成 HashMap，然后再定义规则，用 itertools 的 fold_while（类似 reduce_while）处理。这样代码非常简单可扩展，以后我们修改 session-id 的来源，只需要更新 options 即可，甚至，我们还可以把这样的规则写入配置文件中，更加灵活地处理。代码如下：

后记

这个 proxy server 我在圣诞节前一两天开始开发（难得没什么会议），圣诞假期晚上哄娃睡觉后抽空陆陆续续写了一些，今早出门前玩雪前基本功能完工，写了有 1000 行代码。本来只想做个 POC，验证了思路后后续的事情交给同事来完成，结果一不小心写 high 了，POC 变成了一个相对完善的，可以开箱即用的工具（我连 CLI 都写好了）。不过，也正因为一开始冲着 POC 来写的，这个项目也成了我少有的没有 unit testing 的项目，我需要抑制住在节日期间继续做功能的冲动，争取在元旦前把 unit testing 和代码重构做完。

嗯，项目的名字叫：wormhole，暂时还没开源，也许以后会开源。