教你在 Pisa-Proxy 中，如何利用 Rust 实现 MySQL

一、前言

背景

在 Database Mesh 中，Pisanix 是一套以数据库为中心的治理框架，为用户提供了诸多治理能力，例如：数据库流量治理，SQL 防火墙，负载均衡和审计等。在 Pisanix 中，Pisa-Proxy 是作为整个 Database Mesh 实现中数据平面的核心组件。Pisa-Proxy 服务本身需要具备 MySQL 协议感知，理解 SQL 语句，能对后端代理的数据库做一些特定的策略，SQL 并发控制和断路等功能。在这诸多特性当中，能够理解 MySQL 协议就尤为重要，本篇将主要介绍 MySQL 协议和在 Pisa-Proxy 中 MySQL 协议的 Rust 实现。

Why Rust

为什么要选用 Rust 语言呢？我们的考量有以下几个必要条件。

• 安全性：首先作为数据库治理的核心组件，其语言的安全性是居首位的。Rust 中，类型安全实现内存安全，如所有权机制、借用、生命周期等特性避免了程序开发过程中的空指针、悬垂指针等问题，从而保证了服务在语言层面的安全性。
• 优秀的性能表现：Rust 的目标在性能方面对标 C 语言，但在安全和生产力方面则比 C 更胜一筹。其无 GC，不需要开发人员手动分配内存等特性，极大程度地减少内存碎片，简化内存管理。
• 低开销：从开发效率和可读可维护性上来说，有足够的抽象能力，并且这种抽象没有运行时开销（runtime cost）。零开销抽象，通过泛型和 Trait 在编译期展开并完成抽象解释。
• 实用性：有优秀的包管理器工具 Crate、文档注释支持、详细的编译器提示、友好的错误处理等，在开发过程中能够高效帮助程序员快速开发出可靠、高性能的应用。

二、整体架构，模块设计

整体架构
如图 1，代理服务包含服务端和客户端的协议解析、SQL 解析、访问控制、连接池等模块。

工作流程
在图 1 中我们可以看出整个 Proxy 服务可以概括为以下几个阶段。

1. 首先 Pisa-Proxy 支持 MySQL 协议，将自己伪装为数据库服务端，应用连接配置只需修改访问地址即可建连 Pisa-Proxy 通过读取应用发来的握手请求和数据包；
2. 得到应用发来的 SQL 语句后对该 SQL 进行语法解析，并得到该 SQL 的 AST；
3. 得到对应 AST 后，基于 AST 实现高级访问控制和 SQL 防火墙能力；
4. 访问控制和防火墙通过后 SQL 提交执行，SQL 执行阶段的指标将采集为 Prometheus Metrics，最后根据负载均衡策略获取执行该语句的后端数据库连接；
5. 如果连接池为空将根据配置建立和后端数据库的连接，SQL 将从该连接发往后端数据库；
6. 最后读取 SQL 执行结果，组装后返回给客户端。

三、如何用 Rust 快速实现 MySQL 代理服务

如图 2，在整个代理服务中总体分为两部分：服务端和客户端，即代理服务作为服务端处理来自客户端的请求。和服务端，需要对服务端发起认证，并将客户端端命令发送给 MySQL 数据库。在这两部分中我们需要创建两套 Socket 来完成网络数据包的处理。

技术选型

介绍

在网络报处理和运行时处理上，我们选用了 Rust 实现的 Tokio（https://github.com/tokio-rs/tokio）框架。Tokio 框架是 Rust 编写的可靠、异步和精简应用程序的运行时。并且 Tokio 是一个事件驱动的非阻塞 I/O 平台，用于使用 Rust 编程语言编写异步应用程序。在高层次上，它提供了几个主要组件：

• 一个多线程、基于工作窃取的任务调度程序（https://docs.rs/tokio/latest/tokio/runtime/index.html）。
• 由操作系统的事件队列（epoll、kqueue、IOCP 等）支持的响应式编程。
• 异步 TCP 和 UDP（https://docs.rs/tokio/latest/tokio/net/index.html）套接字。

同时，Tokio 还提供了丰富的工具链，例如编解码工具包、分帧包等等，可以使我们更加方便快捷地处理 MySQL 中各种各样的数据报文。

项目地址: https://github.com/tokio-rs/tokio

优势

a. 快速：Tokio 的设计旨在使应用程序尽可能都快。

b. 零成本抽象：Tokio 以 Future 为基础。虽然 Future 并非 Rust 独创，但与其他语言的 Future 不同的是，Tokio Future 编译成了状态机，用 Future 实现常见的同步，不会增加额外开销成本。Tokio 的非阻塞 IO 可以充分发挥系统优势，例如实现类似 Linux Epoll 这种多路复用技术，在单个线程上的多路复用允许套接字并批量接收操作系统消息，从而减少系统调用，所有这些都可以减少应用程序的开销。

c. 可靠：Rust 的所有权模型和类型系统可以实现系统级应用程序，而不必担心内存不安全。

d. 轻量级：没有垃圾收集器，因为 Tokio 是基于 Rust 构建的，所以编译后的可执行文件包含最少的语言运行时。这意味着，没有垃圾收集器，没有虚拟机，没有 JIT 编译，也没有堆栈操作。这样在编写多线程并发的系统时，能够有效避免阻塞。

e. 模块化：每个组件都位于一个单独的库中。如果需要，应用程序可以挑选所需的组件，避免依赖不需要的组件。

代码实现

Rust 中数据包处理

#[derive(Debug)]
pub struct Packet {
    pub sequence: u8,
    pub conn: BufStream<LocalStream>,
    pub header: [u8; 4],
}

以上为 Proxy 数据包处理逻辑中核心的结构体，结构体中包含了三个字段分别为：

• sequence：报文中的序列号 ID
• conn：处理链接的 Socket
• header：存储消息头报文的数组，长度为 4 字节

在包处理逻辑中主要定义了以下函数，在整个代理服务中网络数据交换都由以下方法来完成。

Pisa-Proxy 作为服务端

pub struct Connection {
    salt: Vec<u8>,
    status: u16,
    collation: CollationId,
    capability: u32,
    connection_id: u32,
    _charset: String,
    user: String,
    password: String,
    auth_data: BytesMut,
    pub auth_plugin_name: String,
    pub db: String,
    pub affected_rows: i64,
    pub pkt: Packet,
}

上面的结构体描述了 Pisa- Proxy 作为服务端处理来自于客户端请求时所包含的字段。例如，其中包含了和 MySQL 客户端进行认证信息，和所包含数据包处理逻辑的 Packet。

Pisa-Proxy 作为客户端

#[derive(Debug, Default)]
pub struct ClientConn {
    pub framed: Option<Box<ClientCodec>>,
    pub auth_info: Option<ClientAuth>,
    user: String,
    password: String,
    endpoint: String,
}

Tokio 提供的编解码器

在 Pisa-Proxy 中，大量使用了 Tokio 工具包中的编解码器，使用 codec Rust 会自动帮助开发者将原始字节转化为 Rust 的数据类型，方便开发者处理数据。使用编解码器，只需要在代码中为定义好的类型实现 Decoder 和 Encoder 两个 Trait，就可以通过 stream 和 sink 进行数据的读写。下面通过一个简单的示例来看一下 Tokio 编解码器的使用步骤。

使用 Tokio 编解码器一共分为三步：

1. 首先要自定义一个错误类型，这里定义了一个 ProtocolError，并为它实现一个 from 方法，能够让它接收错误处理。

pub enum ProtocolError {
    Io(io::Error),
}

impl From<io::Error> for ProtocolError {
    fn from(err: io::Error) -> Self {
        ProtocolError::Io(err)
    }
}

2. 定义一个数据类型，这里我们声明一个 Message 为 String，然后定义一个结构体，也就是我们要解析的是原始字节流要实际转换成的结构体，也即 Tokio 中 framed（帧的概念），这里定义一个空结构体 PisaProxy；

type Message = String;
struct PisaProxy;

接下来就是为 PisaProxy 分别实现 Encoder 和 Decoder 这两个 Trait。这里的示例，实现的功能为将数据转为 byte 数组，追加到 buf 中。在编码器中，我们首先要指定 Item 类型为 Message 和错误类型，编码处理逻辑这里将字符串进行拼接并返回给客户端。

在这里，Encoder 编码是指将用户自定义类型转换成 BytesMut 类型，写到 TcpStream 中，Decoder 解码指将读到的字节数据序列化为 Rust 的结构体。

impl Encoder<Message> for PisaProxy {
    type Error = ProtocolError;

    fn encode(&mut self, item: Message, dst: &mut BytesMut) ->Result<(), Self::Error> {
        dst.extend(item.as_bytes());
        Ok(())
    }
}

impl Decoder for PisaProxy {
    type Item = Message;
    type Error = ProtocolError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        if src.is_empty() {
            return Ok(None);
        }
        let data = src.split();
        let mut buf = BytesMut::from(&b"hello:"[..]);
        buf.extend_from_slice(&data[..]);
        let data = String::from_utf8_lossy(&buf[..]).to_string();

        Ok(Some(data))
    }
}

3. 当实例化 PisaProxy 结构体后，就可以调用 framed 方法，codec 的 framed 方法（codec.framed(socket)）将 TcpStream 转换为 Framed<TcpStream，PisaProxy>，这个 Framed 就是实现了 tokio 中的 Stream 和 Sink 这两个 trait，实现的这两个 Trait 的实例就具有了接收（通过 Stream）和发送（通过 Sink）数据的功能，这样我们就可以调用 send 方法发送数据了。

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let addr = "127.0.0.1:9088";
    let listener = TcpListener::bind(addr).await?;
    println!("listen on: {:?}", addr);

    loop {
        let (socket, addr) = listener.accept().await?;

        println!("accepted connect from: {}", addr);

        tokio::spawn(async move {
            let codec = PisaProxy {};
            let mut conn = codec.framed(socket);
            loop {
                match conn.next().await {
                    Some(Ok(None)) => println("waiting for data..."),
                    Some(Ok(data)) => {
                        println!("data {:?}", data);
                        conn.send(data).await;
                    },
                    Some(Err(e)) => {
                    },
                    None => {},
                }
            }
        });
    }
}

四、MySQL 协议在 Pisa-Proxy 中的实现

MySQL 协议简介
MySQL 数据库本身是一个很典型的 C/S 结构的服务，客户端和服务端通信可以通过 Tcp 和 Unix Socket 的方式进行交互。在本篇中，我们主要说明通过网络 Tcp 的方式来实现 MySQL 代理。

MySQL 客户端和服务端的交互主要包含了两个重要的过程：1. 握手认证，2.发送命令。本篇会主要围绕这两个过程来介绍相关的实现。在客户端和服务端交互的过程中主要包含了以下几种类型报文：数据包、数据结束包、成功报告包以及错误消息包，在后面的章节中会为大家详细介绍这几种报文。

交互过程
MySQL 客户端和服务端在交互的过程中主要包含了两个过程，即握手认证阶段和执行命令阶段，当然在这两个过程之前首先要经历 TCP 三次握手的过程。在三次握手结束后首先进入握手认证阶段，在交换完信息并且客户端正确登录服务端后，进入执行命令阶段，图 3 完整描述了整个交互过程。

代码链接：https://github.com/database-mesh/pisanix/blob/master/pisa-proxy/protocol/mysql/src/server/conn.rs

握手认证

在握手认证阶段是 MySQL 客户端和服务端建联非常重要的阶段，该阶段发生在 TCP 三次握手之后。首先服务端会给客户端发送服务端信息，其中包括协议版本号、服务版本信息、挑战随机数、权能标志位等等。当客户端接收到服务端发来的响应之后，客户端开始发起认证请求。认证请求中，会携带客户端用户名、数据库名以及通过服务端响应中的挑战随机数将客户端密码加密后，一同发送给服务端进行校验。校验过程中，除了会对用户名密码进行校验，还会匹配客户端所使用的认证插件，如果不匹配则会发生插件的自动切换，以及判断客户端是否使用了加密链接。当以上阶段全部正常完成后，客户端则登录成功，服务端返回客户端 OK 数据报文。

上述过程分别在 runtime 和 protocol 的 server 中实现。在 runtime 中的 start 函数等待请求进入，Tcp 三次握手结束后，从 handshake 函数如图 3，开始握手阶段。在 handshake 中分别包含了三个过程，首先由 write_initial_handshake 给客户端发送服务端信息，然后在 read_handshake_response 客户端拿着服务端信息开始认证请求。

pub async fn handshake(&mut self) -> Result<(), ProtocolError> {
        match self.write_initial_handshake().await {
            Err(err) => return Err(err::ProtocolError::Io(err)),
            Ok(_) => debug!("it is ok"),
        }

        match self.read_handshake_response().await {
            Err(err) => {
                return Err(err);
            }
            Ok(_) => {
                self.pkt.write_ok().await?;
                debug!("handshake response ok")
            }
        }

        self.pkt.sequence = 0;

        Ok(())
    }

执行命令

当握手和认证阶段完成后，此时客户端才算是真正意义上的与服务端完成建立链接。那么此时则进入执行命令阶段。在 MySQL 中，能够发送命令的指令类型有很多种，我们会在下文中为大家进行介绍。

代码链接: https://github.com/database-mesh/pisanix/blob/master/pisa-proxy/runtime/mysql/src/server/server.rs

在下面的代码中可以看到，Pisa-Proxy 在这里会对不同类型的指令进行不同的逻辑处理。例如，初始化 db 的处理逻辑为 handle_init_db，处理查询的逻辑为 handle_query。

match cmd {
          COM_INIT_DB => self.handle_init_db(&payload, true).await,
          COM_QUERY => self.handle_query(&payload).await,
          COM_FIELD_LIST => self.handle_field_list(&payload).await,
          COM_QUIT => self.handle_quit().await,
          COM_PING => self.handle_ok().await,
          COM_STMT_PREPARE => self.handle_prepare(&payload).await,
          COM_STMT_EXECUTE => self.handle_execute(&payload).await,
          COM_STMT_CLOSE => self.handle_stmt_close(&payload).await,
          COM_STMT_RESET => self.handle_ok().await,
          _ => self.handle_err(format!("command {} not support", cmd)).await,
}

MySQL 协议基本数据类型
在 MySQL 协议中主要有以下几种数据类型：

• 整型值：MySQL 报文中整型值分别有 1、2、3、4、8 字节长度，使用小端传输。
• 字符串（以 NULL 结尾）（Null-Terminated String）：字符串长度不固定，当遇到’NULL’（0x00）字符时结束。
• 二进制数据（长度编码）（Length Coded Binary）

参考链接：https://dev.mysql.com/doc/internals/en/basic-types.html

报文结构
在 MySQL 客户端和服务端所交互的数据最大长度为 16MByte，其基本数据报文结构类型如下：

图 4

图 5

在图 4 和图 5 中描述了 MySQL 报文的基本结构。报文包括了两部分，消息头和消息体。其中在消息头中，3 字节表示数据报文长度，1 字节存储序列号，消息体中存储实际的报文数据。

消息头

用于标记当前请求消息的实际数据长度值，以字节为单位，占用 3 个字节，最大值为 0xFFFFFF，即接 2^24-1。

序列号

序列 ID 从 0 开始随每个数据包递增，并且当进入新的命令时重置为 0。序列号 ID 有可能会发生回绕，当发生回绕时，需将序列号 ID 重置为 0，并且重新开始计数递增。

报文数据

消息体用于存放请求的内容及响应的数据，长度由消息头中的长度值决定。

客户端请求命令报文

该指令用于标识客户所要执行命令的类型，以字节为单位占用 1 个字节。请求命令报文常用到的有 Text 文本协议和 Binary 二进制协议。这里可以参考【执行命令】中的代码，代码中描述了对不同指令的处理逻辑。

在文本协议中常用到的有以下指令：

更多请参考：https://dev.mysql.com/doc/internals/en/text-protocol.html

在二进制协议，即 Prepare Statement 中常用到的有以下指令：

更多请参考：https://dev.mysql.com/doc/internals/en/prepared-statements.html

响应报文

• OK 响应报文
• Error 响应报文
• Field 结构
• EOF 结构
• Row Data 结构

…

例如，下面代码中展示了如何给客户端写 OK 和 EOF 报文。

#[inline]
pub async fn write_ok(&mut self) -> Result<(), Error> {
  let mut data = [7, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0];
  self.set_seq_id(&mut data);
  self.write_buf(&data).await
}
  
#[inline]
pub async fn write_eof(&mut self) -> Result<(), Error> {
  let mut eof = [5, 0, 0, 0, 0xfe, 0, 0, 2, 0];
  self.set_seq_id(&mut eof);
  self.write_buf(&eof).await
}

com_query 请求流程

如图 6 为 com_query 的请求流程，com_query 指令可能会返回以下结果集：

• ERR 报文
• OK 报文
• Protocol::LOCAL_INFILE_Request 报文
• Resultset 报文

在 (https://github.com/database-mesh/pisanix/blob/master/pisa-proxy/protocol/mysql/src/client/resultset.rs) 文件中主要为对 ResultSet 结果集的处理，可以看到在这里定义了 ResultSet 结构，同样对 ResultSetCodec 分别实现了 Encoder 和 Decoder 两个 Trait，这样就可以通过编解码器来处理 ResultSet 的报文。

#[derive(Debug, Default)]
pub struct ResultsetCodec {
  pub next_state: DecodeResultsetState,
  pub col: u64,
  pub is_binary: bool,
  pub seq: u8,
  pub auth_info: Option<ClientAuth>,
}

图 6

五、总结

以上，就是本篇文章的全部内容。在本篇文章中我们介绍了使用 Rust 实现 MySQL 代理的动机，介绍了 MySQL 协议中一些常用到的概念和 MySQL 中数据报文在网络中是如何进行交换数据的；并在最后介绍如何使用 Rust 去快速实现一个 MySQL 代理服务。更多实现细节请关注 Pisanix 代码仓库。

欢迎点击链接查看相关教学视频：https://www.bilibili.com/video/BV1B54y1o7EC?spm_id_from=333.999.0.0

六、相关链接

Pisanix
项目地址：https://github.com/database-mesh/pisanix

官网地址：https://www.pisanix.io/

Database Mesh：https://www.database-mesh.io/

Mini-Proxy：一个最小化的 MySQL Rust 代理实现

项目地址：https://github.com/wbtlb/mini-proxy

社区
目前 Pisanix 社区每两周都会组织线上讨论，详细安排如下，欢迎各位小伙伴一起参与进来。