diff --git a/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware.go b/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware.go index 2cabb894c1..3c8e09399f 100644 --- a/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware.go +++ b/staging/src/k8s.io/apimachinery/pkg/util/proxy/upgradeaware.go @@ -273,7 +273,7 @@ func (h *UpgradeAwareHandler) tryUpgrade(w http.ResponseWriter, req *http.Reques defer backendConn.Close() // determine the http response code from the backend by reading from rawResponse+backendConn - rawResponseCode, headerBytes, err := getResponseCode(io.MultiReader(bytes.NewReader(rawResponse), backendConn)) + backendHTTPResponse, headerBytes, err := getResponse(io.MultiReader(bytes.NewReader(rawResponse), backendConn)) if err != nil { klog.V(6).Infof("Proxy connection error: %v", err) h.Responder.Error(w, req, err) @@ -300,6 +300,22 @@ func (h *UpgradeAwareHandler) tryUpgrade(w http.ResponseWriter, req *http.Reques } defer requestHijackedConn.Close() + if backendHTTPResponse.StatusCode != http.StatusSwitchingProtocols { + // If the backend did not upgrade the request, echo the response from the backend to the client and return, closing the connection. + klog.V(6).Infof("Proxy upgrade error, status code %d", backendHTTPResponse.StatusCode) + // set read/write deadlines + deadline := time.Now().Add(10 * time.Second) + backendConn.SetReadDeadline(deadline) + requestHijackedConn.SetWriteDeadline(deadline) + // write the response to the client + err := backendHTTPResponse.Write(requestHijackedConn) + if err != nil && !strings.Contains(err.Error(), "use of closed network connection") { + klog.Errorf("Error proxying data from backend to client: %v", err) + } + // Indicate we handled the request + return true + } + // Forward raw response bytes back to client. if len(rawResponse) > 0 { klog.V(6).Infof("Writing %d bytes to hijacked connection", len(rawResponse)) @@ -308,17 +324,6 @@ func (h *UpgradeAwareHandler) tryUpgrade(w http.ResponseWriter, req *http.Reques } } - if rawResponseCode != http.StatusSwitchingProtocols { - // If the backend did not upgrade the request, finish echoing the response from the backend to the client and return, closing the connection. - klog.V(6).Infof("Proxy upgrade error, status code %d", rawResponseCode) - _, err := io.Copy(requestHijackedConn, backendConn) - if err != nil && !strings.Contains(err.Error(), "use of closed network connection") { - klog.Errorf("Error proxying data from backend to client: %v", err) - } - // Indicate we handled the request - return true - } - // Proxy the connection. This is bidirectional, so we need a goroutine // to copy in each direction. Once one side of the connection exits, we // exit the function which performs cleanup and in the process closes @@ -380,17 +385,17 @@ func (h *UpgradeAwareHandler) DialForUpgrade(req *http.Request) (net.Conn, error return dial(updatedReq, h.UpgradeTransport) } -// getResponseCode reads a http response from the given reader, returns the status code, +// getResponseCode reads a http response from the given reader, returns the response, // the bytes read from the reader, and any error encountered -func getResponseCode(r io.Reader) (int, []byte, error) { +func getResponse(r io.Reader) (*http.Response, []byte, error) { rawResponse := bytes.NewBuffer(make([]byte, 0, 256)) // Save the bytes read while reading the response headers into the rawResponse buffer resp, err := http.ReadResponse(bufio.NewReader(io.TeeReader(r, rawResponse)), nil) if err != nil { - return 0, nil, err + return nil, nil, err } - // return the http status code and the raw bytes consumed from the reader in the process - return resp.StatusCode, rawResponse.Bytes(), nil + // return the http response and the raw bytes consumed from the reader in the process + return resp, rawResponse.Bytes(), nil } // dial dials the backend at req.URL and writes req to it.