golang webSocket

  1 package info_websocket
  2 
  3 import (
  4     "crypto/sha1"
  5     "encoding/base64"
  6     "errors"
  7     "io"
  8     "log"
  9     "net"
 10     "strings"
 11 )
 12 
 13 func main() {
 14     ln, err := net.Listen("tcp", ":8000")//监听端口
 15     if err != nil {
 16         log.Panic(err)
 17     }
 18     for {
 19         log.Println("wss")
 20         conn, err := ln.Accept()//等待客户的连接
 21         if err != nil {
 22             log.Println("Accept err:", err)
 23         }
 24         for {
 25             handleConnection(conn)
 26         }
 27     }
 28 }
 29 
 30 func handleConnection(conn net.Conn) {
 31     content := make([]byte, 1024)
 32     _, err := conn.Read(content)
 33     log.Println(string(content))
 34     if err != nil {
 35         log.Println(err)
 36     }
 37     isHttp := false
 38     // 先暂时这么判断
 39     if string(content[0:3]) == "GET" {
 40         isHttp = true
 41     }
 42     log.Println("isHttp:", isHttp)
 43     if isHttp {
 44         headers := parseHandshake(string(content))
 45         log.Println("headers", headers)
 46         secWebsocketKey := headers["Sec-WebSocket-Key"]
 47         // NOTE：这里省略其他的验证
 48         guid := "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
 49         // 计算Sec-WebSocket-Accept
 50         h := sha1.New()
 51         log.Println("accept raw:", secWebsocketKey+guid)
 52         io.WriteString(h, secWebsocketKey+guid)
 53         accept := make([]byte, 28)
 54         base64.StdEncoding.Encode(accept, h.Sum(nil))
 55         log.Println(string(accept))
 56         response := "HTTP/1.1 101 Switching Protocols\r\n"
 57         response = response + "Sec-WebSocket-Accept: " + string(accept) + "\r\n"
 58         response = response + "Connection: Upgrade\r\n"
 59         response = response + "Upgrade: websocket\r\n\r\n"
 60         log.Println("response:", response)
 61         if lenth, err := conn.Write([]byte(response)); err != nil {
 62             log.Println(err)
 63         } else {
 64             log.Println("send len:", lenth)
 65         }
 66         wssocket := NewWsSocket(conn)
 67         for {
 68             data, err := wssocket.ReadIframe()
 69             if err != nil {
 70                 log.Println("readIframe err:", err)
 71             }
 72             log.Println("read data:", string(data))
 73             err = wssocket.SendIframe([]byte("good"))
 74             if err != nil {
 75                 log.Println("sendIframe err:", err)
 76             }
 77             log.Println("send data")
 78         }
 79     } else {
 80         log.Println(string(content))
 81         // 直接读取
 82     }
 83 }
 84 
 85 type WsSocket struct {
 86     MaskingKey []byte
 87     Conn       net.Conn
 88 }
 89 
 90 func NewWsSocket(conn net.Conn) *WsSocket {
 91     return &WsSocket{Conn: conn}
 92 }
 93 
 94 func (this *WsSocket) SendIframe(data []byte) error {
 95     // 这里只处理data长度<125的
 96     if len(data) >= 125 {
 97         return errors.New("send iframe data error")
 98     }
 99     lenth := len(data)
100     maskedData := make([]byte, lenth)
101     for i := 0; i < lenth; i++ {
102         if this.MaskingKey != nil {
103             maskedData[i] = data[i] ^ this.MaskingKey[i%4]
104         } else {
105             maskedData[i] = data[i]
106         }
107     }
108     this.Conn.Write([]byte{0x81})
109     var payLenByte byte
110     if this.MaskingKey != nil && len(this.MaskingKey) != 4 {
111         payLenByte = byte(0x80) | byte(lenth)
112         this.Conn.Write([]byte{payLenByte})
113         this.Conn.Write(this.MaskingKey)
114     } else {
115         payLenByte = byte(0x00) | byte(lenth)
116         this.Conn.Write([]byte{payLenByte})
117     }
118     this.Conn.Write(data)
119     return nil
120 }
121 
122 func (this *WsSocket) ReadIframe() (data []byte, err error) {
123     err = nil
124     //第一个字节：FIN + RSV1-3 + OPCODE
125     opcodeByte := make([]byte, 1)
126     this.Conn.Read(opcodeByte)
127     FIN := opcodeByte[0] >> 7
128     RSV1 := opcodeByte[0] >> 6 & 1
129     RSV2 := opcodeByte[0] >> 5 & 1
130     RSV3 := opcodeByte[0] >> 4 & 1
131     OPCODE := opcodeByte[0] & 15
132     log.Println(RSV1, RSV2, RSV3, OPCODE)
133 
134     payloadLenByte := make([]byte, 1)
135     this.Conn.Read(payloadLenByte)
136     payloadLen := int(payloadLenByte[0] & 0x7F)
137     mask := payloadLenByte[0] >> 7
138     if payloadLen == 127 {
139         extendedByte := make([]byte, 8)
140         this.Conn.Read(extendedByte)
141     }
142     maskingByte := make([]byte, 4)
143     if mask == 1 {
144         this.Conn.Read(maskingByte)
145         this.MaskingKey = maskingByte
146     }
147 
148     payloadDataByte := make([]byte, payloadLen)
149     this.Conn.Read(payloadDataByte)
150     log.Println("data:", payloadDataByte)
151     dataByte := make([]byte, payloadLen)
152     for i := 0; i < payloadLen; i++ {
153         if mask == 1 {
154             dataByte[i] = payloadDataByte[i] ^ maskingByte[i%4]
155         } else {
156             dataByte[i] = payloadDataByte[i]
157         }
158     }
159     if FIN == 1 {
160         data = dataByte
161         return
162     }
163     nextData, err := this.ReadIframe()
164     if err != nil {
165         return
166     }
167     data = append(data, nextData...)
168     return
169 }
170 
171 func parseHandshake(content string) map[string]string {
172     headers := make(map[string]string, 10)
173     lines := strings.Split(content, "\r\n")
174     for _, line := range lines {
175         if len(line) >= 0 {
176             words := strings.Split(line, ":")
177             if len(words) == 2 {
178                 headers[strings.Trim(words[0], " ")] = strings.Trim(words[1], " ")
179             }
180         }
181     }
182     return headers
183 }