Go Decorator With Reflect

本部分将通过用例的方式, 展示Reflect构造Go装饰器的用例.

背景

在云环境中存在若干租户(Project)与若干虚拟机(VM), 往往一个Project名下存在多个VM, 并且Project仅能够操作自己名下的VM.

代码示例如下:

租户与虚拟机Model示例:

package cloud

type VM struct {
	ID        string
	Name      string
	CPU       string
	ProjectID string
}

type Project struct {
	ID       string
	Name     string
	AZID     string
	ParentID string
}

云环境内资源信息:

package cloud
// ProjectToVMMap 租户与虚拟机的所属关系
var ProjectToVMMap = map[string]string{
	"A": "001",
	"B": "002",
}
// GlobleVMs 全局虚拟机信息
var GlobleVMs map[string]VM

工具类方法:

package cloud

import "fmt"

// PrintVMs 输出全局虚拟机信息
func PrintVMs() {
	fmt.Printf("VMs num: %v\n", len(GlobleVMs))
	for _, vm := range GlobleVMs {
		fmt.Printf("VM info: %+v \n", vm)
	}
	fmt.Printf("\n")
}

// InitVMs 初始化全局虚拟机信息
func InitVMs() {
	GlobleVMs = map[string]VM{
		"001": {
			ID:        "001",
			Name:      "vm_001",
			CPU:       "x86",
			ProjectID: "A",
		},
		"002": {
			ID:        "002",
			Name:      "vm_002",
			CPU:       "arm",
			ProjectID: "B",
		},
	}
}

代码中提供两个接口, 支持Project操作VM:

修改虚拟机名称:

package cloud

import "fmt"

type PutVMParams struct {
	ProjectID string
	VMID      string
	VMName    string
}

func PutVM(params PutVMParams) error {
	vm, ok := GlobleVMs[params.VMID]
	if !ok {
		fmt.Printf("VMID: %v missing \n", params.VMID)
		return fmt.Errorf("VMID MISSING ERROR")
	}
	vm.Name = params.VMName
	GlobleVMs[params.VMID] = vm
	return nil
}

删除虚拟机:

package cloud

import "fmt"

type DeleteVMParams struct {
	ProjectID string
	VMID      string
}

func DeleteVM(params DeleteVMParams) error {
	_, ok := GlobleVMs[params.VMID]
	if !ok {
		fmt.Printf("VMID: %v missing \n", params.VMID)
		return fmt.Errorf("VMID MISSING ERROR")
	}
	delete(GlobleVMs, params.VMID)
	return nil
}

主逻辑:

package main

import (
	"fmt"
	"self/pkg/cloud"
)

func main() {
	cloud.InitVMs()
	cloud.PrintVMs()

	putParams := cloud.PutVMParams{
		ProjectID: "B",
		VMID:      "001",
		VMName:    "MyServer",
	}
	err := cloud.PutVM(putParams)
	if err != nil {
		fmt.Printf("Put vm failed, err: %v \n", err.Error())
		return
	}
	cloud.PrintVMs()

	deleteParams := cloud.DeleteVMParams{
		ProjectID: "A",
		VMID:      "002",
	}
	err = cloud.DeleteVM(deleteParams)
	if err != nil {
		fmt.Printf("Delete vm failed, err: %v \n", err.Error())
		return
	}
	cloud.PrintVMs()
}

上面的代码中,

优先初始化了两个VM信息, 其中A租户拥有VM 001, B租户拥有VM 002, 并输出VM信息查看下当前的VM情况;
其次B租户修改VM 001名称, 然而其操作的是A租户的VM, 并输出VM信息查看下当前的VM情况;
最后A租户删除VM 002名称, 然而其操作的是B租户的VM, 并输出VM信息查看下当前的VM情况;

输出为:

[root self]# go run ./cmd/main.go 
VMs num: 2
VM info: {ID:002 Name:vm_002 CPU:arm ProjectID:B} 
VM info: {ID:001 Name:vm_001 CPU:x86 ProjectID:A} 

VMs num: 2
VM info: {ID:001 Name:MyServer CPU:x86 ProjectID:A} 
VM info: {ID:002 Name:vm_002 CPU:arm ProjectID:B} 

VMs num: 1
VM info: {ID:001 Name:MyServer CPU:x86 ProjectID:A}

需求

上面的业务代码PutVM, DeleteVM中缺失了Project权限检验逻辑, 导致任意Project的权限泄露. 需要添加权限检验逻辑.

思路

业务嵌入法

可以直接在业务代码头部添加权限检验逻辑, 如下:

// 权限检验逻辑
func checkPermission(projectID string, vmID string) bool {
	return ProjectToVMMap[projectID] == vmID
}

func PutVM(params PutVMParams) error {
	if !checkPermission(params.ProjectID, params.VMID) {
		return fmt.Errorf("Permission check failed")
	}
	// ...
	return nil
}

上面的代码中, 我直接选择在业务代码头部嵌入了权限检验逻辑, 其

优点: 逻辑简单
缺点: 代码耦合: 业务代码与权限代码耦合在一起;

输出为:

[root@sangfor self]# go run ./cmd/main.go 
VMs num: 2
VM info: {ID:001 Name:vm_001 CPU:x86 ProjectID:A} 
VM info: {ID:002 Name:vm_002 CPU:arm ProjectID:B} 

Put vm failed, err: Permission check failed

发现上面的代码修改已经达到了需求, 然而实际项目中的修改量往往不止于此, 可能存在n个业务接口需要适配, 极易造成改动引发. 我们期望可以将业务代码与权限检验逻辑解耦, 这样即使出现了改动引发错误, 也不会改动到业务代码.

`reflect`装饰器法

有一个很合适的解决方案：Python 装饰器

Python 的装饰器是一种非常有用的功能，它允许你在不修改原有函数代码的情况下，给函数添加额外的功能。装饰器本质上是一个函数，它接受一个函数作为参数，并返回一个新的函数。当你将装饰器应用于某个函数时，你实际上是在该函数周围“包裹”了一层额外的逻辑。

然而我们使用的是Go; 此时就需要reflect登场了.

可行性, 在Go中Func也是值(Value), 既然是值那么就可以通过reflect修改.

装饰器具体代码如下:

定义一个Model. 用于权限检验

// decroateCheckVMParams 虚拟机权限检验Model
type decroateCheckVMParams struct {
	ProjectID string
	VMID      string
}

定义一个装饰函数, 嵌入权限检验逻辑

// decorateCheckVM 装饰函数, 其会添加虚拟机权限检验逻辑到函数原代码中
func decorateCheckVM(newFunc, oldFunc interface{}, injectFunc func(in reflect.Value) *decroateCheckVMParams) {
	var decoratedFunc, targetFunc reflect.Value
	decoratedFunc = reflect.ValueOf(newFunc).Elem()
	targetFunc = reflect.ValueOf(oldFunc)
	v := reflect.MakeFunc(
		targetFunc.Type(),
		func(in []reflect.Value) (out []reflect.Value) {
			param := injectFunc(in[0])
			if !checkPermission(param.ProjectID, param.VMID) {
				err := fmt.Errorf("Check permission failed")
				out = []reflect.Value{reflect.ValueOf(err)}
				return
			}
			out = targetFunc.Call(in)
			return
		},
	)
	decoratedFunc.Set(v)
}

// DecoratedPutVM 装饰器逻辑
func DecoratedPutVM() func(params PutVMParams) error {
	decoratedHandleFunc := PutVM
	decorateCheckVM(&decoratedHandleFunc, decoratedHandleFunc, convertToDecoratedCheckVM)
	return decoratedHandleFunc
}

decorateCheckVM 函数是一个装饰器工厂，它接受三个参数：newFunc 是一个将被装饰的函数的引用，oldFunc 是原始的未装饰函数，injectFunc 是一个将 reflect.Value 转换为 *decroateCheckVMParams 的函数。decorateCheckVM 函数的目的是在 oldFunc 的调用前插入权限检查逻辑。

DecoratedPutVM 函数是一个具体的装饰器实现，它使用 decorateCheckVM 来装饰 PutVM 函数（假设 PutVM 是一个已经定义的函数，它接受 PutVMParams 类型的参数并返回一个 error）。DecoratedPutVM 函数返回一个新的函数，这个新函数在调用 PutVM 之前会执行权限检查。

这里有几个需要注意的地方：

decorateCheckVM 函数中的 decoratedFunc 和 targetFunc 是通过反射创建的函数值，decoratedFunc 是一个指向 newFunc 的指针，而 targetFunc 是 oldFunc 的值。
reflect.MakeFunc 用于创建一个新的函数，这个函数在调用时会先执行 injectFunc 来获取参数，然后检查权限，如果权限检查失败，它会返回一个错误；如果权限检查通过，它会调用原始的 targetFunc 函数。

这段代码的目的是通过反射来动态地添加权限检查逻辑到现有的函数中，而不需要修改原始函数的代码。这是一种常见的装饰器模式的实现方式，它允许开发者在不改变原始代码的情况下，为函数添加额外的功能。

定义一个注入函数, 实现将业务参数转换为权限检验参数:

// convertToDecoratedCheckVM 参数转换逻辑, 将业务参数转换为虚拟机权限检验参数
func convertToDecoratedCheckVM(in reflect.Value) *decroateCheckVMParams {
	var param *decroateCheckVMParams
	switch in.Interface().(type) {
	case PutVMParams:
		_p, _ := in.Interface().(PutVMParams)
		param = &decroateCheckVMParams{
			ProjectID: _p.ProjectID,
			VMID:      _p.VMID,
		}
	case DeleteVMParams:
		_p, _ := in.Interface().(DeleteVMParams)
		param = &decroateCheckVMParams{
			ProjectID: _p.ProjectID,
			VMID:      _p.VMID,
		}
	default:
		fmt.Println("Unimplement param type, build check project id param failed")
	}
	return param
}

convertToDecoratedCheckVM 函数接受一个 reflect.Value 类型的参数，并根据这个反射值的实际类型，将其转换为一个指向 decroateCheckVMParams 结构体的指针。

这个函数检查 reflect.Value 中的实际类型，如果是 PutVMParams 或 DeleteVMParams 类型，它就创建一个新的 decroateCheckVMParams 实例，并从输入参数中复制 ProjectID 和 VMID 字段。如果输入的类型不是这两种之一，它将打印一条错误消息。

适配主函数, 查看效果:

package main

import (
	"fmt"
	"self/pkg/cloud"
)

func main() {
	// ...
	putVMFunc := cloud.DecoratedPutVM()
	// err := cloud.PutVM(putParams)
	err := putVMFunc(putParams)
	if err != nil {
		fmt.Printf("Put vm failed, err: %v \n", err.Error())
		return
	}
	// ...
}

发现输出一致:

[root@ self]# go run ./cmd/main.go 
VMs num: 2
VM info: {ID:001 Name:vm_001 CPU:x86 ProjectID:A} 
VM info: {ID:002 Name:vm_002 CPU:arm ProjectID:B} 

Put vm failed, err: Check permission failed

这种方式的特点为:其利用反射与闭包的方式, 实现了类似于Python装饰器的效果,使得在进入业务逻辑前进行动态的权限检验;通过依赖注入方式, 将业务参数修改为由业务控制, 实现权限检验与业务参数解耦.

扩展性良好, 如果我们想实现DeleteVM逻辑, 只需要添加一段简短的逻辑即可:

// DecoratedDeleteVM 装饰器逻辑
func DecoratedDeleteVM() func(params DeleteVMParams) error {
	decoratedHandleFunc := DeleteVM
	decorateCheckVM(&decoratedHandleFunc, decoratedHandleFunc, convertToDecoratedCheckVM)
	return decoratedHandleFunc
}

读者可以自行检验, 效果.

背景

需求

思路

业务嵌入法

reflect装饰器法

`reflect`装饰器法