websocket封裝

EddieChan發表於2018-01-12
package edd_socket

    import (
        "time"
        "log"
        "encoding/json"
        "golang.org/x/net/websocket"
    )

    type ws struct {
        coon *websocket.Conn
    }

    //訊息體
    type Message struct {
        Data      interface{} `json:"content"`
        Type      string      `json:"type"`
        TimeStamp int64       `json:"time_stamp"`
    }

    //使用者體
    type User struct {
        Uid  string
        conn *websocket.Conn
    }

    var (
        member         = make(map[string]*User)
        uidMapWs       = make(map[string]*websocket.Conn)
        groupMapMember = make(map[string][]*User)
    )

    //例項化ws
    func NewWS(wss *websocket.Conn) *ws {
        return &ws{coon: wss}
    }

    //繫結uid和coon
    func (this *ws) BindUid(uid string) {
        client := User{Uid: uid,conn: this.coon}

        member[uid] = &client
        uidMapWs[uid] = this.coon
    }

    //是否線上
    func (this *ws) IsOnline(uid string) bool {
        _, exits := member[uid]
        return exits
    }

    //斷開連線
    func (this *ws) CloseUid(uid string) {
        msg := Message{
            Data: "離開房間",
            Type: "del_user",
        }
        this.SendToAll(msg)

        delete(member, uid)
        delete(uidMapWs, uid)
        this.coon.Close()
    }

    //群發訊息
    func (this *ws) SendToAll(msg Message) {
        msg.TimeStamp = time.Now().Unix()
        sendMess, _ := json.Marshal(msg)

        for k, v := range member {
            if v.conn != this.coon {
                if err := websocket.Message.Send(v.conn, string(sendMess)); err != nil {
                    delete(member, k)
                    delete(uidMapWs, k)
                    continue
                }
            }
        }
    }

    //獲取當前組人數
    func (this *ws) GetClientCountByGroup(groupName string) int {
        return len(groupMapMember[groupName])
    }

    func (this *ws) GetClientByGroup(groupName string) []*User {
        return groupMapMember[groupName]
    }

    //加入某個群
    func (this *ws) JoinGroup(groupName, uid string) {
        groupMapMember[groupName] = append(groupMapMember[groupName], member[uid])
    }

    //給指定群發訊息
    func (this *ws) SendToGroup(groupName string, msg Message) {
        msg.TimeStamp = time.Now().Unix()
        sendMess, _ := json.Marshal(msg)

        for k, v := range groupMapMember[groupName] {
            if v.conn != this.coon {
                if err := websocket.Message.Send(v.conn, string(sendMess)); err != nil {
                    kk := k + 1
                    groupMapMember[groupName] = append(groupMapMember[groupName][:k], groupMapMember[groupName][kk:]...)
                    continue
                }
            }
        }
    }

    //離開某個群
    func (this *ws) LeaveGroup(groupName, uid string) {
        for k, v := range groupMapMember[groupName] {
            if v.Uid == uid {
                kk := k + 1
                groupMapMember[groupName] = append(groupMapMember[groupName][:k], groupMapMember[groupName][kk:] ...)
                break
            }
        }
    }

    //傳送給指定uid
    func (this *ws) SendToUid(uid string, msg Message) {
        toWsCoon := uidMapWs[uid]
        msg.TimeStamp = time.Now().Unix()
        sendMess, _ := json.Marshal(msg)

        if err := websocket.Message.Send(toWsCoon, string(sendMess)); err != nil {
            delete(member, uid)
            log.Println(err)
        }
    }

    //解析客戶端訊息
    func (this *ws) GetMsg(msg *Message) error {
        var reply string
        var err error
        if err = websocket.Message.Receive(this.coon, &reply); err == nil {
            json.Unmarshal([]byte(reply), msg)
        }
        return err
    }

使用:

package edd_socket

import (
    "fmt"
    "net/http"
    "crypto/md5"
    "encoding/hex"
    "io"
    "crypto/rand"
    "encoding/base64"
    "log"
    "golang.org/x/net/websocket"
)

//獲取MD5字串
func getMd5String(s string) string {
    h := md5.New()
    h.Write([]byte(s))
    return hex.EncodeToString(h.Sum(nil))
}

//用於將當前使用者資訊和當前websocket.conn關聯
//正式上線,可用資料庫使用者Id代替
func getUid() string {
    b := make([]byte, 48)
    if _, err := io.ReadFull(rand.Reader, b); err != nil {
        return ""
    }

    return getMd5String(base64.URLEncoding.EncodeToString(b))
}

func Echo(ws *websocket.Conn) {
    wss := NewWS(ws)
    uid := getUid()

    defer func() {
        wss.CloseUid(uid)
    }()

    for {
        var mess Message
        if err := wss.GetMsg(&mess); err != nil {
            //若當前socket接收不到訊息,則已掉線,主動退出for
            fmt.Println(err)
            break
        }

        switch mess.Type {
        case "connect":
            //連線準備,客戶端識別符號name和服務端識別符號繫結
            wss.BindUid(uid)
            msg := Message{
                Data: mess.Data,
                Type: "join_room",
            }

            wss.SendToAll(msg)
        case "all":
            //群發
            msg := Message{
                Data: mess.Data,
                Type: "send_all",
            }
            wss.SendToAll(msg)
        case "join_group":
            //wss.JoinGroup("one", uid)
        case "who":
        }
    }
}

func main() {
    var port = "8022"
    fmt.Println("listening port:" + port)

    go func() {
        http.Handle("/chat", websocket.Handler(Echo))
    }()

    if err := http.ListenAndServe(":"+port, nil); err != nil {
        log.Fatal("ListenAndServer: ", err)
    }
}

客戶端

  var ws =new WebSocket(websocket_domain)
    ws.onopen = function (res) {
        ws.send(sendMessAll(name,"進入房間",'connect'))
    };

    ws.onmessage = function (res) {
        var data = jQuery.parseJSON(res.data);
    };
ws.onclose = function (res) {
        console.log(res)
};
function sendMessAll(name, data, type) {
    return JSON.stringify(
        {
            data: JSON.stringify({
                name:name,
                age:123
            }),
            type: type
        }
    )
}

歡迎指正我的開源地址

相關文章