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 }
因为登录是用websocket,返回前端token的也是websoket(websoket验证成功后马上返回token),所以websocket在,token就在,websocket关闭连接后token也结束了生命周期