go 随websocket存亡的token

 1 package token
 2 
 3 import (
 4     "crypto/md5"
 5     "fmt"
 6     "io"
 7     "math/rand"
 8     "sync"
 9     "time"
10 )
11 
12 type User struct {
13     ID    string
14     Timer *time.Timer
15 }
16 
17 var Token map[string]*User = make(map[string]*User)
18 var maxLiveTime time.Duration = time.Hour * 24 // Token更新时间暂定为24小时
19 var lock sync.RWMutex
20 
21 const randStringLen = 16
22 
23 func randString() string {
24     str := make([]byte, randStringLen)
25     for i := 0; i < randStringLen; i++ {
26         str[i] = byte(rand.Uint32() & ((1 << 7) - 1))
27     }
28     return string(str)
29 }
30 
31 func New(userID string, token *string, writeToken func()) {
32     lock.Lock()
33     delete(Token, *token)
34     hash := md5.New()
35     io.WriteString(hash, userID)
36     io.WriteString(hash, time.Now().String())
37     *token = fmt.Sprintf("%x", hash.Sum(nil))
38     timer := time.AfterFunc(maxLiveTime, func() { New(userID, token, writeToken) })
39     Token[*token] = &User{
40         ID:    userID,
41         Timer: timer,
42     }
43     lock.Unlock()
44     writeToken()
45 }
46 
47 func Del(token string) {
48     lock.Lock()
49     defer lock.Unlock()
50     user := Token[token]
51     user.Timer.Stop()
52     delete(Token, token)
53 }
54 
55 func GetUser(token string) (user *User, exisToken bool) {
56     lock.RLock()
57     defer lock.RUnlock()
58     user, exisToken = Token[token]
59     return user, exisToken
60 }

 1 package Network
 2 
 3 import (
 4     "logs"
 5     "time"
 6     "token"
 7 
 8     "github.com/gorilla/websocket"
 9 )
10 
11 type ClientData struct {
12     conn  *websocket.Conn
13     token string
14 }
15 
16 var wsClients = make(map[string]ClientData)
17 
18 // 是否在线
19 func isOnline(clientID string) bool {
20     lock.RLock()
21     _, isOnline := wsClients[clientID]
22     lock.RUnlock()
23     return isOnline
24 }
25 
26 // --------------------------------------------- 连接管理部分
27 
28 // 增加WebSocket连接
29 func addClient(clientID string, con *websocket.Conn) (canAdd bool) {
30 
31     if isOnline(clientID) {
32         writeResponseToClient(clientID, ServerResponse{
33             Method:  "/wsLogin",
34             Code:    noticeForceQuit,
35             Message: mForceQuit,
36             Result:  nil,
37         })
38         deleteClient(clientID)
39         time.Sleep(time.Millisecond) // 如果不睡眠或者睡眠时间太短，两个连接都会被同时关掉
40         logs.Print(clientID, mForceQuit)
41         addConnection(clientID, con)
42         return false
43     }
44     addConnection(clientID, con)
45     return true
46 }
47 
48 func addConnection(clientID string, con *websocket.Conn) {
49     clientData := ClientData{
50         conn:  con,
51         token: "",
52     }
53     token.New(clientID, &clientData.token, func() { //此函数用于每次token变化后将token写回客户端
54         //更新完client后才能writeResponse,注意死锁
55         lock.Lock()
56         wsClients[clientID] = clientData
57         lock.Unlock()
58         writeResponseToClient(clientID, ServerResponse{
59             Code:    Success,
60             Method:  "/token",
61             Message: "",
62             Result:  map[string]string{"Token": clientData.token},
63         })
64     })
65 }
66 
67 // 获取客户端连接
68 func getClient(clientID string) (clientConn *websocket.Conn, isOnline bool) {
69 
70     lock.RLock()
71     clientData, ok := wsClients[clientID]
72     lock.RUnlock()
73     return clientData.conn, ok
74 }
75 
76 // 删除WebSocket连接
77 func deleteClient(clientID string) {
78     clientData, ok := wsClients[clientID]
79     if !ok {
80         return
81     }
82     lock.Lock()
83     clientData.conn.Close()
84     token.Del(clientData.token)
85     delete(wsClients, clientID)
86     lock.Unlock()
87 }