Fix header timeouts on operation wait

cmd/incusd/operations.go

@@ -947,19 +947,37 @@ func operationWaitGet(d *Daemon, r *http.Request) response.Response {
 			ctx, cancel = context.WithCancel(r.Context())
 		}
 
-		defer cancel()
+		waitResponse := func(w http.ResponseWriter) error {
+			defer cancel()
+
+			// Write header to avoid client side timeouts.
+			w.Header().Set("Connection", "keep-alive")
+			w.Header().Set("Content-Type", "application/json")
+			w.Header().Set("X-Content-Type-Options", "nosniff")
+			w.WriteHeader(http.StatusOK)
+			f, ok := w.(http.Flusher)
+			if ok {
+				f.Flush()
+			}
 
-		err = op.Wait(ctx)
-		if err != nil {
-			return response.SmartError(err)
-		}
+			// Wait for the operation.
+			err = op.Wait(ctx)
+			if err != nil {
+				_ = response.SmartError(err).Render(w)
+				return nil
+			}
 
-		_, body, err := op.Render()
-		if err != nil {
-			return response.SmartError(err)
+			_, body, err := op.Render()
+			if err != nil {
+				_ = response.SmartError(err).Render(w)
+				return nil
+			}
+
+			_ = response.SyncResponse(true, body).Render(w)
+			return nil
 		}
 
-		return response.SyncResponse(true, body)
+		return response.ManualResponse(waitResponse)
 	}
 
 	// Then check if the query is from an operation on another node, and, if so, forward it

internal/server/response/response.go

@@ -187,7 +187,9 @@ func (r *syncResponse) Render(w http.ResponseWriter) error {
 		code = http.StatusOK
 	}
 
-	w.WriteHeader(code)
+	if w.Header().Get("Connection") != "keep-alive" {
+		w.WriteHeader(code)
+	}
 
 	// Handle plain text responses.
 	if r.plaintext {
@@ -347,7 +349,9 @@ func (r *errorResponse) Render(w http.ResponseWriter) error {
 	w.Header().Set("Content-Type", "application/json")
 	w.Header().Set("X-Content-Type-Options", "nosniff")
 
-	w.WriteHeader(r.code) // Set the error code in the HTTP header response.
+	if w.Header().Get("Connection") != "keep-alive" {
+		w.WriteHeader(r.code) // Set the error code in the HTTP header response.
+	}
 
 	_, err = fmt.Fprintln(w, buf.String())
 
@@ -524,7 +528,10 @@ func (r *forwardedResponse) Render(w http.ResponseWriter) error {
 		w.Header().Set(key, response.Header.Get(key))
 	}
 
-	w.WriteHeader(response.StatusCode)
+	if w.Header().Get("Connection") != "keep-alive" {
+		w.WriteHeader(response.StatusCode)
+	}
+
 	_, err = io.Copy(w, response.Body)
 	return err
 }