Golang 搭建 WebSocket 应用（四） - jwt 认证

在上一篇文章中，我们已经搭建起了基本可用的一个 WebSocket 推送中心，但是有一个比较大的问题是，
我们并没有对进行连接的客户端进行认证，这样就会有一定的风险，如果被恶意攻击，
可能会影响我们的 WebSocket 服务器的正常运作。

本文我们就来把认证这个很关键的功能给补一下，在本文中，我们将会使用 jwt 来对我们的客户端进行认证。

什么是 jwt？

JWT 是 JSON Web Token 的缩写，是一种用于在网络中安全传递信息的开放标准。它是一种紧凑且自包含的方式，用于在各方之间传递信息，通常用于身份验证和授权机制。

JWT 主要由三个部分组成：

1. 头部（Header）： 包含关于令牌的元数据，例如令牌的类型（typ）和签名算法（alg）。头部是一个 JSON 对象，通常会经过 Base64 编码。

2. 载荷（Payload）： 包含要传递的信息，通常包括用户身份信息以及其他声明。载荷也是一个 JSON 对象，同样经过 Base64 编码。

3. 签名（Signature）： 使用头部和载荷以及密钥生成的签名，用于验证令牌的真实性和完整性。签名是对头部和载荷的哈希值进行签名后的结果。

这三个部分通过点号（.）连接起来形成一个字符串，即 JWT。最终的 JWT 结构如下：

header.payload.signature

一个简单的 jwt 例子

jwt 的使用会分为两部分：生成 token，使用 token。本文中将会使用 golang-jwt/jwt 来做 jwt 的验证。

生成 token

生成 token 的操作有以下两步：

1. 创建一个 token 对象：使用的是 jwt.NewWithClaims 方法，它第一个参数指定了签名算法，这里使用的是 HMAC，第二个参数接受一个 jwt.MapClaims，也就是上面提到的 payload。

token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
    "foo": "bar",
    "nbf": time.Date(2015, 10, 10, 12, 0, 0, 0, time.UTC).Unix(),
})

2. 利用上一步的 token 对象生成成一个 jwt 的签名字符串：使用的是 token 的 SignedString 方法，它接受一个 key 作为参数，这个 key 也会用于解析这一步生成的 token 字符串。

// 注意：这里的 secret 在实际使用的时候需要替换为自己的 key
//（一般为一个随机字符串）
tokenString, err := token.SignedString([]byte("secret"))

我们生成 token 的操作一般发生在用户登录成功之后，这个 token 会作为用户后续发起请求的凭证。

使用 token

在用户登录成功拿到 token 之后，会使用这个 token 去从服务器获取其他资源，服务器会解析这个 token 并校验。

使用 token 的步骤如下：

1. 解析 token：使用的是 jwt.Parse 方法，它第一个参数接受 token 字符串，第二个参数是一个函数（函数的参数就是解析出来的 token 对象，函数返回解密的 key）

token, err = jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
    // 返回签名密钥，需要是 []byte 类型的
    return []byte("secret"), nil
})

2. 使用 token 中的 payload：payload 也就是我们业务实际使用的数据，其他的东西只是使用 jwt token 这门技术附带的一些东西。

这里说的 payload 实际上就是 token 对象的 Claims 属性，它是一个 map，保存了我们一些业务的数据还有其他一些 jwt 本身的字段。

// claims: map[foo:bar nbf:1.4444784e+09]
fmt.Println("claims:", token.Claims)

在上面这行代码中，foo:bar 是我们的业务数据，而 nbf:1.4444784e+09 是 jwt 本身使用的字段。nbf 表示的是在这个时间之前，这个 token 不应该被处理。

例子

我们把上面的生成 token 的代码和解析 token 的代码放到一起看看效果：

注意：实际使用中这两个操作是分开的。

func TestHmac(t *testing.T) {
    // part 1:
	// 创建一个新的 token 对象，指定签名方法和你想要包含的 claims
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
		"foo": "bar",
		"nbf": time.Date(2015, 10, 10, 12, 0, 0, 0, time.UTC).Unix(),
	})

	// 生成签名字符串，secret 为签名密钥
	tokenString, err := token.SignedString([]byte("secret"))

	// eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJmb28iOiJiYXIiLCJuYmYiOjE3Mjg1NjE2MDB9.9AxzBmYOuOnWfKUul57ATzjQ-sMzbggaoIdDjVzjm2Y, nil
	fmt.Println(tokenString, err)

    // part 2:
	// 解析 token
	token, err = jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}

		// 返回签名密钥
		return []byte("secret"), nil
	})
	if err != nil {
		panic(err)
	}

	// header: map[alg:HS256 typ:JWT]
	fmt.Println("header:", token.Header)
	// claims: map[foo:bar nbf:1.4444784e+09]
	fmt.Println("claims:", token.Claims)
	// signature: Nv24hvNy238QMrpHvYw-BxyCp00jbsTqjVgzk81PiYA
	fmt.Println("signature:", base64.RawURLEncoding.EncodeToString(token.Signature))

	if claims, ok := token.Claims.(jwt.MapClaims); ok {
		// bar 1.4444784e+09
		fmt.Println(claims["foo"], claims["nbf"])
	} else {
		panic(err)
	}
}

本文不是讲解 JWT 的文章，关于 JWT 的更多细节可以参考 rfc7519。

在消息推送中心 demo 加上 jwt 认证

我们使用 jwt 的目的是为了杜绝一些未知来源的连接，而在我们上一篇文章的实现中，
是先建立起连接，然后再进行 “登录” 操作的，这样就会导致即使是未授权的客户端也可以先进行连接，
这样就会在认证之前就启动了两个协程。

而如果这些连接并不是正常的连接，它们只连接但是不登录，那样就会有很多僵尸连接，这显然不是我们想要的结果。

我们可以在客户端打开连接的时候就去验证客户端的 token，如果 token 校验失败，则直接断开连接，
就可以解决上述问题了，从而避免不必要的开销。

如何在建立连接的时候就认证

要实现这个很简单，我们只需要在连接的 url 后加一个 queryString 即可，如下：

ws = new WebSocket('ws://127.0.0.1:8181/ws?token=123')

然后在 serveWs 中通过 r.FormValue("token") 来获取客户端传递过来的 token，
再对其进行认证，认证不通过则拒绝连接。

具体来说，我们的 serveWs 会添加以下几行代码：

jwt := NewJwt(r.FormValue("token"))
err := jwt.Parse()
if err != nil {
    log.Println(fmt.Errorf("jwt parse error: %w", err))
    return
}

在函数入口的地方就进行验证，验证不通过则不进行连接。这里的 jwt 定义如下：

type Jwt struct {
	Secret string
	Token  string
}

func NewJwt(token string) *Jwt {
	return &Jwt{
		Token:  token,
		Secret: os.Getenv("JWT_SECRET"),
	}
}

func (j *Jwt) Parse() error {
	_, err := jwt.Parse(j.Token, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}

		return []byte(j.Secret), nil
	})

	return err
}

在 NewJwt 中，我们从 env 中获取 JWT_SECRET，这让我们的配置可以更加灵活。

token 中加上 uid

我们知道了，在 JWT 中的 payload 是可以加入我们的自定义数据的，所以我们的 uid 其实是可以加入到 jwt 的 token 中的，我们只需要在用户第一次获取 token 的时候加上即可：

token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
    "uid": "123",
    "nbf": time.Date(2015, 10, 10, 12, 0, 0, 0, time.UTC).Unix(),
})

同样的，在解析 jwt token 的时候，可以从中取出这个 uid：

// 第一个返回值是 uid，第二个是 error
func (j *Jwt) Parse() (string, error) {
	token, err := jwt.Parse(j.Token, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", token.Header["alg"])
		}

		return []byte(j.Secret), nil
	})

	if err != nil {
		return "", err
	}

    // 获取 uid
	if claims, ok := token.Claims.(jwt.MapClaims); ok {
		return claims["uid"].(string), nil
	} else {
		return "", fmt.Errorf("jwt parse error")
	}
}

最终，我们的 serveWs 演化成了如下这样：

func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
    // 解析 jwt token，从中取得 uid
	jwt := NewJwt(r.FormValue("token"))
	uid, err := jwt.Parse()
	if err != nil {
		log.Println(fmt.Errorf("jwt parse error: %w", err))
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(fmt.Errorf("upgrade error: %w", err))
		return
	}

	client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256), uid: uid}
	client.hub.register <- client

	go client.writePump()
	go client.readPump()
}

这样一来，我们的 readPump 里面就不再需要处理登录消息了，那就暂时先把 readPump 中的逻辑去掉先：

func (c *Client) readPump() {
	defer func() {
		c.hub.unregister <- c
		_ = c.conn.Close()
	}()
	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetReadDeadline(time.Time{}) // 永不超时
	for {
		// 从客户端接收消息
		_, _, err := c.conn.ReadMessage()
		if err != nil {
			log.Println("readPump error: ", err)
			break
		}
	}
}

register 里的 uid 关联

我们之前是在 readPump 方法中将 uid 和 WebSocket 建立起关联的，但由于我们已经去掉了 readPump 中
的登录消息处理逻辑。

因此我们需要在 Hub 的 register 中将 uid 和 WebSocket 建立起关联：

case client := <-h.register:
    h.Lock()
    h.clients[client] = true
    // 建立起 uid 跟 WebSocket 的关联
    h.userClients[client.uid] = client
    h.Unlock()

jti

在此前使用过的一些 jwt 封装中，会有些使用 jwt 规范中的 jti 字段来传输 token 的唯一 ID，
在本文的实现中，uid 也是同样的功能。如果我们之后看到了 jti 这个字段，不要惊讶，其实这个才是规范。

测试

我们将生成 token 的代码中的 foo: bar 这个键值对修改为 uid: 123 之后，得到了如下 token：

eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJuYmYiOjE0NDQ0Nzg0MDAsInVpZCI6IjEyMyJ9.9yJ-ABQGJkdnDqHo-wV-vojQFEQGt-I0dyva1w6EQ7E

我们在 ws 链接后加上这个 token 即可连接成功：

ws = new WebSocket('ws://127.0.0.1:8181/ws?token=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJuYmYiOjE0NDQ0Nzg0MDAsInVpZCI6IjEyMyJ9.9yJ-ABQGJkdnDqHo-wV-vojQFEQGt-I0dyva1w6EQ7E')

我们可以看到，ws 连接返回的状态码是 101，这就表明我们的连接成功了。

推送消息

同之前一样的，在控制台执行一个 curl 命令即可推送消息了：

curl "http://localhost:8181/send?uid=123&message=Hello%20World"

总结

本文我们使用 jwt 对我们的 WebSocket 进行了认证。
jwt 认证我们可以已经用过很多次了，但如果还没有用过或者还没有在 go 中用过 jwt 认证的话，
那么本文就是一个很好的入门文章。

最后，再来回顾一下本文的内容：

• jwt 包含了三部分：头部（Header）、载荷（Payload）、签名（Signature）。
• 在 go 中可以使用 golang-jwt/jwt 来做 jwt 的验证。
• 创建 token 分两步：使用 jwt.NewWithClaims 创建 token 对象、使用 SignedString 生成签名字符串。
• 使用 token 分两步：使用 jwt.Parse 解析 token、使用 token.Claims 获取 payload。
• 在建立连接的时候就进行认证，可以避免非法连接导致的开销。
• jwt 中可以加入我们的自定义数据，比如 uid，在 jwt.NewWithClaims 中加上即可。