Create & Init Project...

app/service/main/broadcast/libs/websocket/BUILD

@@ -0,0 +1,33 @@
+package(default_visibility = ["//visibility:public"])
+
+load(
+    "@io_bazel_rules_go//go:def.bzl",
+    "go_library",
+)
+
+go_library(
+    name = "go_default_library",
+    srcs = [
+        "conn.go",
+        "request.go",
+        "server.go",
+    ],
+    importpath = "go-common/app/service/main/broadcast/libs/websocket",
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+    deps = ["//app/service/main/broadcast/libs/bufio:go_default_library"],
+)
+
+filegroup(
+    name = "package-srcs",
+    srcs = glob(["**"]),
+    tags = ["automanaged"],
+    visibility = ["//visibility:private"],
+)
+
+filegroup(
+    name = "all-srcs",
+    srcs = [":package-srcs"],
+    tags = ["automanaged"],
+    visibility = ["//visibility:public"],
+)

app/service/main/broadcast/libs/websocket/conn.go

@@ -0,0 +1,256 @@
+package websocket
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+
+	"go-common/app/service/main/broadcast/libs/bufio"
+)
+
+const (
+	// Frame header byte 0 bits from Section 5.2 of RFC 6455
+	finBit  = 1 << 7
+	rsv1Bit = 1 << 6
+	rsv2Bit = 1 << 5
+	rsv3Bit = 1 << 4
+	opBit   = 0x0f
+
+	// Frame header byte 1 bits from Section 5.2 of RFC 6455
+	maskBit = 1 << 7
+	lenBit  = 0x7f
+
+	continuationFrame        = 0
+	continuationFrameMaxRead = 100
+)
+
+// The message types are defined in RFC 6455, section 11.8.
+const (
+	// TextMessage denotes a text data message. The text message payload is
+	// interpreted as UTF-8 encoded text data.
+	TextMessage = 1
+
+	// BinaryMessage denotes a binary data message.
+	BinaryMessage = 2
+
+	// CloseMessage denotes a close control message. The optional message
+	// payload contains a numeric code and text. Use the FormatCloseMessage
+	// function to format a close message payload.
+	CloseMessage = 8
+
+	// PingMessage denotes a ping control message. The optional message payload
+	// is UTF-8 encoded text.
+	PingMessage = 9
+
+	// PongMessage denotes a ping control message. The optional message payload
+	// is UTF-8 encoded text.
+	PongMessage = 10
+)
+
+var (
+	// ErrMessageClose close control message
+	ErrMessageClose = errors.New("close control message")
+	// ErrMessageMaxRead continuation frrame max read
+	ErrMessageMaxRead = errors.New("continuation frame max read")
+)
+
+// Conn represents a WebSocket connection.
+type Conn struct {
+	rwc io.ReadWriteCloser
+	r   *bufio.Reader
+	w   *bufio.Writer
+}
+
+// new connection
+func newConn(rwc io.ReadWriteCloser, r *bufio.Reader, w *bufio.Writer) *Conn {
+	return &Conn{rwc: rwc, r: r, w: w}
+}
+
+// WriteMessage write a message by type.
+func (c *Conn) WriteMessage(msgType int, msg []byte) (err error) {
+	if err = c.WriteHeader(msgType, len(msg)); err != nil {
+		return
+	}
+	err = c.WriteBody(msg)
+	return
+}
+
+// WriteHeader write header frame.
+func (c *Conn) WriteHeader(msgType int, length int) (err error) {
+	var h []byte
+	if h, err = c.w.Peek(2); err != nil {
+		return
+	}
+	// 1.First byte. FIN/RSV1/RSV2/RSV3/OpCode(4bits)
+	h[0] = 0
+	h[0] |= finBit | byte(msgType)
+	// 2.Second byte. Mask/Payload len(7bits)
+	h[1] = 0
+	switch {
+	case length <= 125:
+		// 7 bits
+		h[1] |= byte(length)
+	case length < 65536:
+		// 16 bits
+		h[1] |= 126
+		if h, err = c.w.Peek(2); err != nil {
+			return
+		}
+		binary.BigEndian.PutUint16(h, uint16(length))
+	default:
+		// 64 bits
+		h[1] |= 127
+		if h, err = c.w.Peek(8); err != nil {
+			return
+		}
+		binary.BigEndian.PutUint64(h, uint64(length))
+	}
+	return
+}
+
+// WriteBody write a message body.
+func (c *Conn) WriteBody(b []byte) (err error) {
+	if len(b) > 0 {
+		_, err = c.w.Write(b)
+	}
+	return
+}
+
+// Peek write peek.
+func (c *Conn) Peek(n int) ([]byte, error) {
+	return c.w.Peek(n)
+}
+
+// Flush flush writer buffer
+func (c *Conn) Flush() error {
+	return c.w.Flush()
+}
+
+// ReadMessage read a message.
+func (c *Conn) ReadMessage() (op int, payload []byte, err error) {
+	var (
+		fin         bool
+		finOp, n    int
+		partPayload []byte
+	)
+	for {
+		// read frame
+		if fin, op, partPayload, err = c.readFrame(); err != nil {
+			return
+		}
+		switch op {
+		case BinaryMessage, TextMessage, continuationFrame:
+			if fin && len(payload) == 0 {
+				return op, partPayload, nil
+			}
+			// continuation frame
+			payload = append(payload, partPayload...)
+			if op != continuationFrame {
+				finOp = op
+			}
+			// final frame
+			if fin {
+				op = finOp
+				return
+			}
+		case PingMessage:
+			// handler ping
+			if err = c.WriteMessage(PongMessage, partPayload); err != nil {
+				return
+			}
+		case PongMessage:
+			// handler pong
+		case CloseMessage:
+			// handler close
+			err = ErrMessageClose
+			return
+		default:
+			err = fmt.Errorf("unknown control message, fin=%t, op=%d", fin, op)
+			return
+		}
+		if n > continuationFrameMaxRead {
+			err = ErrMessageMaxRead
+			return
+		}
+		n++
+	}
+}
+
+func (c *Conn) readFrame() (fin bool, op int, payload []byte, err error) {
+	var (
+		b          byte
+		p          []byte
+		mask       bool
+		maskKey    []byte
+		payloadLen int64
+	)
+	// 1.First byte. FIN/RSV1/RSV2/RSV3/OpCode(4bits)
+	b, err = c.r.ReadByte()
+	if err != nil {
+		return
+	}
+	// final frame
+	fin = (b & finBit) != 0
+	// rsv MUST be 0
+	if rsv := b & (rsv1Bit | rsv2Bit | rsv3Bit); rsv != 0 {
+		return false, 0, nil, fmt.Errorf("unexpected reserved bits rsv1=%d, rsv2=%d, rsv3=%d", b&rsv1Bit, b&rsv2Bit, b&rsv3Bit)
+	}
+	// op code
+	op = int(b & opBit)
+	// 2.Second byte. Mask/Payload len(7bits)
+	b, err = c.r.ReadByte()
+	if err != nil {
+		return
+	}
+	// is mask payload
+	mask = (b & maskBit) != 0
+	// payload length
+	switch b & lenBit {
+	case 126:
+		// 16 bits
+		if p, err = c.r.Pop(2); err != nil {
+			return
+		}
+		payloadLen = int64(binary.BigEndian.Uint16(p))
+	case 127:
+		// 64 bits
+		if p, err = c.r.Pop(8); err != nil {
+			return
+		}
+		payloadLen = int64(binary.BigEndian.Uint64(p))
+	default:
+		// 7 bits
+		payloadLen = int64(b & lenBit)
+	}
+	// read mask key
+	if mask {
+		maskKey, err = c.r.Pop(4)
+		if err != nil {
+			return
+		}
+	}
+	// read payload
+	if payloadLen > 0 {
+		if payload, err = c.r.Pop(int(payloadLen)); err != nil {
+			return
+		}
+		if mask {
+			maskBytes(maskKey, 0, payload)
+		}
+	}
+	return
+}
+
+// Close close the connection.
+func (c *Conn) Close() error {
+	return c.rwc.Close()
+}
+
+func maskBytes(key []byte, pos int, b []byte) int {
+	for i := range b {
+		b[i] ^= key[pos&3]
+		pos++
+	}
+	return pos & 3
+}

app/service/main/broadcast/libs/websocket/request.go

@@ -0,0 +1,115 @@
+package websocket
+
+import (
+	"bytes"
+	"fmt"
+	"net/http"
+	"strings"
+
+	"go-common/app/service/main/broadcast/libs/bufio"
+)
+
+// Request request.
+type Request struct {
+	Method     string
+	RequestURI string
+	Proto      string
+	Host       string
+	Header     http.Header
+
+	reader *bufio.Reader
+}
+
+// ReadRequest reads and parses an incoming request from b.
+func ReadRequest(r *bufio.Reader) (req *Request, err error) {
+	var (
+		b  []byte
+		ok bool
+	)
+	req = &Request{reader: r}
+	if b, err = req.readLine(); err != nil {
+		return
+	}
+	if req.Method, req.RequestURI, req.Proto, ok = parseRequestLine(string(b)); !ok {
+		return nil, fmt.Errorf("malformed HTTP request %s", b)
+	}
+	if req.Header, err = req.readMIMEHeader(); err != nil {
+		return
+	}
+	req.Host = req.Header.Get("Host")
+	return req, nil
+}
+
+func (r *Request) readLine() ([]byte, error) {
+	var line []byte
+	for {
+		l, more, err := r.reader.ReadLine()
+		if err != nil {
+			return nil, err
+		}
+		// Avoid the copy if the first call produced a full line.
+		if line == nil && !more {
+			return l, nil
+		}
+		line = append(line, l...)
+		if !more {
+			break
+		}
+	}
+	return line, nil
+}
+
+func (r *Request) readMIMEHeader() (header http.Header, err error) {
+	var (
+		line []byte
+		i    int
+		k, v string
+	)
+	header = make(http.Header, 16)
+	for {
+		if line, err = r.readLine(); err != nil {
+			return
+		}
+		line = trim(line)
+		if len(line) == 0 {
+			return
+		}
+		if i = bytes.IndexByte(line, ':'); i <= 0 {
+			err = fmt.Errorf("malformed MIME header line: " + string(line))
+			return
+		}
+		k = string(line[:i])
+		// Skip initial spaces in value.
+		i++ // skip colon
+		for i < len(line) && (line[i] == ' ' || line[i] == '\t') {
+			i++
+		}
+		v = string(line[i:])
+		header.Add(k, v)
+	}
+}
+
+// parseRequestLine parses "GET /foo HTTP/1.1" into its three parts.
+func parseRequestLine(line string) (method, requestURI, proto string, ok bool) {
+	s1 := strings.Index(line, " ")
+	s2 := strings.Index(line[s1+1:], " ")
+	if s1 < 0 || s2 < 0 {
+		return
+	}
+	s2 += s1 + 1
+	return line[:s1], line[s1+1 : s2], line[s2+1:], true
+}
+
+// trim returns s with leading and trailing spaces and tabs removed.
+// It does not assume Unicode or UTF-8.
+func trim(s []byte) []byte {
+	i := 0
+	for i < len(s) && (s[i] == ' ' || s[i] == '\t') {
+		i++
+	}
+	n := len(s)
+	for n > i && (s[n-1] == ' ' || s[n-1] == '\t') {
+		n--
+	}
+	return s[i:n]
+}

app/service/main/broadcast/libs/websocket/server.go

@@ -0,0 +1,56 @@
+package websocket
+
+import (
+	"crypto/sha1"
+	"encoding/base64"
+	"errors"
+	"io"
+	"strings"
+
+	"go-common/app/service/main/broadcast/libs/bufio"
+)
+
+var (
+	keyGUID = []byte("258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
+	// ErrBadRequestMethod bad request method
+	ErrBadRequestMethod = errors.New("bad method")
+	// ErrNotWebSocket not websocket protocal
+	ErrNotWebSocket = errors.New("not websocket protocol")
+	// ErrBadWebSocketVersion bad websocket version
+	ErrBadWebSocketVersion = errors.New("missing or bad WebSocket Version")
+	// ErrChallengeResponse mismatch challenge response
+	ErrChallengeResponse = errors.New("mismatch challenge/response")
+)
+
+// Upgrade Switching Protocols
+func Upgrade(rwc io.ReadWriteCloser, rr *bufio.Reader, wr *bufio.Writer, req *Request) (conn *Conn, err error) {
+	if req.Method != "GET" {
+		return nil, ErrBadRequestMethod
+	}
+	if req.Header.Get("Sec-Websocket-Version") != "13" {
+		return nil, ErrBadWebSocketVersion
+	}
+	if strings.ToLower(req.Header.Get("Upgrade")) != "websocket" {
+		return nil, ErrNotWebSocket
+	}
+	if !strings.Contains(strings.ToLower(req.Header.Get("Connection")), "upgrade") {
+		return nil, ErrNotWebSocket
+	}
+	challengeKey := req.Header.Get("Sec-Websocket-Key")
+	if challengeKey == "" {
+		return nil, ErrChallengeResponse
+	}
+	wr.WriteString("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n")
+	wr.WriteString("Sec-WebSocket-Accept: " + computeAcceptKey(challengeKey) + "\r\n\r\n")
+	if err = wr.Flush(); err != nil {
+		return
+	}
+	return newConn(rwc, rr, wr), nil
+}
+
+func computeAcceptKey(challengeKey string) string {
+	h := sha1.New()
+	h.Write([]byte(challengeKey))
+	h.Write(keyGUID)
+	return base64.StdEncoding.EncodeToString(h.Sum(nil))
+}