Skip to content

Commit 818579b

Browse files
committed
TestDialViaProxy: Fix bug in forward proxy
Closes #395 Confirmed library works correctly with a working forward proxy.
1 parent 249edb2 commit 818579b

File tree

2 files changed

+22
-13
lines changed

2 files changed

+22
-13
lines changed

Diff for: conn_test.go

+1
Original file line numberDiff line numberDiff line change
@@ -535,6 +535,7 @@ func assertEcho(tb testing.TB, ctx context.Context, c *websocket.Conn) {
535535
})
536536

537537
var act interface{}
538+
c.SetReadLimit(1 << 30)
538539
err := wsjson.Read(ctx, c, &act)
539540
assert.Success(tb, err)
540541
assert.Equal(tb, "read msg", exp, act)

Diff for: dial_test.go

+21-13
Original file line numberDiff line numberDiff line change
@@ -361,21 +361,29 @@ func (fc *forwardProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
361361
}
362362
w.Header().Set("PROXIED", "true")
363363
w.WriteHeader(resp.StatusCode)
364-
errc1 := xsync.Go(func() error {
365-
_, err := io.Copy(w, resp.Body)
366-
return err
367-
})
368-
var errc2 <-chan error
369-
if bodyw, ok := resp.Body.(io.Writer); ok {
370-
errc2 = xsync.Go(func() error {
371-
_, err := io.Copy(bodyw, r.Body)
364+
if resprw, ok := resp.Body.(io.ReadWriter); ok {
365+
c, brw, err := w.(http.Hijacker).Hijack()
366+
if err != nil {
367+
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
368+
return
369+
}
370+
brw.Flush()
371+
372+
errc1 := xsync.Go(func() error {
373+
_, err := io.Copy(c, resprw)
372374
return err
373375
})
374-
}
375-
select {
376-
case <-errc1:
377-
case <-errc2:
378-
case <-r.Context().Done():
376+
errc2 := xsync.Go(func() error {
377+
_, err := io.Copy(resprw, c)
378+
return err
379+
})
380+
select {
381+
case <-errc1:
382+
case <-errc2:
383+
case <-r.Context().Done():
384+
}
385+
} else {
386+
io.Copy(w, resp.Body)
379387
}
380388
}
381389

0 commit comments

Comments
 (0)