go 随websocket存亡的token

package token

import (
    "crypto/md5"
    "fmt"
    "io"
    "math/rand"
    "sync"
    "time"
)

type User struct {
    ID    string
    Timer *time.Timer
}

var Token map[string]*User = make(map[string]*User)
var maxLiveTime time.Duration = time.Hour * 24 // Token更新时间暂定为24小时
var lock sync.RWMutex

const randStringLen = 16

func randString() string {
    str := make([]byte, randStringLen)
    for i := 0; i < randStringLen; i++ {
        str[i] = byte(rand.Uint32() & ((1 << 7) - 1))
    }
    return string(str)
}

func New(userID string, token *string, writeToken func()) {
    lock.Lock()
    delete(Token, *token)
    hash := md5.New()
    io.WriteString(hash, userID)
    io.WriteString(hash, time.Now().String())
    *token = fmt.Sprintf("%x", hash.Sum(nil))
    timer := time.AfterFunc(maxLiveTime, func() { New(userID, token, writeToken) })
    Token[*token] = &User{
        ID:    userID,
        Timer: timer,
    }
    lock.Unlock()
    writeToken()
}

func Del(token string) {
    lock.Lock()
    defer lock.Unlock()
    user := Token[token]
    user.Timer.Stop()
    delete(Token, token)
}

func GetUser(token string) (user *User, exisToken bool) {
    lock.RLock()
    defer lock.RUnlock()
    user, exisToken = Token[token]
    return user, exisToken
}

package Network

import (
    "logs"
    "time"
    "token"

    "github.com/gorilla/websocket"
)

type ClientData struct {
    conn  *websocket.Conn
    token string
}

var wsClients = make(map[string]ClientData)

// 是否在线
func isOnline(clientID string) bool {
    lock.RLock()
    _, isOnline := wsClients[clientID]
    lock.RUnlock()
    return isOnline
}

// --------------------------------------------- 连接管理部分

// 增加WebSocket连接
func addClient(clientID string, con *websocket.Conn) (canAdd bool) {

    if isOnline(clientID) {
        writeResponseToClient(clientID, ServerResponse{
            Method:  "/wsLogin",
            Code:    noticeForceQuit,
            Message: mForceQuit,
            Result:  nil,
        })
        deleteClient(clientID)
        time.Sleep(time.Millisecond) // 如果不睡眠或者睡眠时间太短，两个连接都会被同时关掉
        logs.Print(clientID, mForceQuit)
        addConnection(clientID, con)
        return false
    }
    addConnection(clientID, con)
    return true
}

func addConnection(clientID string, con *websocket.Conn) {
    clientData := ClientData{
        conn:  con,
        token: "",
    }
    token.New(clientID, &clientData.token, func() { //此函数用于每次token变化后将token写回客户端
        //更新完client后才能writeResponse,注意死锁
        lock.Lock()
        wsClients[clientID] = clientData
        lock.Unlock()
        writeResponseToClient(clientID, ServerResponse{
            Code:    Success,
            Method:  "/token",
            Message: "",
            Result:  map[string]string{"Token": clientData.token},
        })
    })
}

// 获取客户端连接
func getClient(clientID string) (clientConn *websocket.Conn, isOnline bool) {

    lock.RLock()
    clientData, ok := wsClients[clientID]
    lock.RUnlock()
    return clientData.conn, ok
}

// 删除WebSocket连接
func deleteClient(clientID string) {
    clientData, ok := wsClients[clientID]
    if !ok {
        return
    }
    lock.Lock()
    clientData.conn.Close()
    token.Del(clientData.token)
    delete(wsClients, clientID)
    lock.Unlock()
}

因为登录是用websocket，返回前端token的也是websoket(websoket验证成功后马上返回token)，所以websocket在，token就在，websocket关闭连接后token也结束了生命周期