send pingFrame from frps to frpc

frps使用websocket协议运行在nginx后面的时候,如果没有数据传输会被nginx判定读超时,断开websocket链接。
This commit is contained in:
FeelGo 2018-09-09 22:20:41 +08:00 committed by GitHub
parent 66a69f873f
commit 50d19c71eb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -11,16 +11,45 @@ import (
"github.com/fatedier/frp/utils/log" "github.com/fatedier/frp/utils/log"
"golang.org/x/net/websocket" "golang.org/x/net/websocket"
"sync"
) )
var ( var (
ErrWebsocketListenerClosed = errors.New("websocket listener closed") 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 ( const (
FrpWebsocketPath = "/~!frp" 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 { type WebsocketListener struct {
net.Addr net.Addr
ln net.Listener ln net.Listener
@ -43,8 +72,10 @@ func NewWebsocketListener(ln net.Listener) (wl *WebsocketListener) {
muxer.Handle(FrpWebsocketPath, websocket.Handler(func(c *websocket.Conn) { muxer.Handle(FrpWebsocketPath, websocket.Handler(func(c *websocket.Conn) {
notifyCh := make(chan struct{}) notifyCh := make(chan struct{})
conn := WrapCloseNotifyConn(c, func() { conn := WrapCloseNotifyConn(c, func() {
websoketClientMap.Delete(c.LocalAddr().String())
close(notifyCh) close(notifyCh)
}) })
websoketClientMap.Store(c.LocalAddr().String(), c)
wl.accept <- conn wl.accept <- conn
<-notifyCh <-notifyCh
})) }))
@ -53,7 +84,6 @@ func NewWebsocketListener(ln net.Listener) (wl *WebsocketListener) {
Addr: ln.Addr().String(), Addr: ln.Addr().String(),
Handler: muxer, Handler: muxer,
} }
go wl.server.Serve(ln) go wl.server.Serve(ln)
return return
} }