go-zero 自定义中间件的几种方式

首先 go-zero 已经为我们提供了很多的中间件的实现，但有时难免有需求需要自定义，这里介绍几种自定义的方法，供参考。

1.通过 api 文件生成并填入具体逻辑

定义 api

首先你需要在 api 文件中定义你需要的路由位置和中间件的名字：

// test for ping
@server (
	timeout:    3s
	middleware: PingMiddleware // 自定义的中间件
)
service userapi-api {
	@handler Ping
	get /ping returns (Reply)
}

生成对应的模板

接下来通过 goctl 生成对应的模板： goctl api go --api xxx.api --dir .

填充中间件逻辑

然后通过代码生成，goctl 会在 project/internal 下生成 middleware 文件夹，你的 pingmiddleware.go 就在里面，然后填充 中间件逻辑：

package middleware

import (
	"net/http"

	"github.com/zeromicro/go-zero/core/logx"
)

type PingMiddleware struct {
}

func NewPingMiddleware() *PingMiddleware {
	return &PingMiddleware{}
}

// need to implement logic
func (m *PingMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		logx.Info("This is pingMiddleware before logic") // 填充逻辑

		// Passthrough to next handler if need
		next(w, r)

		logx.Info("This is pingMiddleware before logic") // 填充逻辑
	}
}

另外，你还需要在 svc/servicecontext.go 中填充你的中间件：

package svc

import (
	"github.com/zeromicro/go-zero/rest"
	
	"userapiv1/internal/config"
	"userapiv1/internal/middleware"
)

type ServiceContext struct {
	Config config.Config
	PingMiddleware rest.Middleware // manual added
}

func NewServiceContext(c config.Config) *ServiceContext {
	return &ServiceContext{
		Config: c,
		PingMiddleware: middleware.NewPingMiddleware().Handle, // manual added
	}
}

另外也看下 路由注册部分 handler/routers.go 是否加入中间件：

func RegisterHandlers(server *rest.Server, serverCtx *svc.ServiceContext) {
	server.AddRoutes(
		rest.WithMiddlewares(
			[]rest.Middleware{serverCtx.PingMiddleware}, // auto added by goctl
			[]rest.Route{
				{
					Method:  http.MethodGet,
					Path:    "/ping",
					Handler: PingHandler(serverCtx),
				},
			}...,
		),
		rest.WithTimeout(3000*time.Millisecond),
	)
}

完成以上部分，一个自定义的中间件就算完成了。

2.在 server 启动前完成 中间件 的注册

用这种方法添加自定义的中间件就更简单了，你只需要在 internal/middleware/xxxMiddleware.go 定义自己的中间件，然后注册到 server 中即可。

定义中间件：

package middleware

import (
	"net/http"

	"github.com/zeromicro/go-zero/core/logx"
)

func LogMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		logx.Info("here is LogMiddleware handle before logic")

		next(w, r)

		logx.Info("here is LogMiddleware handle after logic")

	}
}

注册到 server 中

package main

import (
	"flag"
	"fmt"

	"userapiv1/internal/config"
	"userapiv1/internal/handler"
	"userapiv1/internal/middleware"
	"userapiv1/internal/svc"

	"github.com/zeromicro/go-zero/core/conf"
	"github.com/zeromicro/go-zero/rest"
)

var configFile = flag.String("f", "etc/userapi-api.yaml", "the config file")

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)

	// go-zero version >= v1.7.0
	//logx.AddWriter(logx.NewWriter(os.Stdout))  // 添加控制台输出

	// register middleware in global scope
	server.Use(middleware.LogMiddleware)

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

以上就是两种自定义中间件的添加方法，希望对你有用。