|
1 | 1 | package client
|
2 | 2 |
|
3 | 3 | import (
|
| 4 | + "bufio" |
4 | 5 | "crypto/tls"
|
5 | 6 | "errors"
|
6 | 7 | "fmt"
|
@@ -74,9 +75,47 @@ func (cli *Client) postHijacked(ctx context.Context, path string, query url.Valu
|
74 | 75 | // Server hijacks the connection, error 'connection closed' expected
|
75 | 76 | _, err = clientconn.Do(req)
|
76 | 77 |
|
77 |
| - rwc, br := clientconn.Hijack() |
| 78 | + c, br := clientconn.Hijack() |
| 79 | + if br.Buffered() > 0 { |
| 80 | + // If there is buffered content, wrap the connection. We return an |
| 81 | + // object that implements CloseWrite iff the underlying connection |
| 82 | + // implements it. |
| 83 | + if _, ok := c.(types.CloseWriter); ok { |
| 84 | + c = &hijackedConnCloseWriter{c, br} |
| 85 | + } else { |
| 86 | + c = &hijackedConn{c, br} |
| 87 | + } |
| 88 | + } else { |
| 89 | + br.Reset(nil) |
| 90 | + } |
| 91 | + |
| 92 | + return types.HijackedResponse{Conn: c, Reader: bufio.NewReader(c)}, err |
| 93 | +} |
| 94 | + |
| 95 | +// hijackedConn wraps a net.Conn and is returned by setupHijackConn in the case |
| 96 | +// that a) there was already buffered data in the http layer when Hijack() was |
| 97 | +// called, and b) the underlying net.Conn does *not* implement CloseWrite(). |
| 98 | +// hijackedConn does not implement CloseWrite() either. |
| 99 | +type hijackedConn struct { |
| 100 | + net.Conn |
| 101 | + r *bufio.Reader |
| 102 | +} |
| 103 | + |
| 104 | +func (c *hijackedConn) Read(b []byte) (int, error) { |
| 105 | + return c.r.Read(b) |
| 106 | +} |
| 107 | + |
| 108 | +// hijackedConnCloseWriter is a hijackedConn which additionally implements |
| 109 | +// CloseWrite(). It is returned by setupHijackConn in the case that a) there |
| 110 | +// was already buffered data in the http layer when Hijack() was called, and b) |
| 111 | +// the underlying net.Conn *does* implement CloseWrite(). |
| 112 | +type hijackedConnCloseWriter hijackedConn |
| 113 | + |
| 114 | +var _ types.CloseWriter = &hijackedConnCloseWriter{} |
78 | 115 |
|
79 |
| - return types.HijackedResponse{Conn: rwc, Reader: br}, err |
| 116 | +func (c *hijackedConnCloseWriter) CloseWrite() error { |
| 117 | + conn := c.Conn.(types.CloseWriter) |
| 118 | + return conn.CloseWrite() |
80 | 119 | }
|
81 | 120 |
|
82 | 121 | func tlsDial(network, addr string, config *tls.Config) (net.Conn, error) {
|
|
0 commit comments