golang websocket简单实现

包

go get github.com/gorilla/websocket

相关方法

获取连接

func (u *Upgrader) Upgrade(w http.ResponseWriter,r *http.Request,responseHeader http.Header)(*Conn , error)

关闭连接

func(c *Conn) Close() error

代码实现

package main

import (
	"fmt"
	"github.com/gorilla/websocket"
	"net/http"
	"runtime/debug"
)

var msg chan []byte
var isClose = false
func main()  {

		http.HandleFunc("/", func(writer http.ResponseWriter, request *http.Request) {

			isClose = false

			wsClient , _ := (&websocket.Upgrader{
					CheckOrigin: func (r *http.Request) bool {
						return true
					},
			}).Upgrade(writer, request, nil)


			msg = make(chan []byte, 200)


			// 设置关闭处理器
			wsClient.SetCloseHandler(func (closeCode int,text string) error {
				isClose = true
				wsClient.Close()
				fmt.Println("关闭处理代码")
				return nil
			})
			
			//启用两个协程分别处理读取消息与写消息
			go read(wsClient, msg)
			go write(wsClient)
	})


	http.ListenAndServe("127.0.0.1:8010", nil)
}

func read(ws *websocket.Conn, msg chan []byte)  {
	defer func() {
		if r := recover(); r != nil {
			fmt.Println("write stop", string(debug.Stack()), r)
		}
	}()

	for  {
		_, message , ok :=  ws.ReadMessage()

		if ok != nil {
			ws.Close()
			return
		}

		fmt.Println("读消息")
		fmt.Println(string(message))
		msg <- message
	}

}



/**

 */
func write(ws *websocket.Conn)  {

	defer func() {
		if r := recover(); r != nil {
			fmt.Println("write stop", string(debug.Stack()), r)
		}
	}()

	for {
		if isClose{
			break
		}

		select {
			case message , err  := <- msg:

				if !err {
					ws.Close() // 关闭连接
					fmt.Println("断开连接")
					return
				}
				fmt.Println("写消息")
				_ = ws.WriteMessage(websocket.TextMessage, message)
				fmt.Println("读取客户端数据 处理:", string(message))

		}
		// 处理程序
	}
	//fmt.Println(err)
}