B站微服务框架Kratos详细教程（3）- 中间件

背景

基于bm的handler机制，可以自定义很多middleware(中间件)进行通用的业务处理，比如用户登录鉴权。
接下来就以鉴权为例，说明middleware的写法和用法。

创建新项目

kratos new middledemo --http
gine, err error) {
	var (
		cfg bm.ServerConfig
		ct paladin.TOML
	)
	if err = paladin.Get("http.toml").Unmarshal(&ct); err != nil {
		return
	}
	if err = ct.Get("Server").UnmarshalTOML(&cfg); err != nil {
		return
	}
	svc = s
	engine = bm.DefaultServer(&cfg) //Ctrl+点击，查看bm server代码

	pb.RegisterDemoBMServer(engine, s)
	initRouter(engine)
	err = engine.Start()
	return
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

Ctrl+点击bm.DefaultServer，查看blademaster/server.go代码（注意是kratos框架blademaster的server.go）：

func DefaultServer(conf *ServerConfig) *Engine {
	engine := NewServer(conf)
	engine.Use(Recovery(), Trace(), Logger())
	return engine
}
1
2
3
4
5

会默认创建一个bm engine，并注册Recovery(), Trace(), Logger()三个middlerware用于全局handler处理，优先级从前到后。

自定义中间件

如果想要将自定义的middleware注册进全局，可以继续调用Use方法如下：

engine.Use(MyMiddleware())
1

此方法会将MyMiddleware追加到已有的全局middleware后执行。
middleware本质上就是一个handler，接口和方法声明如下：

// Handler responds to an HTTP request.
type Handler interface {
	ServeHTTP(c *Context)
}

// HandlerFunc http request handler function.
type HandlerFunc func(*Context)

// ServeHTTP calls f(ctx).
func (f HandlerFunc) ServeHTTP(c *Context) {
	f(c)
}
1
2
3
4
5
6
7
8
9
10
11
12

只要实现了Handler接口，就可以作为engine的全局中间件使用：

engine.Use(YourHandler)
1

声明为HandlerFunc方法，可以作为engine的全局中间件使用：

engine.UseFunc(YourHandlerFunc)
1

也可以作为router的局部中间件使用：

e.GET("/path", YourHandlerFunc)
1

示例代码：

package http

import (
	"net/http"

	pb "middledemo/api"
	"middledemo/internal/model"
	"github.com/go-kratos/kratos/pkg/conf/paladin"
	"github.com/go-kratos/kratos/pkg/log"
	bm "github.com/go-kratos/kratos/pkg/net/http/blademaster"
)

var svc pb.DemoServer

// New new a bm server.
func New(s pb.DemoServer) (engine *bm.Engine, err error) {
	var (
		cfg bm.ServerConfig
		ct paladin.TOML
	)
	if err = paladin.Get("http.toml").Unmarshal(&ct); err != nil {
		return
	}
	if err = ct.Get("Server").UnmarshalTOML(&cfg); err != nil {
		return
	}
	svc = s
	engine = bm.DefaultServer(&cfg)

	myware := &MyMiddleware{}
	engine.Use(myware) //注册自定义全局中间件
	//engine.UseFunc(myware.ServeHTTP)
	//engine.GET("/path",myware.ServeHTTP)

	//只有方法的中间件，这个和howToStart一样，只不过在这里是全局的
	engine.UseFunc(MyMiddleHandler)
	//也可以像howToStart指定范围使用
	//engine.GET("/path",MyMiddleHandler)

	pb.RegisterDemoBMServer(engine, s)
	initRouter(engine)
	err = engine.Start()
	return
}

func initRouter(e *bm.Engine) {
	e.Ping(ping)
	g := e.Group("/middledemo", MyGroupHandler)
	{
		g.GET("/start", howToStart)
	}
}

func ping(ctx *bm.Context) {
	if _, err := svc.Ping(ctx, nil); err != nil {
		log.Error("ping error(%v)", err)
		ctx.AbortWithStatus(http.StatusServiceUnavailable)
	}
}

// example for http request handler.
func howToStart(c *bm.Context) {
	log.Info("输出 Handler: howToStart")

	k := &model.Kratos{
		Hello: "Golang 大法好 !!!",
	}
	c.JSON(k, nil)
}

//自定义中间件
type MyMiddleware struct {
	Key   string
	Value   string
}
// ServeHTTP implements from Handler interface
func (d *MyMiddleware) ServeHTTP(c *bm.Context) {
	c.Set(d.Key, d.Value)

	log.Info("全局自定义中间件 MyMiddleware: ServeHTTP")

	//TODO: 中间件业务代码, 例如授权验证、限流等操作

	c.Next()
}

func MyMiddleHandler(c *bm.Context) {
	// some code
	log.Info("全局方法中间件: MyMiddleHandler")
	c.Next()
}

func MyGroupHandler(c *bm.Context) {
	// some code
	log.Info("分组中间件: MyGroupHandler")
	c.Next()
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

运行代码：

kratos run
1

打开浏览器访问：

http://localhost:8000/middledemo/start
1

查看控制台输出：

可以看到我们自定义的中间件被按顺序执行了一次。

如果需要全部自定义全局执行的middleware，可以使用NewServer方法创建一个无middleware的engine对象，然后使用engine.Use/UseFunc进行注册。

示例代码：

// New new a bm server.
func New(s pb.DemoServer) (engine *bm.Engine, err error) {
	var (
		cfg bm.ServerConfig
		ct paladin.TOML
	)
	if err = paladin.Get("http.toml").Unmarshal(&ct); err != nil {
		return
	}
	if err = ct.Get("Server").UnmarshalTOML(&cfg); err != nil {
		return
	}
	svc = s
	
	//全局自定义, 不使用框架内置中间件: Recovery(), Trace(), Logger()
	engine = bm.NewServer(&cfg) 

	myware := &MyMiddleware{}
	engine.Use(myware) //注册自定义全局中间件
	//engine.UseFunc(myware.ServeHTTP)
	//engine.GET("/path",myware.ServeHTTP)

	//只有方法的中间件，这个和howToStart一样，只不过在这里是全局的
	engine.UseFunc(MyMiddleHandler)
	//也可以像howToStart指定范围使用
	//engine.GET("/path",MyMiddleHandler)

	pb.RegisterDemoBMServer(engine, s)
	initRouter(engine)
	err = engine.Start()
	return
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

局部中间件

先来看一段鉴权伪代码示例：

func Example() {
	myHandler := func(ctx *bm.Context) {
		mid := metadata.Int64(ctx, metadata.Mid)
		ctx.JSON(fmt.Sprintf("%d", mid), nil)
	}

	authn := auth.New(&auth.Config{DisableCSRF: false})

	e := bm.DefaultServer(nil)

	// "/user"接口必须保证登录用户才能访问，那么我们加入"auth.User"来确保用户鉴权通过，才能进入myHandler进行业务逻辑处理
	e.GET("/user", authn.User, myHandler)
	// "/guest"接口访客用户就可以访问，但如果登录用户我们需要知道mid，那么我们加入"auth.Guest"来尝试鉴权获取mid，但肯定会继续执行myHandler进行业务逻辑处理
	e.GET("/guest", authn.Guest, myHandler)

    // "/owner"开头的所有接口，都需要进行登录鉴权才可以被访问，那可以创建一个group并加入"authn.User"
	o := e.Group("/owner", authn.User)
	o.GET("/info", myHandler) // 该group创建的router不需要再显示的加入"authn.User"
	o.POST("/modify", myHandler) // 该group创建的router不需要再显示的加入"authn.User"

	e.Start()
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

完整示例代码：
首先复制示例代码kratos/example/blademaster/middleware/auth/auth.go到自己的项目，或者可以直接import引入示例代码

package http

import (
	"middledemo/internal/server/auth"
	"net/http"

	pb "middledemo/api"
	"middledemo/internal/model"
	"github.com/go-kratos/kratos/pkg/conf/paladin"
	"github.com/go-kratos/kratos/pkg/log"
	bm "github.com/go-kratos/kratos/pkg/net/http/blademaster"
)

var svc pb.DemoServer

// New new a bm server.
func New(s pb.DemoServer) (engine *bm.Engine, err error) {
	var (
		cfg bm.ServerConfig
		ct paladin.TOML
	)
	if err = paladin.Get("http.toml").Unmarshal(&ct); err != nil {
		return
	}
	if err = ct.Get("Server").UnmarshalTOML(&cfg); err != nil {
		return
	}
	svc = s
	engine = bm.DefaultServer(&cfg)

	//engine = bm.NewServer(&cfg) //全局自定义, 不使用框架内置中间件: Recovery(), Trace(), Logger()

	myware := &MyMiddleware{}
	engine.Use(myware) //注册自定义全局中间件
	//engine.UseFunc(myware.ServeHTTP)
	//engine.GET("/path",myware.ServeHTTP)

	//只有方法的中间件，这个和howToStart一样，只不过在这里是全局的
	engine.UseFunc(MyMiddleHandler)
	//也可以像howToStart指定范围使用
	//engine.GET("/path",MyMiddleHandler)

	pb.RegisterDemoBMServer(engine, s)
	initRouter(engine)
	err = engine.Start()
	return
}

func initRouter(e *bm.Engine) {
	e.Ping(ping)

	//按分组验证
	g := e.Group("/middledemo", MyGroupHandler)
	{
		g.GET("/start", howToStart)
	}

	authn := auth.New(&auth.Config{DisableCSRF: false})

	//添加多个Handler，按顺序执行
	e.GET("/userinfo", authn.User, userInfo)

	//按分组验证
	g2 := e.Group("/user", authn.User)
	{
		g2.GET("/info", userInfo2)
	}
}

func ping(ctx *bm.Context) {
	if _, err := svc.Ping(ctx, nil); err != nil {
		log.Error("ping error(%v)", err)
		ctx.AbortWithStatus(http.StatusServiceUnavailable)
	}
}

// example for http request handler.
func howToStart(c *bm.Context) {
	log.Info("输出 Handler: howToStart")

	k := &model.Kratos{
		Hello: "Golang 大法好 !!!",
	}
	c.JSON(k, nil)
}


func userInfo(c *bm.Context) {
	log.Info("输出 Handler: userInfo")

	k := &model.Kratos{
		Hello: "用户信息 !!!",
	}
	c.JSON(k, nil)
}

func userInfo2(c *bm.Context) {
	log.Info("输出 Handler: userInfo")

	k := &model.Kratos{
		Hello: "用户信息22222 !!!",
	}
	c.JSON(k, nil)
}

//自定义中间件
type MyMiddleware struct {
	Key   string
	Value   string
}
// ServeHTTP implements from Handler interface
func (d *MyMiddleware) ServeHTTP(c *bm.Context) {
	c.Set(d.Key, d.Value)

	log.Info("全局自定义中间件 MyMiddleware: ServeHTTP")

	//TODO: 中间件业务代码, 例如授权验证、限流等操作

	c.Next()
}

func MyMiddleHandler(c *bm.Context) {
	// some code
	log.Info("全局方法中间件: MyMiddleHandler")
	c.Next()
}

func MyGroupHandler(c *bm.Context) {
	// some code
	log.Info("分组中间件: MyGroupHandler")
	c.Next()
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

打开浏览器分别访问以下地址：

http://localhost:8000/userinfo
http://localhost:8000/userinfo?access_token