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

Some changes after debugging websockets

Browse Source
This commit is contained in:
Melon 2023-08-17 15:23:23 +01:00
parent fbbcf0440c
commit b5ff809345
Signed by: melon
GPG Key ID: 6C9D970C50D26A25
2 changed files with 9 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
cmd/violet/serve.go
View File

@ -170,7 +170,7 @@ func normalLoad(startUp startUpConfig, wd string) {
_ = srvHttp.Close() _ = srvHttp.Close()
} }
if srvHttps != nil { if srvHttps != nil {
_ = srvHttps.Shutdown(context.Background()) _ = srvHttps.Close()
} }
}) })
} }

11
proxy/websocket/server.go
View File

@ -33,6 +33,7 @@ func NewServer() *Server {
} }
func (s *Server) Upgrade(rw http.ResponseWriter, req *http.Request) { func (s *Server) Upgrade(rw http.ResponseWriter, req *http.Request) {
req.URL.Scheme = "ws"
log.Printf("[Websocket] Upgrading request to '%s' from '%s'\n", req.URL.String(), req.Header.Get("Origin")) log.Printf("[Websocket] Upgrading request to '%s' from '%s'\n", req.URL.String(), req.Header.Get("Origin"))
c, err := upgrader.Upgrade(rw, req, nil) c, err := upgrader.Upgrade(rw, req, nil)
@ -55,16 +56,17 @@ func (s *Server) Upgrade(rw http.ResponseWriter, req *http.Request) {
log.Printf("[Websocket] Dialing: '%s'\n", req.URL.String()) log.Printf("[Websocket] Dialing: '%s'\n", req.URL.String())
// dial for internal connection // dial for internal connection
ic, _, err := websocket.DefaultDialer.DialContext(req.Context(), req.URL.String(), req.Header) ic, _, err := websocket.DefaultDialer.DialContext(req.Context(), req.URL.String(), nil)
if err != nil { if err != nil {
log.Printf("[Websocket] Failed to dial '%s': %s\n", req.URL.String(), err)
s.Remove(c) s.Remove(c)
return return
} }
done := make(chan struct{}, 1) done := make(chan struct{}, 1)
// relay messages each way // relay messages each way
s.wsRelay(done, c, ic) go s.wsRelay(done, c, ic)
s.wsRelay(done, ic, c) go s.wsRelay(done, ic, c)
// wait for done signal and close both connections // wait for done signal and close both connections
go func() { go func() {
@ -72,6 +74,8 @@ func (s *Server) Upgrade(rw http.ResponseWriter, req *http.Request) {
_ = c.Close() _ = c.Close()
_ = ic.Close() _ = ic.Close()
}() }()
log.Println("[Websocket] Completed websocket hijacking")
} }
func (s *Server) wsRelay(done chan struct{}, a, b *websocket.Conn) { func (s *Server) wsRelay(done chan struct{}, a, b *websocket.Conn) {
@ -81,6 +85,7 @@ func (s *Server) wsRelay(done chan struct{}, a, b *websocket.Conn) {
for { for {
mt, message, err := a.ReadMessage() mt, message, err := a.ReadMessage()
if err != nil { if err != nil {
log.Println("Websocket read message error: ", err)
return return
} }
if b.WriteMessage(mt, message) != nil { if b.WriteMessage(mt, message) != nil {