Fix progress bar for DFU/Recovery image upload

2 files changed, 4 insertions, 2 deletions

diff --git a/src/dfu.c b/src/dfu.c
index 4650e9a..6ba23e7 100644
--- a/src/dfu.c
+++ b/src/dfu.c
@@ -61,7 +61,6 @@ int dfu_client_new(struct idevicerestore_client_t* client)
 
 	irecv_event_subscribe(dfu, IRECV_PROGRESS, &dfu_progress_callback, NULL);
 	client->dfu->client = dfu;
-	register_progress('DFUP', "Uploading");
 	return 0;
 }
 
@@ -200,7 +199,9 @@ int dfu_send_component(struct idevicerestore_client_t* client, plist_t build_ide
 
 	logger(LL_INFO, "Sending %s (%zu bytes)...\n", component, size);
 
+	register_progress('DFUP', "Uploading");
 	irecv_error_t err = irecv_send_buffer(client->dfu->client, data, size, IRECV_SEND_OPT_DFU_NOTIFY_FINISH);
+	finalize_progress('DFUP');
 	if (err != IRECV_E_SUCCESS) {
 		logger(LL_ERROR, "Unable to send %s component: %s\n", component, irecv_strerror(err));
 		free(data);
diff --git a/src/recovery.c b/src/recovery.c
index 4f7520a..b843f1e 100644
--- a/src/recovery.c
+++ b/src/recovery.c
@@ -99,7 +99,6 @@ int recovery_client_new(struct idevicerestore_client_t* client)
 	}
 
 	irecv_event_subscribe(recovery, IRECV_PROGRESS, &recovery_progress_callback, NULL);
-	register_progress('RECV', "Uploading");
 	client->recovery->client = recovery;
 	return 0;
 }
@@ -317,8 +316,10 @@ int recovery_send_component(struct idevicerestore_client_t* client, plist_t buil
 	logger(LL_INFO, "Sending %s (%zu bytes)...\n", component, size);
 
 	// FIXME: Did I do this right????
+	register_progress('RECV', "Uploading");
 	err = irecv_send_buffer(client->recovery->client, data, size, 0);
 	free(data);
+	finalize_progress('RECV');
 	if (err != IRECV_E_SUCCESS) {
 		logger(LL_ERROR, "Unable to send %s component: %s\n", component, irecv_strerror(err));
 		return -1;