Golang 编写 MySQL UDF

一、MySQL UDF

这玩意全称 “MySQL user-definable function”, 从名字就可以看出来叫 “用户定义的方法”; 那么 UDF 到底是干啥的呢?

简单一句话说就是说: 你可以自己写点代码处理数据, 然后把这段代码编译成动态链接库(so), 最后在 MySQL 中动态加载后用户就可以用了.

二、解决方案

由于要检查数据库, 但是实际上审查并不会关注每个表甚至数据库细节; 所以想到最简单的方案就是在读取和写入时通过 UDF 定义一个 SM4 的加密算法把数据动态加密和解密, 关于其他细节这里不做详细说明, 本文主要阐述如何用 Go 搓一个简单的 UDF 并使用.

三、UDF 方法

由于 UDF 官方支持是 C/C++, 所以在 Go 中需要使用 CGO; 一个 UDF 实现通常包含两个 func:

func xxx_init(initid *C.UDF_INIT, args *C.UDF_ARGS, message *C.char) C.int {
    // ... 逻辑实现
}

func xxx(initid *C.UDF_INIT, args *C.UDF_ARGS, result *C.char, length *C.ulong, is_null *C.char, error *C.char) *C.char {
    // ...逻辑实现
}

其中 xxx_init 用于预检查, xxx 作为真正的逻辑实现; 当 xxx 方法被调用之前会先通过 xxx_init 方法做一次参数、内存分配等预处理.

注意: 从 MySQL 8.0.1 开始 xxx_init 的返回值从 my_bool 变更为 int, 网上很多代码写 my_bool 的会导致无法通过编译; 具体参考 https://bugs.mysql.com/bug.php?id=85131

四、Go 实现 UDF

知道了方法签名以后, 就不多废话直接上代码实现:

package main

// #include <stdio.h>
// #include <sys/types.h>
// #include <sys/stat.h>
// #include <stdlib.h>
// #include <string.h>
// #include <mysql.h>
// #cgo CFLAGS: -D ENVIRONMENT=0 -I/usr/include/mysql -fno-omit-frame-pointer
import "C"

import (
	"github.com/tjfoc/gmsm/sm4"
)

//export xsm4_enc_init
func xsm4_enc_init(initid *C.UDF_INIT, args *C.UDF_ARGS, message *C.char) C.int {
	if args.arg_count != 1 {
		msg := "xsm4_enc() 仅支持单个字符串参数\n"
		C.strcpy(message, C.CString(msg))
		return 1
	}

	return 0
}

//export xsm4_enc
func xsm4_enc(initid *C.UDF_INIT, args *C.UDF_ARGS, result *C.char, length *C.ulong, is_null *C.char, error *C.char) *C.char {
	// 将 C 的指针转换为 Go 的类型
	var str = C.GoString(*args.args)

	var resp string
	enc, err := sm4.Sm4Ecb([]byte("1234567890abcdef"), []byte(str), true)
	if err != nil {
		resp = err.Error()
	} else {
		resp = string(enc)
	}

	// 将结果转换为 C 的类型
	var res = C.CString(resp)

	// 设置输出参数
	*length = C.ulong(len(resp))
	*is_null = 0

	// 返回结果
	return res
}

func main() {}

xsm4_enc_init 方法做一下检查, 当前只支持单个字段参数, xsm4_enc 通过开源的 gmsm 库对传入的字段进行简单的 SM4 加密并返回; 在真实环境中需要调用加密机来实现相关加密, 这里只演示直接使用开源库+固定密码.

五、编译并加载

将上面的代码保存为 xsm4_enc.go, 然后在安装有 MySQL 头文件的的服务器上使用以下命令编译:

go build -o xsm4_enc.so -buildmode=c-shared xsm4_enc.go

如果没问题将会生成一个 xsm4_enc.so 文件, 如果提示 C.xxx 类型没找到等问题说明头文件没有加载, 自行检查或修改 -I/usr/include/mysql 位置.

生成好 so 文件以后将其复制到 MySQL 的插件目录(插件目录可通过 SHOW VARIABLES LIKE 'plugin_dir'; 查询到):

cp xsm4_enc.so /usr/lib/mysql/plugin/

最后在 MySQL 中创建 UDF:

# 创建
CREATE FUNCTION xsm4_enc RETURNS STRING SONAME 'xsm4_enc.so';

# 删除
DROP FUNCTION xsm4_enc;

六、UDF 使用

使用就简单了, 在查询的时候直接把你的 func 名称写上就行:

SELECT id, xsm4_enc(username), username FROM users;

同理也可以创建一个解密 UDF, 当然这些 UDF 最终配合视图啥的做啥、怎么用就不做过多赘述了.