Improve websocket closing code

This commit is contained in:
Melon 2023-09-10 21:49:57 +01:00
parent ef5a15f5c3
commit e87809849e
Signed by: melon
GPG Key ID: 6C9D970C50D26A25

View File

@ -40,12 +40,13 @@ func (s *Server) Upgrade(rw http.ResponseWriter, req *http.Request) {
if err != nil {
return
}
defer c.Close()
s.connLock.Lock()
// no more connections allowed
if s.connStop {
s.connLock.Unlock()
_ = c.Close()
return
}
@ -62,6 +63,8 @@ func (s *Server) Upgrade(rw http.ResponseWriter, req *http.Request) {
s.Remove(c)
return
}
defer ic.Close()
d1 := make(chan struct{}, 1)
d2 := make(chan struct{}, 1)
@ -70,16 +73,13 @@ func (s *Server) Upgrade(rw http.ResponseWriter, req *http.Request) {
go s.wsRelay(d2, ic, c)
// wait for done signal and close both connections
go func() {
log.Println("[Websocket] Completed websocket hijacking")
// waiting until d1 or d2 close then automatically defer close both connections
select {
case <-d1:
case <-d2:
}
_ = c.Close()
_ = ic.Close()
}()
log.Println("[Websocket] Completed websocket hijacking")
}
func (s *Server) wsRelay(done chan struct{}, a, b *websocket.Conn) {