update README, allowed custom handle ping event

README.md

@@ -6,4 +6,101 @@
 go get github.com/RealFax/q-websocket@latest
 ```
 
+## quick start
+
+```go
+// echo server
+
+package main
+
+import (
+	qWebsocket "github.com/RealFax/q-websocket"
+
+	"github.com/gobwas/ws/wsutil"
+	"github.com/panjf2000/gnet/v2"
+
+	"log"
+)
+
+func echo(req *qWebsocket.HandlerParams) {
+	wsutil.WriteServerText(req.Writer, req.Request)
+}
+
+func main() {
+	server := qWebsocket.NewServer(
+		"tcp://0.0.0.0:9090",
+		qWebsocket.WithHandler(echo),
+	)
+
+	if err := server.ListenAndServe(gnet.WithMulticore(true)); err != nil {
+		log.Fatal("server error:", err)
+	}
+}
+```
+
+## using proto
+_proto is a low-level handler built into q-websocket, providing simple router_
+
+```go
+// using proto 
+
+package main
+
+import (
+	qWebsocket "github.com/RealFax/q-websocket"
+	"github.com/RealFax/q-websocket/proto"
+
+	"github.com/gobwas/ws"
+	"github.com/panjf2000/gnet/v2"
+
+	"log"
+	"strings"
+)
+
+// define proto struct
+
+type Proto struct {
+	Type    uint32 `json:"type"`
+	Message string `json:"message"`
+}
+
+func (p *Proto) Key() uint32  { return p.Type }
+func (p *Proto) Value() Proto { return *p }
+func (p *Proto) Self() *Proto { return p }
+
+const (
+	Echo uint32 = iota
+	Pong
+)
+
+func main() {
+	engine := proto.New[Proto, uint32](func() proto.Proto[Proto, uint32] {
+		return new(Proto)
+	})
+
+	// echo handler
+	engine.Register(Echo, func(r *proto.Request[Proto]) {
+		r.WriteText([]byte(r.Request.Message))
+	})
+
+	// ping pong handler
+	engine.Register(Pong, func(r *proto.Request[Proto]) {
+		if strings.ToLower(r.Request.Message) == "ping" {
+			r.WriteText([]byte("pong"))
+			return
+		}
+		r.WriteClose(ws.StatusGoingAway, "")
+	})
+
+	server := qWebsocket.NewServer(
+		"tcp://127.0.0.1:9090",
+		qWebsocket.WithHandler(engine.UseHandler()),
+	)
+
+	if err := server.ListenAndServe(gnet.WithMulticore(true)); err != nil {
+		log.Fatal("server error:", err)
+	}
+}
+```
+
 more usage see: [Example](https://github.com/RealFax/q-websocket/tree/master/example)

example/simple-chat/chat.go

@@ -141,13 +141,13 @@ func main() {
 	server := qWebsocket.NewServer(
 		"tcp://127.0.0.1:8080",
 		qWebsocket.WithHandler(engine.UseHandler()),
-		qWebsocket.WithOnCloseHandler(func(conn *qWebsocket.GNetUpgraderConn, _ error) {
+		qWebsocket.WithOnCloseHandler(func(conn *qWebsocket.Conn, _ error) {
 			s.offline(conn.ID)
 		}),
 	)
 
 	go func() {
-		if err := server.ListenAndServer(gnet.WithMulticore(true)); err != nil {
+		if err := server.ListenAndServe(gnet.WithMulticore(true)); err != nil {
 			fmt.Println("[-] gnet error:", err)
 		}
 	}()

handler.go

@@ -2,19 +2,26 @@ package qWebsocket
 
 import (
 	"github.com/gobwas/ws"
+	"github.com/gobwas/ws/wsutil"
 	"io"
 )
 
-type OnCloseHandlerFunc func(conn *GNetUpgraderConn, err error)
+type (
+	OnCloseHandlerFunc func(conn *Conn, err error)
+	OnPingHandlerFunc  func(conn *Conn)
+	HandlerFunc        func(req *HandlerParams)
 
-type HandlerParams struct {
-	OpCode  ws.OpCode
-	Request []byte
-	Writer  io.Writer
-	WsConn  *GNetUpgraderConn
-}
+	HandlerParams struct {
+		OpCode  ws.OpCode
+		Request []byte
+		Writer  io.Writer
+		WsConn  *Conn
+	}
+)
 
-type HandlerFunc func(req *HandlerParams)
+func EmptyOnCloseHandler(_ *Conn, _ error) {}
+func EmptyHandler(_ *HandlerParams)        {}
 
-func EmptyOnCloseHandler(_ *GNetUpgraderConn, _ error) {}
-func EmptyHandler(_ *HandlerParams)                    {}
+func DefaultOnPingHandler(c *Conn) {
+	wsutil.WriteServerMessage(c, ws.OpPong, nil)
+}

options.go

@@ -47,6 +47,12 @@ func WithOnCloseHandler(handler OnCloseHandlerFunc) OptionFunc {
 	}
 }
 
+func WithOnPingHandler(handler OnPingHandlerFunc) OptionFunc {
+	return func(s *Server) {
+		s.onPingHandler = handler
+	}
+}
+
 func WithHandler(handler HandlerFunc) OptionFunc {
 	return func(s *Server) {
 		s.handler = handler

proto/proto_test.go

@@ -5,7 +5,6 @@ import (
 	"github.com/RealFax/q-websocket/proto"
 	"github.com/gobwas/ws"
 	"github.com/panjf2000/gnet/v2"
-	"log"
 	"testing"
 )
 
@@ -24,20 +23,20 @@ var engine = proto.New[Proto, uint32](func() proto.Proto[Proto, uint32] {
 
 func init() {
 	engine.Register(1, func(r *proto.Request[Proto]) {
-		log.Println(r.Request.Message)
+		r.WriteText([]byte(r.Request.Message))
 	})
 	engine.Register(2, func(r *proto.Request[Proto]) {
-		ws.WriteFrame(r.Writer, ws.NewTextFrame([]byte(r.Request.Message)))
+		r.WriteClose(ws.StatusGoingAway, "")
 	})
 }
 
 func TestEngine_Handler(t *testing.T) {
 	server := qWebsocket.NewServer(
-		"tcp://127.0.0.1:8080",
+		"tcp://127.0.0.1:10001",
 		qWebsocket.WithHandler(engine.UseHandler()),
 	)
 
-	if err := server.ListenAndServer(gnet.WithMulticore(true)); err != nil {
+	if err := server.ListenAndServe(gnet.WithMulticore(true)); err != nil {
 		t.Fatal(err)
 	}
 }

proto/request.go

@@ -4,13 +4,14 @@ import (
 	"bytes"
 	qWebsocket "github.com/RealFax/q-websocket"
 	"github.com/gobwas/ws"
+	"github.com/gobwas/ws/wsutil"
 	"io"
 )
 
 type Request[T any] struct {
 	OpCode     ws.OpCode
 	Writer     io.Writer
-	Conn       *qWebsocket.GNetUpgraderConn
+	Conn       *qWebsocket.Conn
 	Request    *T
 	RawRequest []byte
 }
@@ -20,13 +21,14 @@ func (t Request[T]) Reader() io.Reader {
 }
 
 func (t Request[T]) WriteText(p []byte) error {
-	return ws.WriteFrame(t.Writer, ws.NewTextFrame(p))
+	return wsutil.WriteServerMessage(t.Writer, ws.OpText, p)
 }
 
 func (t Request[T]) WriteBinary(p []byte) error {
-	return ws.WriteFrame(t.Writer, ws.NewBinaryFrame(p))
+	return wsutil.WriteServerMessage(t.Writer, ws.OpBinary, p)
 }
 
 func (t Request[T]) WriteClose(statusCode ws.StatusCode, reason string) error {
-	return ws.WriteFrame(t.Writer, ws.NewCloseFrame(ws.NewCloseFrameBody(statusCode, reason)))
+	defer t.Conn.Close()
+	return wsutil.WriteServerMessage(t.Writer, ws.OpClose, ws.NewCloseFrameBody(statusCode, reason))
 }

server.go

@@ -23,6 +23,7 @@ type Server struct {
 	upgrader       *ws.Upgrader
 	keepConnTable  *ttlcache.Cache[string, gnet.Conn]
 	onCloseHandler OnCloseHandlerFunc
+	onPingHandler  OnPingHandlerFunc
 	handler        HandlerFunc
 }
 
@@ -52,12 +53,16 @@ func (s *Server) withDefault() {
 		WithOnCloseHandler(EmptyOnCloseHandler)(s)
 	}
 
+	if s.onPingHandler == nil {
+		WithOnPingHandler(DefaultOnPingHandler)(s)
+	}
+
 	if s.handler == nil {
 		WithHandler(EmptyHandler)(s)
 	}
 }
 
-func (s *Server) closeWS(conn *GNetUpgraderConn, statusCode ws.StatusCode, reason error) error {
+func (s *Server) closeWS(conn *Conn, statusCode ws.StatusCode, reason error) error {
 	defer s.onCloseHandler(conn, reason)
 	return ws.WriteFrame(conn, ws.NewCloseFrame(ws.NewCloseFrameBody(statusCode, func() string {
 		if reason != nil {
@@ -75,7 +80,7 @@ func (s *Server) setupTimeoutHandler() {
 		item *ttlcache.Item[string, gnet.Conn],
 	) {
 		atomic.AddInt64(&s.connNum, -1)
-		upgraderConn, ok := item.Value().Context().(*GNetUpgraderConn)
+		upgraderConn, ok := item.Value().Context().(*Conn)
 		if !ok {
 			item.Value().Close()
 			return
@@ -137,7 +142,7 @@ func (s *Server) OnTraffic(c gnet.Conn) gnet.Action {
 		c.SetContext(NewUpgraderConn(c))
 	}
 
-	upgraderConn, ok := c.Context().(*GNetUpgraderConn)
+	upgraderConn, ok := c.Context().(*Conn)
 	if !ok {
 		log.Printf("[-] invalid context, remote addr: %s", c.RemoteAddr())
 		return gnet.None
@@ -167,7 +172,10 @@ func (s *Server) OnTraffic(c gnet.Conn) gnet.Action {
 	for _, message := range messages {
 		switch message.OpCode {
 		case ws.OpPing:
-			wsutil.WriteServerMessage(upgraderConn, ws.OpPong, nil)
+			// async handle
+			s.workerPool.Submit(func() {
+				s.onPingHandler(upgraderConn)
+			})
 			upgraderConn.UpdateActive()
 		case ws.OpText, ws.OpBinary:
 			// async handle
@@ -189,7 +197,7 @@ func (s *Server) OnTraffic(c gnet.Conn) gnet.Action {
 	return gnet.None
 }
 
-func (s *Server) ListenAndServer(opts ...gnet.Option) error {
+func (s *Server) ListenAndServe(opts ...gnet.Option) error {
 	s.setupTimeoutHandler()
 	return gnet.Run(s, s.addr, opts...)
 }

server_test.go

@@ -42,7 +42,7 @@ func TestServer_ListenAndServer(t *testing.T) {
 		http.ListenAndServe("localhost:9090", nil)
 	}()
 
-	if err := server.ListenAndServer(gnet.WithMulticore(true)); err != nil {
+	if err := server.ListenAndServe(gnet.WithMulticore(true)); err != nil {
 		t.Fatal(err)
 	}
 }

upgrader.go

@@ -13,7 +13,8 @@ var (
 	emptyUpgrader = &ws.Upgrader{}
 )
 
-type GNetUpgraderConn struct {
+// Conn is upgraded websocket conn
+type Conn struct {
 	gnet.Conn
 	LastActive      int64 // atomic
 	successUpgraded bool
@@ -22,23 +23,24 @@ type GNetUpgraderConn struct {
 	ctx context.Context
 }
 
-func (c *GNetUpgraderConn) Context() context.Context {
+func (c *Conn) Context() context.Context {
 	return c.ctx
 }
 
-func (c *GNetUpgraderConn) SetContext(ctx context.Context) {
+func (c *Conn) SetContext(ctx context.Context) {
 	c.ctx = ctx
 }
 
-func (c *GNetUpgraderConn) UpdateActive() {
+func (c *Conn) UpdateActive() {
 	atomic.StoreInt64(&c.LastActive, time.Now().Unix())
 }
 
-func NewUpgraderConn(conn gnet.Conn) *GNetUpgraderConn {
-	return &GNetUpgraderConn{
+func NewUpgraderConn(conn gnet.Conn) *Conn {
+	return &Conn{
 		Conn:            conn,
 		LastActive:      time.Now().Unix(),
 		successUpgraded: false,
 		ID:              uuid.New().String(),
+		ctx:             context.Background(),
 	}
 }