1 files changed, 51 insertions, 6 deletions

diff --git a/src/mobilebackup.c b/src/mobilebackup.c
index 209aee3..03d3920 100644
--- a/src/mobilebackup.c
+++ b/src/mobilebackup.c
@@ -113,8 +113,11 @@ mobilebackup_error_t mobilebackup_client_free(mobilebackup_client_t client)
 {
 	if (!client)
 		return MOBILEBACKUP_E_INVALID_ARG;
-	device_link_service_disconnect(client->parent);
-	mobilebackup_error_t err = mobilebackup_error(device_link_service_client_free(client->parent));
+	mobilebackup_error_t err = MOBILEBACKUP_E_SUCCESS;
+	if (client->parent) {
+		device_link_service_disconnect(client->parent);
+		err = mobilebackup_error(device_link_service_client_free(client->parent));
+	}
 	free(client);
 	return err;
 }
@@ -453,17 +456,59 @@ mobilebackup_error_t mobilebackup_receive_restore_file_received(mobilebackup_cli
 }
 
 /**
- * Tells the device that the restore process is complete.
+ * Tells the device that the restore process is complete and waits for the
+ * device to close the connection. After that, the device should reboot.
  *
  * @param client The connected MobileBackup client to use.
  * 
  * @return MOBILEBACKUP_E_SUCCESS on success, MOBILEBACKUP_E_INVALID_ARG if
- *    client is invalid, or MOBILEBACKUP_E_MUX_ERROR if a communication error
- *    occurs.
+ *    client is invalid, MOBILEBACKUP_E_PLIST_ERROR if the received disconnect
+ *    message plist is invalid, or MOBILEBACKUP_E_MUX_ERROR if a communication
+ *    error occurs.
  */
 mobilebackup_error_t mobilebackup_send_restore_complete(mobilebackup_client_t client)
 {
-	return mobilebackup_send_message(client, "BackupMessageRestoreComplete", NULL);
+	mobilebackup_error_t err = mobilebackup_send_message(client, "BackupMessageRestoreComplete", NULL);
+	if (err != MOBILEBACKUP_E_SUCCESS) {
+		return err;
+	}
+	plist_t dlmsg = NULL;
+	err = mobilebackup_receive(client, &dlmsg);
+	if ((err != MOBILEBACKUP_E_SUCCESS) || !dlmsg || (plist_get_node_type(dlmsg) != PLIST_ARRAY) || (plist_array_get_size(dlmsg) != 2)) {
+		if (dlmsg) {
+			debug_info("ERROR: Did not receive DLMessageDisconnect:");
+			debug_plist(dlmsg);
+			plist_free(dlmsg);
+		}
+		if (err == MOBILEBACKUP_E_SUCCESS) {
+			err = MOBILEBACKUP_E_PLIST_ERROR;
+		}
+		return err;
+	}
+	plist_t node = plist_array_get_item (dlmsg, 0);
+	char *msg = NULL;
+	if (node && (plist_get_node_type(node) == PLIST_STRING)) {
+		plist_get_string_val(node, &msg);
+	}
+
+	if (msg && !strcmp(msg, "DLMessageDisconnect")) {
+		err = MOBILEBACKUP_E_SUCCESS;
+		/* we need to do this here, otherwise mobilebackup_client_free
+		   will fail */
+		device_link_service_client_free(client->parent);
+		client->parent = NULL;
+	} else {
+		debug_info("ERROR: Malformed plist received:");
+		debug_plist(dlmsg);
+		err = MOBILEBACKUP_E_PLIST_ERROR;
+	}
+
+	plist_free(dlmsg);
+
+	if (msg)
+		free(msg);
+
+	return err;
 }
 
 /**