基于gorilla websocket的简易聊天服务器

简言

这篇文章将实现以下功能：
让用户在同"房间"内实现通讯(类似群聊)

原理

首先，我们需要一台"注册机"。用于记录和正在服务器通讯的连接，并且对这些连接进行分组。相同房号的连接会被放到同一个"房间"内
房间池会存放所有房间，这个房间池就相当于一家酒店，下面的房间就是一个一个的房间
注册机工作流程

虚线箭头代表一个新连接刚连到服务端时的注册流程，实线箭头就是注册过的连接的步骤。server中的service就是你要做的事，这里是实现一个聊天室，所以走的就是chatservice，所有连接(不管是刚注册的新连接，还是之前的老连接)走完service过后都要和注册机(hub)寒暄几句，不管你说什么，只要你张口支个声就行。然后注册机就不会把你记入"生死薄"(如果你不吱声把你记入"生死薄"了，也不代表你die了，你等下吱个声，重新去注册机那里登记一遍还是可以的，要注册机把你注销了，你才die了)

实现

import "github.com/gorilla/websocket"

var conpool = make(map[*Connection]struct{})
var roompool = make(map[string][]*Connection)

type basicmsg struct {
	Content string `json:"content"`
	Sign    string `json:"sign"`
}
type Connection struct {
	con  *websocket.Conn
	send chan []byte
}
type Hub struct {
	register   chan *Connection
	unregister chan *Connection
	Broadcast  chan []byte
}

这里的checkargs实现的是检查map中是否有我列出的key，这里可以你自己手动实现

import (
	"fmt"
	"net/http"
	"newstart/server"

	"github.com/gorilla/websocket"
)

var upgrade = websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024}
var hub = &Hub{register: make(chan *Connection), unregister: make(chan *Connection), Broadcast: make(chan []byte, 512)}

func (s *Hub) Run() {
	fmt.Println("start register machine")
	for {
		select {
		case c := <-s.register:
			conpool[c] = struct{}{}
			fmt.Println("register connection")
		case c := <-s.unregister:
			if _, ok := conpool[c]; ok {
				delete(conpool, c)
				fmt.Println("unregister connection")
			}
		case m := <-s.Broadcast:
			for clients := range conpool {
				select {
				case clients.send <- m:
				default:
					delete(conpool, clients)
					fmt.Println("delete connection")
				}
			}
		}
	}
}
func ServerStart() {
	go hub.Run()
	http.HandleFunc("/chat", ChatRoom)
	http.ListenAndServe(":8001", nil)
}
func ChatRoom(w http.ResponseWriter, r *http.Request) {
	upgrade.CheckOrigin = func(r *http.Request) bool { return true }
	ws, err := upgrade.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println(err)
		return
	}
	heads := r.URL.Query()
	if !server.CheckArgs([]string{"roomid", "usrid"}, heads) {
		fmt.Println("args not enough")
		return
	}
	roomid := heads["roomid"][0]
	// usrid:=heads["usrid"][0]
	con := &Connection{con: ws, send: make(chan []byte, 512)}
	if _, ok := roompool[roomid]; !ok {
		roompool[roomid] = []*Connection{}
	}
	conlist := roompool[roomid]
	conlist = append(conlist, con)
	roompool[roomid] = conlist
	hub.register <- con
	defer func() {
		hub.unregister <- con
		con.con.Close()
	}()
	for {
		var recmsg basicmsg
		err = con.con.ReadJSON(&recmsg)
		if err != nil {
			fmt.Println(err)
		}
		fmt.Printf("rec msg:%+v\n", recmsg)
		recmsg.Sign = "from server"
		for _, wscon := range conlist {
			err = wscon.con.WriteJSON(&recmsg)
			fmt.Println(recmsg)
			if err != nil {
				fmt.Println(err)
			}
		}
		hub.Broadcast <- []byte(recmsg.Sign)
	}
}

测试

postman