diff --git a/utils/net/websocket.go b/utils/net/websocket.go index 8ecc17e1..436346a7 100644 --- a/utils/net/websocket.go +++ b/utils/net/websocket.go @@ -11,16 +11,45 @@ import ( "github.com/fatedier/frp/utils/log" "golang.org/x/net/websocket" + "sync" ) var ( ErrWebsocketListenerClosed = errors.New("websocket listener closed") + websoketClientMap = new(sync.Map) + ticker = time.NewTicker(55 * time.Second) + pingCodec = websocket.Codec{func(v interface{}) ([]byte, byte, error) { + return []byte{}, websocket.PingFrame, nil + }, nil} ) const ( FrpWebsocketPath = "/~!frp" ) +func init() { + go pingClient() +} + +func pingClient() { + for { + select { + case <-ticker.C: + willRemove := make([]interface{}, 0) + websoketClientMap.Range(func(k, v interface{}) bool { + conn := v.(*websocket.Conn) + if err := pingCodec.Send(conn, nil); err != nil { + willRemove = append(willRemove, k) + } + return true + }) + for _, value := range willRemove { + websoketClientMap.Delete(value) + } + } + } +} + type WebsocketListener struct { net.Addr ln net.Listener @@ -43,8 +72,10 @@ func NewWebsocketListener(ln net.Listener) (wl *WebsocketListener) { muxer.Handle(FrpWebsocketPath, websocket.Handler(func(c *websocket.Conn) { notifyCh := make(chan struct{}) conn := WrapCloseNotifyConn(c, func() { + websoketClientMap.Delete(c.LocalAddr().String()) close(notifyCh) }) + websoketClientMap.Store(c.LocalAddr().String(), c) wl.accept <- conn <-notifyCh })) @@ -53,7 +84,6 @@ func NewWebsocketListener(ln net.Listener) (wl *WebsocketListener) { Addr: ln.Addr().String(), Handler: muxer, } - go wl.server.Serve(ln) return }