Skip to content

Commit

Permalink
TestDialViaProxy: Fix bug in forward proxy
Browse files Browse the repository at this point in the history
Closes #395

Confirmed library works correctly with a working forward proxy.
  • Loading branch information
nhooyr committed Oct 19, 2023
1 parent 249edb2 commit 818579b
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 13 deletions.
1 change: 1 addition & 0 deletions conn_test.go
Expand Up @@ -535,6 +535,7 @@ func assertEcho(tb testing.TB, ctx context.Context, c *websocket.Conn) {
})

var act interface{}
c.SetReadLimit(1 << 30)
err := wsjson.Read(ctx, c, &act)
assert.Success(tb, err)
assert.Equal(tb, "read msg", exp, act)
Expand Down
34 changes: 21 additions & 13 deletions dial_test.go
Expand Up @@ -361,21 +361,29 @@ func (fc *forwardProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
}
w.Header().Set("PROXIED", "true")
w.WriteHeader(resp.StatusCode)
errc1 := xsync.Go(func() error {
_, err := io.Copy(w, resp.Body)
return err
})
var errc2 <-chan error
if bodyw, ok := resp.Body.(io.Writer); ok {
errc2 = xsync.Go(func() error {
_, err := io.Copy(bodyw, r.Body)
if resprw, ok := resp.Body.(io.ReadWriter); ok {
c, brw, err := w.(http.Hijacker).Hijack()
if err != nil {
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
return
}
brw.Flush()

errc1 := xsync.Go(func() error {
_, err := io.Copy(c, resprw)
return err
})
}
select {
case <-errc1:
case <-errc2:
case <-r.Context().Done():
errc2 := xsync.Go(func() error {
_, err := io.Copy(resprw, c)
return err
})
select {
case <-errc1:
case <-errc2:
case <-r.Context().Done():
}
} else {
io.Copy(w, resp.Body)
}
}

Expand Down

0 comments on commit 818579b

Please sign in to comment.