Merge pull request #71419 from liggitt/backend-error-deadline

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.