Переглянути джерело

Merge pull request #36517 from jim-minter/missing_closewrite

ensure hijackedConn implements CloseWrite function
Sebastiaan van Stijn 7 роки тому
батько
коміт
518d028baa
1 змінених файлів з 25 додано та 2 видалено
  1. 25 2
      client/hijack.go

+ 25 - 2
client/hijack.go

@@ -188,8 +188,14 @@ func (cli *Client) setupHijackConn(req *http.Request, proto string) (net.Conn, e
 
 	c, br := clientconn.Hijack()
 	if br.Buffered() > 0 {
-		// If there is buffered content, wrap the connection
-		c = &hijackedConn{c, br}
+		// If there is buffered content, wrap the connection.  We return an
+		// object that implements CloseWrite iff the underlying connection
+		// implements it.
+		if _, ok := c.(types.CloseWriter); ok {
+			c = &hijackedConnCloseWriter{c, br}
+		} else {
+			c = &hijackedConn{c, br}
+		}
 	} else {
 		br.Reset(nil)
 	}
@@ -197,6 +203,10 @@ func (cli *Client) setupHijackConn(req *http.Request, proto string) (net.Conn, e
 	return c, nil
 }
 
+// hijackedConn wraps a net.Conn and is returned by setupHijackConn in the case
+// that a) there was already buffered data in the http layer when Hijack() was
+// called, and b) the underlying net.Conn does *not* implement CloseWrite().
+// hijackedConn does not implement CloseWrite() either.
 type hijackedConn struct {
 	net.Conn
 	r *bufio.Reader
@@ -205,3 +215,16 @@ type hijackedConn struct {
 func (c *hijackedConn) Read(b []byte) (int, error) {
 	return c.r.Read(b)
 }
+
+// hijackedConnCloseWriter is a hijackedConn which additionally implements
+// CloseWrite().  It is returned by setupHijackConn in the case that a) there
+// was already buffered data in the http layer when Hijack() was called, and b)
+// the underlying net.Conn *does* implement CloseWrite().
+type hijackedConnCloseWriter hijackedConn
+
+var _ types.CloseWriter = &hijackedConnCloseWriter{}
+
+func (c *hijackedConnCloseWriter) CloseWrite() error {
+	conn := c.Conn.(types.CloseWriter)
+	return conn.CloseWrite()
+}