Go language websocket Source code

This is a creation in Article, where the information may have evolved or changed.

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 required\r\nsec-websocket-version: 13\r\n\r\n")      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 ("\ r \ n")     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 //  whether to check origin,  default true    }    func  (Wd wsdefaulthandler)  checkorigin (origin, host string)  bool {    return true     }    func  (Wd wsdefaulthandler)  onopen (ws * Websocket) &NBSP;{&NBSP;&NBSP;&NBSP;&NBSP;}     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 //  processor,  default wsdefaulthandler    maskoutgoing bool //Send frame whether 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 protocols\r\nupgrade: websocket\r\nconnection:  Upgrade\r\nSec-WebSocket-Accept:  ")     buf. WriteString (Accept)     buf. Write (CRLF)     if protocol !=  " {    buf". WriteString ("sec-websocket-protocol: ")     buf. WriteString (protocol)     buf. Write (CRLF) &NBSP;&NBSP;&NBSP;&NBSP;}&NBSP;&NBsp;  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")  !=  " {   "  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]   & NBsP;}     }    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 ())) &NBSP;&NBSP;&NBSP;&NBSP;RETURN&NBSP;NIL,&NBSP;ERR&NBSP;&NBSP;&NBSP;&NBSP;}&NBSP;&NBSP;&NBSP;&NBSP;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,&NBSP;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 {   &nBSP;}     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,&NBSP;ERR&NBSP;ERROR)  { //text  Data message     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    }     //Parsing Frame Length     var frameLength uint64     if payloadlen < 126 {    framelength = payloadlen     } else&nBsp;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, &NBSP;4,&NBSP;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) &NBSP;&NBsp;  }    return data, nil    }     func  (Ws *websocket)  start ()  {    for {     _, err := ws. RECV ()     if err != nil {    ws.conn.close ()      }    }    }