Golang语言websocket源码

package websocket    
import (    
"bufio"    
"bytes"    
"crypto/sha1"    
"encoding/base64"    
"encoding/binary"    
"errors"    
"io"    
"net"    
"net/http"    
"strings"    
)    
var (    
ErrUpgrade = errors.New("Can "Upgrade" only to "WebSocket"")    
ErrConnection = errors.New(""Connection" must be "Upgrade"")    
ErrCrossOrigin = errors.New("Cross origin websockets not allowed")    
ErrSecVersion = errors.New("HTTP/1.1 Upgrade RequiredrnSec-WebSocket-Version: 13rnrn")    
ErrSecKey = errors.New(""Sec-WebSocket-Key" must not be nil")    
ErrHijacker = errors.New("Not implement http.Hijacker")    
)    
var (    
ErrReservedBits = errors.New("Reserved_bits show using undefined extensions")    
ErrFrameOverload = errors.New("Control frame payload overload")    
ErrFrameFragmented = errors.New("Control frame must not be fragmented")    
ErrInvalidOpcode = errors.New("Invalid frame opcode")    
)    
var (    
crlf = []byte("rn")    
challengeKey = []byte("258EAFA5-E914-47DA-95CA-C5AB0DC85B11")    
)    
//referer https://github.com/Skycrab/skynet_websocket/blob/master/websocket.lua    
type WsHandler interface {    
CheckOrigin(origin, host string) bool    
OnOpen(ws *Websocket)    
OnMessage(ws *Websocket, message []byte)    
OnClose(ws *Websocket, code uint16, reason []byte)    
OnPong(ws *Websocket, data []byte)    
}    
type WsDefaultHandler struct {    
checkOriginOr bool // 是否校验origin, default true    
}    
func (wd WsDefaultHandler) CheckOrigin(origin, host string) bool {    
return true    
}    
func (wd WsDefaultHandler) OnOpen(ws *Websocket) {    
}    
func (wd WsDefaultHandler) OnMessage(ws *Websocket, message []byte) {    
}    
func (wd WsDefaultHandler) OnClose(ws *Websocket, code uint16, reason []byte) {    
}    
func (wd WsDefaultHandler) OnPong(ws *Websocket, data []byte) {    
}    
type Websocket struct {    
conn net.Conn    
rw *bufio.ReadWriter    
handler WsHandler    
clientTerminated bool    
serverTerminated bool    
maskOutgoing bool    
}    
type Option struct {    
Handler WsHandler // 处理器, default WsDefaultHandler    
MaskOutgoing bool //发送frame是否mask, default false    
}    
func challengeResponse(key, protocol string) []byte {    
sha := sha1.New()    
sha.Write([]byte(key))    
sha.Write(challengeKey)    
accept := base64.StdEncoding.EncodeToString(sha.Sum(nil))    
buf := bytes.NewBufferString("HTTP/1.1 101 Switching ProtocolsrnUpgrade: websocketrnConnection: UpgradernSec-WebSocket-Accept: ")    
buf.WriteString(accept)    
buf.Write(crlf)    
if protocol != "" {    
buf.WriteString("Sec-WebSocket-Protocol: ")    
buf.WriteString(protocol)    
buf.Write(crlf)    
}    
buf.Write(crlf)    
return buf.Bytes()    
}    
func acceptConnection(r *http.Request, h WsHandler) (challenge []byte, err error) {    
//Upgrade header should be present and should be equal to WebSocket    
if strings.ToLower(r.Header.Get("Upgrade")) != "websocket" {    
return nil, ErrUpgrade    
}    
//Connection header should be upgrade. Some proxy servers/load balancers    
// might mess with it.    
if !strings.Contains(strings.ToLower(r.Header.Get("Connection")), "upgrade") {    
return nil, ErrConnection    
}    
// Handle WebSocket Origin naming convention differences    
// The difference between version 8 and 13 is that in 8 the    
// client sends a "Sec-Websocket-Origin" header and in 13 it's    
// simply "Origin".    
if r.Header.Get("Sec-Websocket-Version") != "13" {    
return nil, ErrSecVersion    
}    
origin := r.Header.Get("Origin")    
if origin == "" {    
origin = r.Header.Get("Sec-Websocket-Origin")    
}    
if origin != "" && !h.CheckOrigin(origin, r.Header.Get("Host")) {    
return nil, ErrCrossOrigin    
}    
key := r.Header.Get("Sec-Websocket-Key")    
if key == "" {    
return nil, ErrSecKey    
}    
protocol := r.Header.Get("Sec-Websocket-Protocol")    
if protocol != "" {    
idx := strings.IndexByte(protocol, ',')    
if idx != -1 {    
protocol = protocol[:idx]    
}    
}    
return challengeResponse(key, protocol), nil    
}    
func websocketMask(mask []byte, data []byte) {    
for i := range data {    
data[i] ^= mask[i%4]    
}    
}    
func New(w http.ResponseWriter, r *http.Request, opt *Option) (*Websocket, error) {    
var h WsHandler    
var maskOutgoing bool    
if opt == nil {    
h = WsDefaultHandler{true}    
maskOutgoing = false    
} else {    
h = opt.Handler    
maskOutgoing = opt.MaskOutgoing    
}    
challenge, err := acceptConnection(r, h)    
if err != nil {    
var code int    
if err == ErrCrossOrigin {    
code = 403    
} else {    
code = 400    
}    
w.WriteHeader(code)    
w.Write([]byte(err.Error()))    
return nil, err    
}    
hj, ok := w.(http.Hijacker)    
if !ok {    
return nil, ErrHijacker    
}    
conn, rw, err := hj.Hijack()    
ws := new(Websocket)    
ws.conn = conn    
ws.rw = rw    
ws.handler = h    
ws.maskOutgoing = maskOutgoing    
if _, err := ws.conn.Write(challenge); err != nil {    
ws.conn.Close()    
return nil, err    
}    
ws.handler.OnOpen(ws)    
return ws, nil    
}    
func (ws *Websocket) read(buf []byte) error {    
_, err := io.ReadFull(ws.rw, buf)    
return err    
}    
func (ws *Websocket) SendFrame(fin bool, opcode byte, data []byte) error {    
//max frame header may 14 length    
buf := make([]byte, 0, len(data)+14)    
var finBit, maskBit byte    
if fin {    
finBit = 0x80    
} else {    
finBit = 0    
}    
buf = append(buf, finBit|opcode)    
length := len(data)    
if ws.maskOutgoing {    
maskBit = 0x80    
} else {    
maskBit = 0    
}    
if length < 126 {    
buf = append(buf, byte(length)|maskBit)    
} else if length < 0xFFFF {    
buf = append(buf, 126|maskBit, 0, 0)    
binary.BigEndian.PutUint16(buf[len(buf)-2:], uint16(length))    
} else {    
buf = append(buf, 127|maskBit, 0, 0, 0, 0, 0, 0, 0, 0)    
binary.BigEndian.PutUint64(buf[len(buf)-8:], uint64(length))    
}    
if ws.maskOutgoing {    
}    
buf = append(buf, data...)    
ws.rw.Write(buf)    
return ws.rw.Flush()    
}    
func (ws *Websocket) SendText(data []byte) error {    
return ws.SendFrame(true, 0x1, data)    
}    
func (ws *Websocket) SendBinary(data []byte) error {    
return ws.SendFrame(true, 0x2, data)    
}    
func (ws *Websocket) SendPing(data []byte) error {    
return ws.SendFrame(true, 0x9, data)    
}    
func (ws *Websocket) SendPong(data []byte) error {    
return ws.SendFrame(true, 0xA, data)    
}    
func (ws *Websocket) Close(code uint16, reason []byte) {    
if !ws.serverTerminated {    
data := make([]byte, 0, len(reason)+2)    
if code == 0 && reason != nil {    
code = 1000    
}    
if code != 0 {    
data = append(data, 0, 0)    
binary.BigEndian.PutUint16(data, code)    
}    
if reason != nil {    
data = append(data, reason...)    
}    
ws.SendFrame(true, 0x8, data)    
ws.serverTerminated = true    
}    
if ws.clientTerminated {    
ws.conn.Close()    
}    
}    
func (ws *Websocket) RecvFrame() (final bool, message []byte, err error) { //text 数据报文    
buf := make([]byte, 8, 8)    
err = ws.read(buf[:2])    
if err != nil {    
return    
}    
header, payload := buf[0], buf[1]    
final = header&0x80 != 0    
reservedBits := header&0x70 != 0    
frameOpcode := header & 0xf    
frameOpcodeIsControl := frameOpcode&0x8 != 0    
if reservedBits {    
// client is using as-yet-undefined extensions    
err = ErrReservedBits    
return    
}    
maskFrame := payload&0x80 != 0    
payloadlen := uint64(payload & 0x7f)    
if frameOpcodeIsControl && payloadlen >= 126 {    
err = ErrFrameOverload    
return    
}    
if frameOpcodeIsControl && !final {    
err = ErrFrameFragmented    
return    
}    
//解析frame长度    
var frameLength uint64    
if payloadlen < 126 {    
frameLength = payloadlen    
} else if payloadlen == 126 {    
err = ws.read(buf[:2])    
if err != nil {    
return    
}    
frameLength = uint64(binary.BigEndian.Uint16(buf[:2]))    
} else { //payloadlen == 127    
err = ws.read(buf[:8])    
if err != nil {    
return    
}    
frameLength = binary.BigEndian.Uint64(buf[:8])    
}    
frameMask := make([]byte, 4, 4)    
if maskFrame {    
err = ws.read(frameMask)    
if err != nil {    
return    
}    
}    
// fmt.Println("final_frame:", final, "frame_opcode:", frameOpcode, "mask_frame:", maskFrame, "frame_length:", frameLength)    
message = make([]byte, frameLength, frameLength)    
if frameLength > 0 {    
err = ws.read(message)    
if err != nil {    
return    
}    
}    
if maskFrame && frameLength > 0 {    
websocketMask(frameMask, message)    
}    
if !final {    
return    
} else {    
switch frameOpcode {    
case 0x1: //text    
case 0x2: //binary    
case 0x8: // close    
var code uint16    
var reason []byte    
if frameLength >= 2 {    
code = binary.BigEndian.Uint16(message[:2])    
}    
if frameLength > 2 {    
reason = message[2:]    
}    
message = nil    
ws.clientTerminated = true    
ws.Close(0, nil)    
ws.handler.OnClose(ws, code, reason)    
case 0x9: //ping    
message = nil    
ws.SendPong(nil)    
case 0xA:    
ws.handler.OnPong(ws, message)    
message = nil    
default:    
err = ErrInvalidOpcode    
}    
return    
}    
}    
func (ws *Websocket) Recv() ([]byte, error) {    
data := make([]byte, 0, 8)    
for {    
final, message, err := ws.RecvFrame()    
if final {    
data = append(data, message...)    
break    
} else {    
data = append(data, message...)    
}    
if err != nil {    
return data, err    
}    
}    
if len(data) > 0 {    
ws.handler.OnMessage(ws, data)    
}    
return data, nil    
}    
func (ws *Websocket) Start() {    
for {    
_, err := ws.Recv()    
if err != nil {    
ws.conn.Close()    
}    
}    
}