1f349/violet
4
0
Fork 0
mirror of https://github.com/1f349/violet.git synced 2024-11-09 22:22:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Improve websocket closing code

Browse Source
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
1 changed files with 10 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
proxy/websocket/server.go
View File

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