From 67dee19f54f118a22d1dd71f1d2761c776a6e1ee Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?St=C3=A9phane=20Graber?= <stgraber@stgraber.org>
Date: Sat, 31 Aug 2024 00:23:16 -0400
Subject: [PATCH] client: Fix error handling in push mode copy
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Stéphane Graber <stgraber@stgraber.org>
Sponsored-by: https://webdock.io
---
 client/incus_instances.go | 36 ++++++++++++++++++++++++++++++++----
 1 file changed, 32 insertions(+), 4 deletions(-)

diff --git a/client/incus_instances.go b/client/incus_instances.go
index 54af0dae74e..7af66fbe6fe 100644
--- a/client/incus_instances.go
+++ b/client/incus_instances.go
@@ -888,7 +888,7 @@ func (r *ProtocolIncus) CopyInstance(source InstanceServer, instance api.Instanc
 		target.Certificate = info.Certificate
 		sourceReq.Target = &target
 
-		return r.tryMigrateInstance(source, instance.Name, sourceReq, info.Addresses)
+		return r.tryMigrateInstance(source, instance.Name, sourceReq, info.Addresses, op)
 	}
 
 	// Get source server connection information
@@ -999,7 +999,7 @@ func (r *ProtocolIncus) RenameInstance(name string, instance api.InstancePost) (
 
 // tryMigrateInstance attempts to migrate a specific instance from a source server to one of the target URLs.
 // The function runs the migration operation asynchronously and returns a RemoteOperation to track the progress and handle any errors.
-func (r *ProtocolIncus) tryMigrateInstance(source InstanceServer, name string, req api.InstancePost, urls []string) (RemoteOperation, error) {
+func (r *ProtocolIncus) tryMigrateInstance(source InstanceServer, name string, req api.InstancePost, urls []string, op Operation) (RemoteOperation, error) {
 	if len(urls) == 0 {
 		return nil, fmt.Errorf("The target server isn't listening on the network")
 	}
@@ -1011,6 +1011,9 @@ func (r *ProtocolIncus) tryMigrateInstance(source InstanceServer, name string, r
 	operation := req.Target.Operation
 
 	// Forward targetOp to remote op
+	chConnect := make(chan error, 1)
+	chWait := make(chan error, 1)
+
 	go func() {
 		success := false
 		var errors []remoteOperationResult
@@ -1044,10 +1047,35 @@ func (r *ProtocolIncus) tryMigrateInstance(source InstanceServer, name string, r
 			break
 		}
 
-		if !success {
-			rop.err = remoteOperationError("Failed instance migration", errors)
+		if success {
+			chConnect <- nil
+			close(chConnect)
+		} else {
+			chConnect <- remoteOperationError("Failed instance migration", errors)
+			close(chConnect)
+
+			if op != nil {
+				_ = op.Cancel()
+			}
 		}
+	}()
 
+	if op != nil {
+		go func() {
+			chWait <- op.Wait()
+			close(chWait)
+		}()
+	}
+
+	go func() {
+		var err error
+
+		select {
+		case err = <-chConnect:
+		case err = <-chWait:
+		}
+
+		rop.err = err
 		close(rop.chDone)
 	}()