PostgreSQL 练习 ---- psql 新增连接参数

目标

添加一个连接参数，默认为 false 。当 psql 连接时，若该连接参数非 “true” 时，用户 “u1“ 对表对象无操作权限，包括自己拥有的表。

连接机制简介

连接过程如下所述：

• 客户端初始化一个空连接，设置连接状态为 CONNECTION_NEEDED 。
• 然后创建 socket，与服务端建立 tcp 连接，状态设置为 CONNECTION_STARTED ，检查 socket 连接是否正常。
• 初始化连接中的 socket 地址，修改连接状态为 CONNECTED_MODE 。
• 给服务端发送协议号等信息，发送启动包给服务端后，设置状态为 CONNECTION_AWAITING_RESPONSE 。
• 处理身份验证交换，等待 postmaster 消息并根据需要作出响应，最终设置状态为 CONNECTION_AUTH_OK。(后续过程再次不再细述，有兴趣可以自己查资料进行学习)。

我们重点关注启动包，因为参数的发送、解析都是通过启动包完成。启动包不止包含了连接协议，还包含了些连接选项等信息( dbname、user 等)。此次连接参数的发送解析就在启动包中。

实现

此处以添加 conninfo 为例：

确定标志为 priv_flag，默认为 false ，并设置为全局变量。确定连接参数为 char* 型的变量 privilege_flag 。

定义全局变量 priv_flag

在 src/backend/utils/init/globals.c 文件中定义

bool priv_flag = false;

并在 src/include/miscadmin.h 中声明此变量

extern bool priv_flag;

在 Port 和 conn 结构体中添加字段

主要目的是在建立连接的过程中，对该参数的赋值发送和接受，进行处理。pg_conn 为客户端存储连接状态结构体，Port 为服务端相关联结构体。

在 src/include/libpq/libpq-be.h 文件的 Port 结构体中添加 privilege_flag 字段：

char        *privilege_flag;

src/interfaces/libpq/libpq-int.h 文件中 pg_conn 结构体同理添加：

char        *privilege_flag;

在 conninfo 中添加 privilege_flag

psql 首先会对命令行 -选项 的命令进行解析，未匹配的“”引号里面的信息，会在客户端填充第一包数据之前，进行解析，此时依赖的就是 conninfo 结构体，去进行匹配。

在 src/interfaces/libpq/fe-connect.c 文件中定义的结构体 PQconninfoOptions 中添加新添参数信息。

internalPQconninfoOption 结构体定义 conninfo 参数及其回退资源。

添加新参数信息至倒数第二个成员，因最后一个为全NULL，遍历时到此停止。

	{"privilege_flag", "PRIVILEGE_FLAG", NULL, NULL,
		"Privilege-Flag", "", 8, /* sizeof("true") = 5 or sizeof("false") = 6 */
	offserof(struct pg_conn, privilege_flag)},

用以参数解析时进行匹配。

在客户端填充启动包数据

在 src/interfaces/libpq/fe-protocol3.c 文件 build_startup_packet 函数中添加相关代码。对 privilege_flag 参数的值进行填充。

	if (conn->privilege_flag && conn->privilege_flag[0])
		ADD_STARTUP_OPTION("privilege_flag", conn->privilege_flag);

服务端解析启动包

在 src/backend/postmaster/postmaster.c 文件的 ProcessStartupPacket 函数对第一包数据进行解析。

当读取到 privilege_flag 参数，且值为 true 时，置全局变量 priv_flag 为 true ，否则继续使用默认值。

	else if(strcmp(nameptr, "privilege_flag") == 0)
	{
		port->privilege_flag = pstrdup(valptr);
		if(strcmp("true",port->privilege_flag) == 0)
		{
			priv_flag = true;
		}
	}

设置权限检查判断

在 src/backend/catalog/aclchk.c 文件的 pg_class_aclmask 函数中，对权限检查增加逻辑判断。

当前用户为“u1”，且全局变量 priv_flag 为 false 时，在权限检查这块“去除”该用户对表对象的 insert/update/delete/truncate 权限(前提是 u1 拥有该表的这些权限)。即通过 privilege_flag 连接参数控制 u1 用户的权限。

	if(strcmp(GetUserNameFromId(roleid, false), "u1") == 0 && !priv_flag)
	{
		mask &= ~(ACL_INSERT | ACL_UPDATE | ACL_DELETE | ACL_TRUNCATE);
		ReleaseSyscache(tuple);
		return mask;
	}

测试验证

1. 首先使用 postgres 用户登录数据库，创建 u1 用户

create user u1 with password '1qaz!QAZ';

2. 然后创建 public 表 test ，并赋其所有权限给 u1 用户

create table public.test (id int);
insert into test values (1),(2),(3),(4),(5);
grant all ON test to u1 ;

3. 通过 pg_class 表查看 u1 用户对表 test 的权限
   

select relname,relacl from pg_class where relname = 'test';

4. 使用 u1 用户进行登录。因当使用 -d、-U 参数时，“”里面的参数就会被忽略，所以使用全 conninfo 的方式进行登录。
   首先不使用 privilege_flag 参数进行连接，进行 select/insert 操作
   
   可以看到 select 权限正常，insert 操作提示没有权限。
   接下来使用 privilege_flag=true 进行连接，再次进行 select/insert 操作：
   可以看到，此时 insert 操作正常。