【Golang】gorilla/websocket实战和底层代码分析开发者社区

【Golang】gorilla/websocket实战和底层代码分析

原创

go get github.com/gorilla/websocket

type Upgrader struct {
    HandshakeTimeout time.Duration
	ReadBufferSize, WriteBufferSize int
	WriteBufferPool BufferPool
	Subprotocols []string
	Error func(w http.ResponseWriter, r *http.Request, status int, reason error)
	CheckOrigin func(r *http.Request) bool
	EnableCompression bool
}

type Conn struct {
	conn        net.Conn
	isServer    bool
	// Write fields
	writeBuf      []byte        
	writePool     BufferPool
	writeBufSize  int
	writer        io.WriteCloser 
	isWriting     bool           
	// Read fields
	readRemaining int64
	readFinal     bool  
	readLength    int64 
	messageReader *messageReader 
}

var upGrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
func main() {
	http.HandleFunc("/ws", wsUpGrader)
	err := http.ListenAndServe("localhost:8080", nil)
	if err != nil {
		log.Println("server start err", err)
func wsUpGrader(w http.ResponseWriter, r *http.Request) {
    //转换为升级为websocket
	conn, err := upGrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
    //释放连接
	defer conn.Close()
	for {
		//接收消息
		messageType, message, err := conn.ReadMessage()
		if err != nil {
			log.Println(err)
			return
		log.Println("server receive messageType", messageType, "message", string(message))
		//发送消息
		err = conn.WriteMessage(messageType, []byte("pong"))
		if err != nil {
			log.Println(err)
			return
}

import (
    "fmt"
    "github.com/gorilla/websocket"
    "log"
    "time"
func main() {
    //服务器地址 websocket 统一使用 ws://
    url := "ws://localhost:8080/ws" 
    //使用默认拨号器，向服务器发送连接请求
    ws, _, err := websocket.DefaultDialer.Dial(url, nil)
    if err != nil {
        log.Fatal(err)
    //关闭连接
	defer conn.Close()
    //发送消息
    go func() {
        for {
            err := ws.WriteMessage(websocket.BinaryMessage, []byte("ping"))
            if err != nil {
                log.Fatal(err)
            //休眠两秒
            time.Sleep(time.Second * 2)
    //接收消息
    for {
        _, data, err := ws.ReadMessage()
        if err != nil {
            log.Fatal(err)
		fmt.Println("client receive message: ", string(data))
}

func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header) (*Conn, error) {
	if !tokenListContainsValue(r.Header, "Connection", "upgrade") {
   		return ...
	if !tokenListContainsValue(r.Header, "Upgrade", "websocket") {
    	return ...
	//必须是get请求方法
	if r.Method != http.MethodGet {
   		return ...
	if !tokenListContainsValue(r.Header, "Sec-Websocket-Version", "13") {
    	return ...
	if _, ok := responseHeader["Sec-Websocket-Extensions"]; ok {
    	return ...
    c := newConn(netConn, true, u.ReadBufferSize, u.WriteBufferSize, u.WriteBufferPool, br, writeBuf)
}

var keyGUID = []byte("258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
func computeAcceptKey(challengeKey string) string {
	h := sha1.New() 
	h.Write([]byte(challengeKey))
	h.Write(keyGUID)
	return base64.StdEncoding.EncodeToString(h.Sum(nil))
}

func (c *Conn) ReadMessage() (messageType int, p []byte, err error) {
	var r io.Reader
	messageType, r, err = c.NextReader()
	if err != nil {
		return messageType, nil, err
    //ReadAll从r读取,直到出现错误或EOF，并返回读取的数据
	p, err = io.ReadAll(r)
	return messageType, p, err
}

func (c *Conn) NextReader() (messageType int, r io.Reader, err error) {
    	for c.readErr == nil {
        //解析数据帧方法advanceFrame
        // frameType : 帧类型
		frameType, err := c.advanceFrame()
		if err != nil {
			c.readErr = hideTempErr(err)
			break
    	//数据类型是 文本或二进制类型
		if frameType == TextMessage || frameType == BinaryMessage {
			c.messageReader = &messageReader{c}
			c.reader = c.messageReader
			if c.readDecompress {
				c.reader = c.newDecompressionReader(c.reader)
			return frameType, c.reader, nil
}

func (r *messageReader) Read(b []byte) (int, error) {
	for c.readErr == nil {
    	//当前帧中剩余的字节
		if c.readRemaining > 0 {
			if int64(len(b)) > c.readRemaining {
				b = b[:c.readRemaining]
            //读取到切片b中
			n, err := c.br.Read(b)
			c.readErr = hideTempErr(err)
            //当Conn是服务端
			if c.isServer {
				c.readMaskPos = maskBytes(c.readMaskKey, c.readMaskPos, b[:n])
            //readRemaining字节数转int64
			rem := c.readRemaining
			rem -= int64(n)
            //跟踪连接上剩余的字节数
			if err := c.setReadRemaining(rem); err != nil {
				return 0, err
			if c.readRemaining > 0 && c.readErr == io.EOF {
				c.readErr = errUnexpectedEOF
            //返回读后字节数
			return n, c.readErr
    	//标记是否最后一个数据帧
		if c.readFinal {
            // messageRader 置为nil
			c.messageReader = nil
			return 0, io.EOF
    	//获取数据帧类型
		frameType, err := c.advanceFrame()
		switch {
		case err != nil:
			c.readErr = hideTempErr(err)
		case frameType == TextMessage || frameType == BinaryMessage:
			c.readErr = errors.New("websocket: internal error, unexpected text or binary in Reader")
	err := c.readErr
	if err == io.EOF && c.messageReader == r {
		err = errUnexpectedEOF
	return 0, err
}

func ReadAll(r Reader) ([]byte, error) {
	b := make([]byte, 0, 512)
	for {
		if len(b) == cap(b) {
			// 给[]byte添加更多容量
			b = append(b, 0)[:len(b)]
		n, err := r.Read(b[len(b):cap(b)])
		b = b[:len(b)+n]
		if err != nil {
			if err == EOF {
				err = nil
			return b, err
}

func (c *Conn) WriteMessage(messageType int, data []byte) error {
    //w 是一个io.WriteCloser
	w, err := c.NextWriter(messageType)
	if err != nil {
		return err
    //将data写入writeBuf中
	if _, err = w.Write(data); err != nil {
		return err
	return w.Close()
}

func (c *Conn) NextWriter(messageType int) (io.WriteCloser, error) {
	var mw messageWriter
	if err := c.beginMessage(&mw, messageType); err != nil {
		return nil, err
	c.writer = &mw
	return c.writer, nil
}

func (w *messageWriter) Write(p []byte) (int, error) {
    //如果字节长度大于初始化的writeBuf空间大小
	if len(p) > 2*len(w.c.writeBuf) && w.c.isServer {
		//写入方法
		err := w.flushFrame(false, p)
	//字节长度不大于初始化的writeBuf空间大小
	nn := len(p)
	for len(p) > 0 {
        //内部也是调用的flushFrame
		n, err := w.ncopy(len(p))
	return nn, nil
}

func (c *Conn) advanceFrame() (int, error) {
    //读取前两个字节
    p, err := c.read(2)
	if err != nil {
		return noFrame, err
    //数据帧类型
	frameType := int(p[0] & 0xf)
    // FIN 标记位
	final := p[0]&finalBit != 0
    //三个扩展用
	rsv1 := p[0]&rsv1Bit != 0
	rsv2 := p[0]&rsv2Bit != 0
	rsv3 := p[0]&rsv3Bit != 0
    //mask ：是否使用掩码
	mask := p[1]&maskBit != 0
    switch c.readRemaining {
	case 126:
		p, err := c.read(2)
		if err != nil {
			return noFrame, err
		if err := c.setReadRemaining(int64(binary.BigEndian.Uint16(p))); err != nil {
			return noFrame, err
	case 127:
		p, err := c.read(8)
		if err != nil {
			return noFrame, err
		if err := c.setReadRemaining(int64(binary.BigEndian.Uint64(p))); err != nil {
			return noFrame, err

【Golang】gorilla/websocket实战和底层代码分析

【Golang】gorilla/websocket实战和底层代码分析

前言

简单使用

Upgrader

Conn

服务端示例

客户端示例

源码走读

Upgrade 协议升级

ReadMessage 读消息

WriteMessage 写消息

advanceFrame 解析数据帧