go实现websocket客户端和服务端,ws连接

服务器端
main.go

package main

import "log"

func main() {
	log.Println("start server")

	server.Start()
}

server.go

package main

import (
	"log"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
)

type Server struct {
	onlineUserMap map[string]*User
	userMapLock   sync.RWMutex
}

var server = Server{
	onlineUserMap: make(map[string]*User),
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

func handler(w http.ResponseWriter, r *http.Request) {
	wsConn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
	}

	NewUser(wsConn)
}

func (server *Server) Start() {
	http.HandleFunc("/", handler)
	http.ListenAndServe(":8888", nil)
}

func (server *Server) addOnlineUserMap(user *User) {
	server.userMapLock.Lock()
	server.onlineUserMap[user.addr] = user
	server.userMapLock.Unlock()
}

func (server *Server) deleteOnlineUserMap(user *User) {
	server.userMapLock.Lock()
	delete(server.onlineUserMap, user.addr)
	server.userMapLock.Unlock()
}

user.go

package main

import (
	"log"

	"github.com/gorilla/websocket"
)

type User struct {
	addr     string
	wsConn   *websocket.Conn
	sendChan chan []byte
}

func NewUser(conn *websocket.Conn) {
	var user = &User{
		addr:     conn.RemoteAddr().String(),
		wsConn:   conn,
		sendChan: make(chan []byte),
	}

	user.online()

	go user.recvMessage()
	go user.sendMessage()
}

func (user *User) recvMessage() {
	defer user.offline()
	for {
		_, p, err := user.wsConn.ReadMessage()
		if err != nil {
			log.Println(err)
			return
		}

		user.sendChan <- p

		log.Printf("Recv [%s] msg:%s", user.addr, p)
	}
}

func (user *User) sendMessage() {
	defer user.offline()
	for {
		buf := <-user.sendChan

		err := user.wsConn.WriteMessage(1, buf)
		if err != nil {
			log.Println(err)
			return
		}

		log.Printf("Send [%s] msg:%s", user.addr, buf)
	}
}

func (user *User) online() {
	server.addOnlineUserMap(user)

	log.Printf("[%s] 上线了", user.addr)
}

func (user *User) offline() {
	user.wsConn.Close()
	server.deleteOnlineUserMap(user)

	log.Printf("[%s] 下线了", user.addr)
}

客户端实现

clint.go

package main

import (
	"bufio"
	"log"
	"os"

	"github.com/gorilla/websocket"
)

func main() {
	dl := websocket.Dialer{}
	conn, _, err := dl.Dial("ws://127.0.0.1:8888", nil)
	if err != nil {
		log.Println("建立链接失败！！")
		return
	}
	// err = conn.WriteMessage(websocket.TextMessage, body)
	go send(conn)
	for {
		m, p, e := conn.ReadMessage()
		if e != nil {
			return
		}
		log.Println(m, string(p))
	}
}

func send(conn *websocket.Conn) {
	for {

		reader := bufio.NewReader(os.Stdin)
		l, _, _ := reader.ReadLine()
		conn.WriteMessage(websocket.TextMessage, l)
	}
}