nnet/ws.go

package nnet

import (
	"git.noahlan.cn/noahlan/nnet/config"
	"git.noahlan.cn/noahlan/nnet/conn"
	"github.com/gorilla/websocket"
)

type WsEventOption func(conf config.WSEvent)

func WithPingHandler(fn func(appData string)) WsEventOption {
	return func(conf config.WSEvent) {
		conf.PingHandler = fn
	}
}

func WithPongHandler(fn func(appData string)) WsEventOption {
	return func(conf config.WSEvent) {
		conf.PongHandler = fn
	}
}

func WithCloseHandler(fn func(closeCode int, closeText string) error) WsEventOption {
	return func(conf config.WSEvent) {
		conf.CloseHandler = fn
	}
}

func (ngin *Engine) handleWS(wc *websocket.Conn, conf config.WSEvent) *conn.Connection {
	wsConn := conn.NewWSConn(wc)

	nc := ngin.handle(wsConn)

	defaultCloseHandler := wsConn.CloseHandler()
	wsConn.SetCloseHandler(func(code int, text string) error {
		result := defaultCloseHandler(code, text)
		_ = wsConn.Close()
		ngin.evtMgr.OnClose(nc)
		return result
	})

	// ping
	defaultPingHandler := wsConn.PingHandler()
	wsConn.SetPingHandler(func(appData string) error {
		if conf.PingHandler != nil {
			conf.PingHandler(appData)
		}
		return defaultPingHandler(appData)
	})

	// pong
	defaultPongHandler := wsConn.PongHandler()
	wsConn.SetPongHandler(func(appData string) error {
		if conf.PongHandler != nil {
			conf.PongHandler(appData)
		}
		return defaultPongHandler(appData)
	})

	return nc
}