当前位置：首页 > article >正文

go-zero学习笔记（五）

article 2025/2/24 13:04:28

api自定义中间件

1. 修改.api文件


	syntax="v1"
	
	type (
		GetInfoReq {
			IDs []string `json:"IDs"`
		}
	
		GetInfoData {
			ID     string `json:"ID"`
			Name   string `json:"Name"`
			MD5    string `json:"md5"`
			Size   int64  `json:"Size"`
			Uptime int64  `json:"uptime"`
		}
	
		GetInfoResp {
			Data      []*GetInfoData `json:"data"`
			Msg       string         `json:"msg"`
			Code      int64          `json:"code"`
			RequestID string         `json:"requestId"`
		}
	)
	
	// 修饰其紧挨的service块，且一个server块只能修饰一个service块
	@server(
		prefix: api/demo/v1         
		group: demo                 
		middleware: CheckMiddleware, TestMiddleware  
	)
	service demo-api {
		@doc(
			summary: "信息获取"
		)
		@handler GetInfoHandler
		post /Info/get (GetInfoReq) returns (GetInfoResp)
	}

在server块中新增checkMiddleware, testMiddleware 两个中间件，其中checkMiddleware, testMiddleware为中间件名称。


    middleware: CheckMiddleware, TestMiddleware

之前说过， server修饰的是其紧挨的service 块，因此在server块中新增中间件，会为其紧挨的service 块中的每个接口都增加中间件checkMiddleware, testMiddleware，因此，请注意作用域，如何部分接口不需要中间件，请分开定义。

2. 执行命令

goctl api go -api ./api/demo.api -dir . -style gozero

生成关于中间件部分的代码主要如下：

3. 编写中间件逻辑代码

接下来就可以在Handle函数中添加对应的逻辑了，在next(w, r)前后都可以添加，

package middleware

import "net/http"

type CheckMiddleware struct {
}

func NewCheckMiddleware() *CheckMiddleware {
	return &CheckMiddleware{}
}

func (m *CheckMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// TODO generate middleware implement function, delete after code implementation

		// Passthrough to next handler if need

		/******
		请求到来后需要执行的代码
		代码块a
		*/

		next(w, r)

		/******
		请求返回时需要执行的代码
		代码块A
		*/
	}
}

package middleware

import "net/http"

type TestMiddleware struct {
}

func NewTestMiddleware() *TestMiddleware {
	return &TestMiddleware{}
}

func (m *TestMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		// TODO generate middleware implement function, delete after code implementation

		// Passthrough to next handler if need
		/******
		请求到来后需要执行的代码
		代码块b
		*/
		
		next(w, r)
		
		/******
		请求返回时需要执行的代码
		代码块B
		*/
	}
}

如果你需要外部传递参数，可以在对应的结构体中添加，如在TestMiddleware 结构体中添加变量，并在New函数中传入

type TestMiddleware struct {
	Num   int
	Class string
}

func NewTestMiddleware(num int, class string) *TestMiddleware {
	return &TestMiddleware{
		Num:   num,
		Class: class,
	}
}

4. 修改servicecontext.go文件

通过函数间的调用关系，可以发现，中间件中的Handle函数是在servicecontext.go被调用的。

type ServiceContext struct {
	Config          config.Config
	CheckMiddleware rest.Middleware
	TestMiddleware  rest.Middleware
	DemoRpc         demo.DemoClient
}

func NewServiceContext(c config.Config) *ServiceContext {
	return &ServiceContext{
		Config:          c,
		CheckMiddleware: middleware.NewCheckMiddleware().Handle,
		TestMiddleware:  middleware.NewTestMiddleware(5, "rongyu").Handle,
		DemoRpc:         demo.NewDemoClient(zrpc.MustNewClient(c.DemoRpcConf).Conn()),
	}
}

如果，你对中间件的结构体进行了修改，新增了变量，则需要修改上述代码的NewTestMiddleware处将值传进去。

通过阅读代码可以发现，中间件在ServiceContext 中定义，在main函数中初始化，随后被传入到RegisterHandlers中使用。

func main() {
	flag.Parse()

	var c config.Config
	conf.MustLoad(*configFile, &c)

	server := rest.MustNewServer(c.RestConf)
	defer server.Stop()

	ctx := svc.NewServiceContext(c)
	handler.RegisterHandlers(server, ctx)

	fmt.Printf("Starting server at %s:%d...\n", c.Host, c.Port)
	server.Start()
}


func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
	server.AddRoutes(
		rest.WithMiddlewares(
			[]rest.Middleware{serverCtx.CheckMiddleware, serverCtx.TestMiddleware},
			[]rest.Route{
				{
					Method:  http.MethodPost,
					Path:    "/Info/get",
					Handler: demo.GetInfoHandler(serverCtx),
				},
			}...,
		),
		rest.WithPrefix("/api/demo/v1"),
	)
}