Merge pull request #9 from lxzan/dev

v1.4.8 release

README.md

@@ -66,7 +66,7 @@ type Event interface {
 #### Examples
 
 - [chat room](examples/chatroom/main.go)
-- [echo](examples/wss-server/server.go)
+- [echo](examples/echo/main.go)
 
 #### Quick Start
 
@@ -92,12 +92,12 @@ func main() {
 	})
 
 	http.HandleFunc("/connect", func(writer http.ResponseWriter, request *http.Request) {
-		socket, err := upgrader.Accept(writer, request)
+		socket, err := upgrader.Upgrade(writer, request)
 		if err != nil {
 			log.Printf("Accept: " + err.Error())
 			return
 		}
-		go socket.Listen()
+		socket.ReadLoop()
 	})
 
 	if err := http.ListenAndServe(":3000", nil); err != nil {
@@ -135,7 +135,7 @@ func main() {
 		log.Printf(err.Error())
 		return
 	}
-	socket.Listen()
+	socket.ReadLoop()
 }
 
 type WebSocket struct {
@@ -175,11 +175,11 @@ func main() {
 	app := gin.New()
 	upgrader := gws.NewUpgrader(new(WebSocket), nil)
 	app.GET("/connect", func(ctx *gin.Context) {
-		socket, err := upgrader.Accept(ctx.Writer, ctx.Request)
+		socket, err := upgrader.Upgrade(ctx.Writer, ctx.Request)
 		if err != nil {
 			return
 		}
-		go upgrader.Listen(socket)
+		go upgrader.ReadLoop(socket)
 	})
 	if err := app.Run(":8080"); err != nil {
 		panic(err)
@@ -190,7 +190,7 @@ func main() {
 - HeartBeat
 
 ```go
-const PingInterval = 5 * time.Second
+const PingInterval = 10 * time.Second
 
 type Websocket struct {
 	gws.BuiltinEventHandler

compress.go

@@ -16,33 +16,25 @@ func newCompressor(level int) *compressor {
 
 // 压缩器
 type compressor struct {
-	buffer *bytes.Buffer
-	fw     *flate.Writer
-}
-
-func (c *compressor) Close() {
-	_bpool.Put(c.buffer)
-	c.buffer = nil
+	fw *flate.Writer
 }
 
 // Compress 压缩
-func (c *compressor) Compress(content *bytes.Buffer) (*bytes.Buffer, error) {
-	c.buffer = _bpool.Get(content.Len() / 3)
-	c.fw.Reset(c.buffer)
-	if err := internal.WriteN(c.fw, content.Bytes(), content.Len()); err != nil {
-		return nil, err
+func (c *compressor) Compress(content []byte, buf *bytes.Buffer) error {
+	c.fw.Reset(buf)
+	if err := internal.WriteN(c.fw, content, len(content)); err != nil {
+		return err
 	}
 	if err := c.fw.Flush(); err != nil {
-		return nil, err
+		return err
 	}
-
-	if n := c.buffer.Len(); n >= 4 {
-		compressedContent := c.buffer.Bytes()
+	if n := buf.Len(); n >= 4 {
+		compressedContent := buf.Bytes()
 		if tail := compressedContent[n-4:]; binary.BigEndian.Uint32(tail) == math.MaxUint16 {
-			c.buffer.Truncate(n - 4)
+			buf.Truncate(n - 4)
 		}
 	}
-	return c.buffer, nil
+	return nil
 }
 
 func newDecompressor() *decompressor { return &decompressor{fr: flate.NewReader(nil)} }

compress_test.go

@@ -17,16 +17,14 @@ func TestFlate(t *testing.T) {
 			var dps = newDecompressor()
 			var n = internal.AlphabetNumeric.Intn(1024)
 			var rawText = internal.AlphabetNumeric.Generate(n)
-			var buf = bytes.NewBufferString("")
-			buf.Write(rawText)
-			compressedText, err := cps.Compress(buf)
-			if err != nil {
+			var compressedBuf = bytes.NewBufferString("")
+			if err := cps.Compress(rawText, compressedBuf); err != nil {
 				as.NoError(err)
 				return
 			}
 
-			buf.Reset()
-			buf.Write(compressedText.Bytes())
+			var buf = bytes.NewBufferString("")
+			buf.Write(compressedBuf.Bytes())
 			plainText, err := dps.Decompress(buf)
 			if err != nil {
 				as.NoError(err)
@@ -41,18 +39,16 @@ func TestFlate(t *testing.T) {
 		var dps = newDecompressor()
 		var n = internal.AlphabetNumeric.Intn(1024)
 		var rawText = internal.AlphabetNumeric.Generate(n)
-		var buf = bytes.NewBufferString("")
-		buf.Write(rawText)
-		compressedText, err := cps.Compress(buf)
-		if err != nil {
+		var compressedBuf = bytes.NewBufferString("")
+		if err := cps.Compress(rawText, compressedBuf); err != nil {
 			as.NoError(err)
 			return
 		}
 
-		buf.Reset()
-		buf.Write(compressedText.Bytes())
+		var buf = bytes.NewBufferString("")
+		buf.Write(compressedBuf.Bytes())
 		buf.WriteString("1234")
-		_, err = dps.Decompress(buf)
+		_, err := dps.Decompress(buf)
 		as.Error(err)
 	})
 }

conn.go

@@ -71,12 +71,19 @@ func serveWebSocket(isServer bool, config *Config, session SessionStorage, netCo
 	return c
 }
 
-// Listen listening to websocket messages through a dead loop
-// 监听websocket消息
+// Listen 监听websocket消息
+// Deprecated: Listen will be deprecated in future versions, please use ReadLoop instead.
 func (c *Conn) Listen() {
+	c.ReadLoop()
+}
+
+// ReadLoop start a read message loop
+// 启动一个读消息的死循环
+func (c *Conn) ReadLoop() {
 	defer c.conn.Close()
 
 	c.handler.OnOpen(c)
+
 	for {
 		if err := c.readMessage(); err != nil {
 			c.emitError(err)

examples/autobahn/client/main.go

@@ -32,7 +32,7 @@ func testCase(id int) {
 		log.Println(err.Error())
 		return
 	}
-	go socket.Listen()
+	go socket.ReadLoop()
 	<-handler.onexit
 }
 

examples/autobahn/server/main.go

@@ -15,11 +15,11 @@ func main() {
 	})
 
 	http.HandleFunc("/connect", func(writer http.ResponseWriter, request *http.Request) {
-		socket, err := upgrader.Accept(writer, request)
+		socket, err := upgrader.Upgrade(writer, request)
 		if err != nil {
 			return
 		}
-		go socket.Listen()
+		socket.ReadLoop()
 	})
 
 	_ = http.ListenAndServe(":3000", nil)

examples/client/client.go

@@ -15,7 +15,7 @@ func main() {
 		log.Printf(err.Error())
 		return
 	}
-	go socket.Listen()
+	go socket.ReadLoop()
 
 	for {
 		var text = ""

examples/echo/main.go

@@ -0,0 +1,27 @@
+package main
+
+import (
+	"github.com/lxzan/gws"
+	"log"
+)
+
+func main() {
+	var app = gws.NewServer(new(Handler), &gws.ServerOption{
+		CompressEnabled:  true,
+		CheckUtf8Enabled: true,
+	})
+	log.Fatalf("%v", app.Run(":3000"))
+}
+
+type Handler struct {
+	gws.BuiltinEventHandler
+}
+
+func (c *Handler) OnPing(socket *gws.Conn, payload []byte) {
+	socket.WritePong(payload)
+}
+
+func (c *Handler) OnMessage(socket *gws.Conn, message *gws.Message) {
+	defer message.Close()
+	_ = socket.WriteMessage(message.Opcode, message.Bytes())
+}

protocol.go

@@ -7,6 +7,8 @@ import (
 	"io"
 )
 
+const frameHeaderSize = 14
+
 type Opcode uint8
 
 const (
@@ -62,7 +64,7 @@ func (b BuiltinEventHandler) OnPong(socket *Conn, payload []byte) {}
 
 func (b BuiltinEventHandler) OnMessage(socket *Conn, message *Message) {}
 
-type frameHeader [internal.FrameHeaderSize]byte
+type frameHeader [frameHeaderSize]byte
 
 func (c *frameHeader) GetFIN() bool {
 	return ((*c)[0] >> 7) == 1

updrader.go

@@ -62,13 +62,19 @@ func (c *Upgrader) connectHandshake(r *http.Request, responseHeader http.Header,
 }
 
 // Accept http upgrade to websocket protocol
+// Deprecated: Accept will be deprecated in future versions, please use Upgrade instead.
 func (c *Upgrader) Accept(w http.ResponseWriter, r *http.Request) (*Conn, error) {
+	return c.Upgrade(w, r)
+}
+
+// Upgrade http upgrade to websocket protocol
+func (c *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request) (*Conn, error) {
 	netConn, br, err := c.hijack(w)
 	if err != nil {
 		return nil, err
 	}
 
-	socket, err := c.doAccept(r, netConn, br)
+	socket, err := c.doUpgrade(r, netConn, br)
 	if err != nil {
 		_ = netConn.Close()
 		return nil, err
@@ -93,7 +99,7 @@ func (c *Upgrader) hijack(w http.ResponseWriter) (net.Conn, *bufio.Reader, error
 	return netConn, brw.Reader, nil
 }
 
-func (c *Upgrader) doAccept(r *http.Request, netConn net.Conn, br *bufio.Reader) (*Conn, error) {
+func (c *Upgrader) doUpgrade(r *http.Request, netConn net.Conn, br *bufio.Reader) (*Conn, error) {
 	var session = new(sliceMap)
 	var header = c.option.ResponseHeader.Clone()
 	if !c.option.CheckOrigin(r, session) {
@@ -138,6 +144,10 @@ func (c *Upgrader) doAccept(r *http.Request, netConn net.Conn, br *bufio.Reader)
 type Server struct {
 	upgrader *Upgrader
 
+	// OnConnect 建立连接事件, 用于处理限流, 熔断和安全问题; 返回错误将会断开连接.
+	// Creates connection events for current limit, fuse and security issues; returning an error will disconnect.
+	OnConnect func(conn net.Conn) error
+
 	// OnError 接收握手过程中产生的错误回调
 	// Receive error callbacks generated during the handshake
 	OnError func(conn net.Conn, err error)
@@ -147,6 +157,7 @@ type Server struct {
 // create a websocket server
 func NewServer(eventHandler Event, option *ServerOption) *Server {
 	var c = &Server{upgrader: NewUpgrader(eventHandler, option)}
+	c.OnConnect = func(conn net.Conn) error { return nil }
 	c.OnError = func(conn net.Conn, err error) {}
 	return c
 }
@@ -236,6 +247,12 @@ func (c *Server) serve(listener net.Listener) error {
 		}
 
 		go func() {
+			if err := c.OnConnect(conn); err != nil {
+				_ = conn.Close()
+				c.OnError(conn, err)
+				return
+			}
+
 			br := bufio.NewReaderSize(conn, c.upgrader.option.ReadBufferSize)
 			r, err := c.parseRequest(conn, br)
 			if err != nil {
@@ -244,13 +261,13 @@ func (c *Server) serve(listener net.Listener) error {
 				return
 			}
 
-			socket, err := c.upgrader.doAccept(r, conn, br)
+			socket, err := c.upgrader.doUpgrade(r, conn, br)
 			if err != nil {
 				_ = conn.Close()
 				c.OnError(conn, err)
 				return
 			}
-			socket.Listen()
+			socket.ReadLoop()
 		}()
 	}
 }