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 }