132 lines
2.9 KiB
Go
132 lines
2.9 KiB
Go
package app
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/coder/websocket"
|
|
"lol.mleku.dev/chk"
|
|
"lol.mleku.dev/log"
|
|
"utils.orly/units"
|
|
)
|
|
|
|
const (
|
|
DefaultWriteWait = 10 * time.Second
|
|
DefaultPongWait = 60 * time.Second
|
|
DefaultPingWait = DefaultPongWait / 2
|
|
DefaultMaxMessageSize = 1 * units.Mb
|
|
|
|
// CloseMessage denotes a close control message. The optional message
|
|
// payload contains a numeric code and text. Use the FormatCloseMessage
|
|
// function to format a close message payload.
|
|
CloseMessage = 8
|
|
|
|
// PingMessage denotes a ping control message. The optional message payload
|
|
// is UTF-8 encoded text.
|
|
PingMessage = 9
|
|
|
|
// PongMessage denotes a pong control message. The optional message payload
|
|
// is UTF-8 encoded text.
|
|
PongMessage = 10
|
|
)
|
|
|
|
func (s *Server) HandleWebsocket(w http.ResponseWriter, r *http.Request) {
|
|
remote := GetRemoteFromReq(r)
|
|
log.T.F("handling websocket connection from %s", remote)
|
|
if len(s.Config.IPWhitelist) > 0 {
|
|
for _, ip := range s.Config.IPWhitelist {
|
|
log.T.F("checking IP whitelist: %s", ip)
|
|
if strings.HasPrefix(remote, ip) {
|
|
log.T.F("IP whitelisted %s", remote)
|
|
goto whitelist
|
|
}
|
|
}
|
|
log.T.F("IP not whitelisted: %s", remote)
|
|
return
|
|
}
|
|
whitelist:
|
|
ctx, cancel := context.WithCancel(s.Ctx)
|
|
defer cancel()
|
|
var err error
|
|
var conn *websocket.Conn
|
|
if conn, err = websocket.Accept(
|
|
w, r, &websocket.AcceptOptions{OriginPatterns: []string{"*"}},
|
|
); chk.E(err) {
|
|
return
|
|
}
|
|
conn.SetReadLimit(DefaultMaxMessageSize)
|
|
defer conn.CloseNow()
|
|
listener := &Listener{
|
|
ctx: ctx,
|
|
Server: s,
|
|
conn: conn,
|
|
remote: remote,
|
|
}
|
|
ticker := time.NewTicker(DefaultPingWait)
|
|
go s.Pinger(ctx, conn, ticker, cancel)
|
|
defer func() {
|
|
log.D.F("closing websocket connection from %s", remote)
|
|
cancel()
|
|
ticker.Stop()
|
|
listener.publishers.Receive(&W{Cancel: true})
|
|
}()
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
default:
|
|
}
|
|
var typ websocket.MessageType
|
|
var msg []byte
|
|
log.T.F("waiting for message from %s", remote)
|
|
if typ, msg, err = conn.Read(ctx); chk.E(err) {
|
|
if strings.Contains(
|
|
err.Error(), "use of closed network connection",
|
|
) {
|
|
return
|
|
}
|
|
status := websocket.CloseStatus(err)
|
|
switch status {
|
|
case websocket.StatusNormalClosure,
|
|
websocket.StatusGoingAway,
|
|
websocket.StatusNoStatusRcvd,
|
|
websocket.StatusAbnormalClosure,
|
|
websocket.StatusProtocolError:
|
|
default:
|
|
log.E.F("unexpected close error from %s: %v", remote, err)
|
|
}
|
|
return
|
|
}
|
|
if typ == PingMessage {
|
|
if err = conn.Write(ctx, PongMessage, msg); chk.E(err) {
|
|
return
|
|
}
|
|
continue
|
|
}
|
|
go listener.HandleMessage(msg, remote)
|
|
}
|
|
}
|
|
|
|
func (s *Server) Pinger(
|
|
ctx context.Context, conn *websocket.Conn, ticker *time.Ticker,
|
|
cancel context.CancelFunc,
|
|
) {
|
|
defer func() {
|
|
cancel()
|
|
ticker.Stop()
|
|
}()
|
|
var err error
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
if err = conn.Ping(ctx); chk.E(err) {
|
|
return
|
|
}
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|