9 files changed, 285 insertions, 204 deletions

diff --git a/adbbu/adbbumain.cpp b/adbbu/adbbumain.cpp
index 050c9bc4a..4f5efec4c 100644
--- a/adbbu/adbbumain.cpp
+++ b/adbbu/adbbumain.cpp
@@ -31,7 +31,8 @@
 
 int main(int argc, char **argv) {
 	int index;
-	int ret = 0, pos = 0;
+	int pos = 0;
+	bool ret = false;
 	int maxpos = sizeof(TWRPARG + 2);
 	std::string command;
 	twrpback tw;
@@ -73,8 +74,9 @@ int main(int argc, char **argv) {
 	else if (command.substr(0, sizeof(TWRP_STREAM_ARG) - 1) == TWRP_STREAM_ARG) {
 		tw.setStreamFileName(argv[3]);
 		tw.threadStream();
+		ret = true;
 	}
-	if (ret == 0)
+	if (ret)
 		tw.adblogwrite("Adb backup/restore completed\n");
 	else
 		tw.adblogwrite("Adb backup/restore failed\n");
@@ -85,5 +87,8 @@ int main(int argc, char **argv) {
 		tw.adblogwrite("Unable to remove TW_ADB_BU_CONTROL: " + str.str());
 	}
 	unlink(TW_ADB_TWRP_CONTROL);
-	return ret;
+	if (ret)
+		return 0;
+	else
+		return -1;
 }
diff --git a/adbbu/libtwadbbu.cpp b/adbbu/libtwadbbu.cpp
index a13ecb2c3..0c7f355fc 100644
--- a/adbbu/libtwadbbu.cpp
+++ b/adbbu/libtwadbbu.cpp
@@ -33,6 +33,7 @@
 
 #include "twadbstream.h"
 #include "libtwadbbu.hpp"
+#include "twrpback.hpp"
 
 bool twadbbu::Check_ADB_Backup_File(std::string fname) {
 	struct AdbBackupStreamHeader adbbuhdr;
@@ -290,3 +291,16 @@ bool twadbbu::Write_TWENDADB() {
 	close(adb_control_bu_fd);
 	return true;
 }
+
+bool twadbbu::Write_TWDATA(FILE* adbd_fp) {
+	struct AdbBackupControlType data_block;
+	memset(&data_block, 0, sizeof(data_block));
+	strncpy(data_block.start_of_header, TWRP, sizeof(data_block.start_of_header));
+	strncpy(data_block.type, TWDATA, sizeof(data_block.type));
+	data_block.crc = crc32(0L, Z_NULL, 0);
+	data_block.crc = crc32(data_block.crc, (const unsigned char*) &data_block, sizeof(data_block));
+	if (fwrite(&data_block, 1, sizeof(data_block), adbd_fp) != sizeof(data_block))  {
+		return false;
+	}
+	return true;
+}
diff --git a/adbbu/libtwadbbu.hpp b/adbbu/libtwadbbu.hpp
index ff9a3eb7f..9244bb54d 100644
--- a/adbbu/libtwadbbu.hpp
+++ b/adbbu/libtwadbbu.hpp
@@ -33,7 +33,6 @@
 #include <sstream>
 
 #include "twadbstream.h"
-#include "twrpback.hpp"
 
 class twadbbu {
 public:
@@ -46,6 +45,7 @@ public:
 	static bool Write_TWEOF();                                                                     //Write ADB End-Of-File marker to stream
 	static bool Write_TWERROR();                                                                   //Write error message occurred to stream
 	static bool Write_TWENDADB();                                                                  //Write ADB End-Of-Stream command to stream
+	static bool Write_TWDATA(FILE* adbd_fp);                                                       //Write TWDATA separator
 };
 
 #endif //__LIBTWADBBU_HPP
diff --git a/adbbu/twadbstream.h b/adbbu/twadbstream.h
index 3c73084d8..fd8eba9f4 100644
--- a/adbbu/twadbstream.h
+++ b/adbbu/twadbstream.h
@@ -39,7 +39,7 @@
 #define TWMD5 "twverifymd5"				//This command is compared to the md5trailer by ORS to verify transfer
 #define TWENDADB "twendadb"				//End Protocol
 #define TWERROR "twerror"				//Send error
-#define ADB_BACKUP_VERSION 1				//Backup Version
+#define ADB_BACKUP_VERSION 2				//Backup Version
 #define DATA_MAX_CHUNK_SIZE 1048576			//Maximum size between each data header
 #define MAX_ADB_READ 512				//align with default tar size for amount to read fom adb stream
 
diff --git a/adbbu/twrpback.cpp b/adbbu/twrpback.cpp
index 2c7ea6966..d88a9c9da 100644
--- a/adbbu/twrpback.cpp
+++ b/adbbu/twrpback.cpp
@@ -36,6 +36,7 @@
 
 #include "twadbstream.h"
 #include "twrpback.hpp"
+#include "libtwadbbu.hpp"
 #include "../twrpDigest/twrpDigest.hpp"
 #include "../twrpDigest/twrpMD5.hpp"
 #include "../twrpAdbBuFifo.hpp"
@@ -124,13 +125,13 @@ void twrpback::close_restore_fds() {
 		unlink(TW_ADB_RESTORE);
 }
 
-int twrpback::backup(std::string command) {
+bool twrpback::backup(std::string command) {
 	twrpMD5 digest;
 	bool breakloop = false;
 	int bytes = 0, errctr = 0;
-	char result[MAX_ADB_READ];
-	uint64_t totalbytes = 0, dataChunkBytes = 0;
-	int64_t count = -1;			// Count of how many blocks set
+	char adbReadStream[MAX_ADB_READ];
+	uint64_t totalbytes = 0, dataChunkBytes = 0, fileBytes = 0;
+	int64_t count = false;			// Count of how many blocks set
 	uint64_t md5fnsize = 0;
 	struct AdbBackupControlType endadb;
 
@@ -143,12 +144,12 @@ int twrpback::backup(std::string command) {
 	adbd_fp = fdopen(adbd_fd, "w");
 	if (adbd_fp == NULL) {
 		adblogwrite("Unable to open adb_fp\n");
-		return -1;
+		return false;
 	}
 
 	if (mkfifo(TW_ADB_BACKUP, 0666) < 0) {
 		adblogwrite("Unable to create TW_ADB_BACKUP fifo\n");
-		return -1;
+		return false;
 	}
 
 	adblogwrite("opening TW_ADB_FIFO\n");
@@ -160,7 +161,7 @@ int twrpback::backup(std::string command) {
 		if (errctr > ADB_BU_MAX_ERROR) {
 			adblogwrite("Unable to open TW_ADB_FIFO\n");
 			close_backup_fds();
-			return -1;
+			return false;
 		}
 	}
 
@@ -168,15 +169,15 @@ int twrpback::backup(std::string command) {
 	if (snprintf(operation, sizeof(operation), "adbbackup %s", command.c_str()) >= sizeof(operation)) {
 		adblogwrite("Operation too big to write to ORS_INPUT_FILE\n");
 		close_backup_fds();
-		return -1;
+		return false;
 	}
 	if (write(write_fd, operation, sizeof(operation)) != sizeof(operation)) {
 		adblogwrite("Unable to write to ORS_INPUT_FILE\n");
 		close_backup_fds();
-		return -1;
+		return false;
 	}
 
-	memset(&result, 0, sizeof(result));
+	memset(&adbReadStream, 0, sizeof(adbReadStream));
 	memset(&cmd, 0, sizeof(cmd));
 
 	adblogwrite("opening TW_ADB_BU_CONTROL\n");
@@ -184,7 +185,7 @@ int twrpback::backup(std::string command) {
 	if (adb_control_bu_fd < 0) {
 		adblogwrite("Unable to open TW_ADB_BU_CONTROL for reading.\n");
 		close_backup_fds();
-		return -1;
+		return false;
 	}
 
 	adblogwrite("opening TW_ADB_BACKUP\n");
@@ -192,7 +193,7 @@ int twrpback::backup(std::string command) {
 	if (adb_read_fd < 0) {
 		adblogwrite("Unable to open TW_ADB_BACKUP for reading.\n");
 		close_backup_fds();
-		return -1;
+		return false;
 	}
 
 	//loop until TWENDADB sent
@@ -208,7 +209,7 @@ int twrpback::backup(std::string command) {
 				writedata = false;
 				adblogwrite("Error received. Quitting...\n");
 				close_backup_fds();
-				return -1;
+				return false;
 			}
 			//we received the end of adb backup stream so we should break the loop
 			else if (cmdtype == TWENDADB) {
@@ -223,13 +224,13 @@ int twrpback::backup(std::string command) {
 			//we recieved the TWSTREAMHDR structure metadata to write to adb
 			else if (cmdtype == TWSTREAMHDR) {
 				writedata = false;
-				adblogwrite("Writing TWSTREAMHDR\n");
+				adblogwrite("writing TWSTREAMHDR\n");
 				if (fwrite(cmd, 1, sizeof(cmd), adbd_fp) != sizeof(cmd)) {
 					std::stringstream str;
 					str << strerror(errno);
 					adblogwrite("Error writing TWSTREAMHDR to adbd" + str.str() + "\n");
 					close_backup_fds();
-					return -1;
+					return false;
 				}
 				fflush(adbd_fp);
 			}
@@ -237,16 +238,17 @@ int twrpback::backup(std::string command) {
 			else if (cmdtype == TWIMG) {
 				struct twfilehdr twimghdr;
 
-				adblogwrite("Writing TWIMG\n");
+				adblogwrite("writing TWIMG\n");
 				digest.init();
 				memset(&twimghdr, 0, sizeof(twimghdr));
 				memcpy(&twimghdr, cmd, sizeof(cmd));
 				md5fnsize = twimghdr.size;
+				compressed = false;
 
-				if (fwrite(cmd, 1, sizeof(cmd), adbd_fp)  != sizeof(cmd)) {
+				if (fwrite(cmd, 1, sizeof(cmd), adbd_fp) != sizeof(cmd)) {
 					adblogwrite("Error writing TWIMG to adbd\n");
 					close_backup_fds();
-					return -1;
+					return false;
 				}
 				fflush(adbd_fp);
 				writedata = true;
@@ -255,7 +257,7 @@ int twrpback::backup(std::string command) {
 			else if (cmdtype == TWFN) {
 				struct twfilehdr twfilehdr;
 
-				adblogwrite("Writing TWFN\n");
+				adblogwrite("writing TWFN\n");
 				digest.init();
 
 				ADBSTRUCT_STATIC_ASSERT(sizeof(twfilehdr) == MAX_ADB_READ);
@@ -269,7 +271,7 @@ int twrpback::backup(std::string command) {
 				if (fwrite(cmd, 1, sizeof(cmd), adbd_fp) != sizeof(cmd)) {
 					adblogwrite("Error writing TWFN to adbd\n");
 					close_backup_fds();
-					return -1;
+					return false;
 				}
 				fflush(adbd_fp);
 				writedata = true;
@@ -284,36 +286,43 @@ int twrpback::backup(std::string command) {
 			*/
 			else if (cmdtype == TWEOF) {
 				adblogwrite("received TWEOF\n");
-				count = totalbytes / MAX_ADB_READ + 1;
-				count = count * MAX_ADB_READ;
-
-				while ((bytes = read(adb_read_fd, &result, sizeof(result))) > 0) {
+				while ((bytes = read(adb_read_fd, &adbReadStream, sizeof(adbReadStream)) != 0)) {
 					totalbytes += bytes;
-					char *writeresult = new char [bytes];
-					memcpy(writeresult, result, bytes);
-					digest.update((unsigned char *) writeresult, bytes);
-					if (fwrite(writeresult, 1, bytes, adbd_fp) != bytes) {
-						adblogwrite("Error writing backup data to adbd\n");
-						close_backup_fds();
-						return -1;
+					fileBytes += bytes;
+					dataChunkBytes += bytes;
+
+					char *writeAdbReadStream = new char [bytes];
+					memcpy(writeAdbReadStream, adbReadStream, bytes);
+
+					digest.update((unsigned char *) writeAdbReadStream, bytes);
+					if (fwrite(writeAdbReadStream, 1, bytes, adbd_fp) < 0) {
+						std::stringstream str;
+						str << strerror(errno);
+						adblogwrite("Cannot write to adbd stream: " + str.str() + "\n");
 					}
 					fflush(adbd_fp);
-					delete [] writeresult;
-					memset(&result, 0, sizeof(result));
+					delete [] writeAdbReadStream;
+					memset(adbReadStream, 0, sizeof(adbReadStream));
 				}
 
-				if ((totalbytes % MAX_ADB_READ) != 0) {
-					adblogwrite("writing padding to stream\n");
-					char padding[count - totalbytes];
-					memset(padding, 0, sizeof(padding));
-					if (fwrite(padding, 1, sizeof(padding), adbd_fp) != sizeof(padding)) {
+				count = fileBytes / DATA_MAX_CHUNK_SIZE + 1;
+				count = count * DATA_MAX_CHUNK_SIZE;
+
+				if (fileBytes % DATA_MAX_CHUNK_SIZE != 0) {
+					char padding[count - fileBytes];
+					int paddingBytes = sizeof(padding);
+					std::stringstream paddingStr;
+					paddingStr << paddingBytes;
+					memset(padding, 0, paddingBytes);
+					adblogwrite("writing padding to stream: " + paddingStr.str() + " bytes\n");
+					if (fwrite(padding, 1, paddingBytes, adbd_fp) != sizeof(padding)) {
 						adblogwrite("Error writing padding to adbd\n");
 						close_backup_fds();
-						return -1;
+						return false;
 					}
-					digest.update((unsigned char *) padding, sizeof(padding));
+					totalbytes += paddingBytes;
+					digest.update((unsigned char *) padding, paddingBytes);
 					fflush(adbd_fp);
-					totalbytes = 0;
 				}
 
 				AdbBackupFileTrailer md5trailer;
@@ -336,11 +345,12 @@ int twrpback::backup(std::string command) {
 				if (fwrite(&md5trailer, 1, sizeof(md5trailer), adbd_fp) != sizeof(md5trailer))  {
 					adblogwrite("Error writing md5trailer to adbd\n");
 					close_backup_fds();
-					return -1;
+					return false;
 				}
 				fflush(adbd_fp);
 				writedata = false;
 				firstDataPacket = true;
+				fileBytes = 0;
 			}
 			memset(&cmd, 0, sizeof(cmd));
 		}
@@ -349,59 +359,76 @@ int twrpback::backup(std::string command) {
 		//to the adb stream.
 		//If the stream is compressed, we need to always write the data.
 		if (writedata || compressed) {
-			while ((bytes = read(adb_read_fd, &result, sizeof(result))) > 0) {
+			while ((bytes = read(adb_read_fd, &adbReadStream, sizeof(adbReadStream))) > 0) {
 				if (firstDataPacket) {
-					struct AdbBackupControlType data_block;
-
-					memset(&data_block, 0, sizeof(data_block));
-					strncpy(data_block.start_of_header, TWRP, sizeof(data_block.start_of_header));
-					strncpy(data_block.type, TWDATA, sizeof(data_block.type));
-					data_block.crc = crc32(0L, Z_NULL, 0);
-					data_block.crc = crc32(data_block.crc, (const unsigned char*) &data_block, sizeof(data_block));
-					if (fwrite(&data_block, 1, sizeof(data_block), adbd_fp) != sizeof(data_block))  {
-						adblogwrite("Error writing data_block to adbd\n");
+					if (!twadbbu::Write_TWDATA(adbd_fp)) {
 						close_backup_fds();
-						return -1;
+						return false;
 					}
 					fflush(adbd_fp);
 					firstDataPacket = false;
+					dataChunkBytes += sizeof(adbReadStream);
 				}
-				char *writeresult = new char [bytes];
-				memcpy(writeresult, result, bytes);
+				char *writeAdbReadStream = new char [bytes];
+				memcpy(writeAdbReadStream, adbReadStream, bytes);
 
-				digest.update((unsigned char *) writeresult, bytes);
+				digest.update((unsigned char *) writeAdbReadStream, bytes);
 
 				totalbytes += bytes;
+				fileBytes += bytes;
 				dataChunkBytes += bytes;
 
-				if (fwrite(writeresult, 1, bytes, adbd_fp) != bytes) {
+				if (fwrite(writeAdbReadStream, 1, bytes, adbd_fp) != bytes) {
 					adblogwrite("Error writing backup data to adbd\n");
 					close_backup_fds();
-					return -1;
+					return false;
 				}
 				fflush(adbd_fp);
+				delete [] writeAdbReadStream;
 
-				delete [] writeresult;
-				memset(&result, 0, sizeof(result));
-				if (dataChunkBytes == DATA_MAX_CHUNK_SIZE - sizeof(result)) {
-					struct AdbBackupControlType data_block;
-
-					memset(&data_block, 0, sizeof(data_block));
-					strncpy(data_block.start_of_header, TWRP, sizeof(data_block.start_of_header));
-					strncpy(data_block.type, TWDATA, sizeof(data_block.type));
-					data_block.crc = crc32(0L, Z_NULL, 0);
-					data_block.crc = crc32(data_block.crc, (const unsigned char*) &data_block, sizeof(data_block));
-					if (fwrite(&data_block, 1, sizeof(data_block), adbd_fp) != sizeof(data_block))  {
-						adblogwrite("Error writing data_block to adbd\n");
-						close_backup_fds();
-						return -1;
+				memset(&adbReadStream, 0, sizeof(adbReadStream));
+
+				if (dataChunkBytes == DATA_MAX_CHUNK_SIZE) {
+					dataChunkBytes = 0;
+					firstDataPacket = true;
+				}
+				else if (dataChunkBytes > (DATA_MAX_CHUNK_SIZE - sizeof(adbReadStream))) {
+					int bytesLeft = DATA_MAX_CHUNK_SIZE - dataChunkBytes;
+					char extraData[bytesLeft];
+
+					memset(&extraData, 0, bytesLeft);
+					while ((bytes = read(adb_read_fd, &extraData, bytesLeft)) != 0) {
+						if (bytes > 0) {
+							totalbytes += bytes;
+							fileBytes += bytes;
+							dataChunkBytes += bytes;
+
+							bytesLeft -= bytes;
+							char *writeAdbReadStream = new char [bytes];
+							memcpy(writeAdbReadStream, extraData, bytes);
+
+							digest.update((unsigned char *) writeAdbReadStream, bytes);
+							if (fwrite(writeAdbReadStream, 1, bytes, adbd_fp) < 0) {
+								std::stringstream str;
+								str << strerror(errno);
+								adblogwrite("Cannot write to adbd stream: " + str.str() + "\n");
+								close_restore_fds();
+								return false;
+							}
+							fflush(adbd_fp);
+							delete [] writeAdbReadStream;
+						}
+						memset(&extraData, 0, bytesLeft);
+						if (bytesLeft == 0) {
+							break;
+						}
 					}
+
 					fflush(adbd_fp);
 					dataChunkBytes = 0;
+					firstDataPacket = true;
 				}
-
 			}
-			compressed = false;
 		}
 	}
 
@@ -409,23 +436,25 @@ int twrpback::backup(std::string command) {
 	if (fwrite(&endadb, 1, sizeof(endadb), adbd_fp) != sizeof(endadb)) {
 		adblogwrite("Error writing endadb to adbd\n");
 		close_backup_fds();
-		return -1;
+		return false;
 	}
 	fflush(adbd_fp);
 	close_backup_fds();
 	return 0;
 }
 
-int twrpback::restore(void) {
+bool twrpback::restore(void) {
 	twrpMD5 digest;
 	char cmd[MAX_ADB_READ];
-	char result[MAX_ADB_READ];
+	char readAdbStream[MAX_ADB_READ];
 	struct AdbBackupControlType structcmd;
-	int adb_control_twrp_fd, errctr = 0;
+	int errctr = 0;
 	uint64_t totalbytes = 0, dataChunkBytes = 0;
 	uint64_t md5fnsize = 0;
 	bool writedata, read_from_adb;
 	bool breakloop, eofsent, md5trsent;
+	bool compressed;
+	bool md5TrailerReceived = false;
 
 	breakloop = false;
 	read_from_adb = true;
@@ -436,13 +465,13 @@ int twrpback::restore(void) {
 	if (adbd_fp == NULL) {
 		adblogwrite("Unable to open adb_fp\n");
 		close_restore_fds();
-		return -1;
+		return false;
 	}
 
 	if(mkfifo(TW_ADB_RESTORE, 0666)) {
 		adblogwrite("Unable to create TW_ADB_RESTORE fifo\n");
 		close_restore_fds();
-		return -1;
+		return false;
 	}
 
 	adblogwrite("opening TW_ADB_FIFO\n");
@@ -454,7 +483,7 @@ int twrpback::restore(void) {
 		if (errctr > ADB_BU_MAX_ERROR) {
 			adblogwrite("Unable to open TW_ADB_FIFO\n");
 			close_restore_fds();
-			return -1;
+			return false;
 		}
 	}
 
@@ -463,10 +492,10 @@ int twrpback::restore(void) {
 	if (write(write_fd, operation, sizeof(operation)) != sizeof(operation)) {
 		adblogwrite("Unable to write to TW_ADB_FIFO\n");
 		close_restore_fds();
-		return -1;
+		return false;
 	}
 
-	memset(&result, 0, sizeof(result));
+	memset(&readAdbStream, 0, sizeof(readAdbStream));
 	memset(&cmd, 0, sizeof(cmd));
 
 	adblogwrite("opening TW_ADB_BU_CONTROL\n");
@@ -476,7 +505,7 @@ int twrpback::restore(void) {
 		str << strerror(errno);
 		adblogwrite("Unable to open TW_ADB_BU_CONTROL for writing. " + str.str() + "\n");
 		close_restore_fds();
-		return -1;
+		return false;
 	}
 
 	adblogwrite("opening TW_ADB_TWRP_CONTROL\n");
@@ -492,7 +521,7 @@ int twrpback::restore(void) {
 			if (errctr > ADB_BU_MAX_ERROR) {
 				adblogwrite("Unable to open TW_ADB_TWRP_CONTROL\n");
 				close_backup_fds();
-				return -1;
+				return false;
 			}
 		}
 	}
@@ -511,7 +540,7 @@ int twrpback::restore(void) {
 				struct AdbBackupControlType tweof;
 
 				memset(&tweof, 0, sizeof(tweof));
-				memcpy(&tweof, result, sizeof(result));
+				memcpy(&tweof, readAdbStream, sizeof(readAdbStream));
 				read_from_adb = true;
 			}
 			//Break when TWRP sends TWENDADB
@@ -524,15 +553,14 @@ int twrpback::restore(void) {
 			else if (cmdtype == TWERROR) {
 				adblogwrite("Error received. Quitting...\n");
 				close_restore_fds();
-				return -1;
+				return false;
 			}
 		}
 		//If we should read from the adb stream, write commands and data to TWRP
 		if (read_from_adb) {
 			int readbytes;
-			if ((readbytes = fread(result, 1, sizeof(result), adbd_fp)) == sizeof(result)) {
-				totalbytes += readbytes;
-				memcpy(&structcmd, result, sizeof(result));
+			if ((readbytes = fread(readAdbStream, 1, sizeof(readAdbStream), adbd_fp)) == sizeof(readAdbStream)) {
+				memcpy(&structcmd, readAdbStream, sizeof(readAdbStream));
 				std::string cmdtype = structcmd.get_type();
 
 				//Tell TWRP we have read the entire adb stream
@@ -540,9 +568,8 @@ int twrpback::restore(void) {
 					struct AdbBackupControlType endadb;
 					uint32_t crc, endadbcrc;
 
-					totalbytes -= sizeof(result);
 					memset(&endadb, 0, sizeof(endadb));
-					memcpy(&endadb, result, sizeof(result));
+					memcpy(&endadb, readAdbStream, sizeof(readAdbStream));
 					endadbcrc = endadb.crc;
 					memset(&endadb.crc, 0, sizeof(endadb.crc));
 					crc = crc32(0L, Z_NULL, 0);
@@ -555,14 +582,14 @@ int twrpback::restore(void) {
 							str << strerror(errno);
 							adblogwrite("Cannot write to ADB_CONTROL_READ_FD: " + str.str() + "\n");
 							close_restore_fds();
-							return -1;
+							return false;
 						}
 						read_from_adb = false;
 					}
 					else {
 						adblogwrite("ADB TWENDADB crc header doesn't match\n");
 						close_restore_fds();
-						return -1;
+						return false;
 					}
 				}
 				//Send TWRP partition metadata
@@ -571,10 +598,9 @@ int twrpback::restore(void) {
 					uint32_t crc, cnthdrcrc;
 
 					ADBSTRUCT_STATIC_ASSERT(sizeof(cnthdr) == MAX_ADB_READ);
-					totalbytes -= sizeof(result);
 
 					memset(&cnthdr, 0, sizeof(cnthdr));
-					memcpy(&cnthdr, result, sizeof(result));
+					memcpy(&cnthdr, readAdbStream, sizeof(readAdbStream));
 					cnthdrcrc = cnthdr.crc;
 					memset(&cnthdr.crc, 0, sizeof(cnthdr.crc));
 					crc = crc32(0L, Z_NULL, 0);
@@ -582,18 +608,18 @@ int twrpback::restore(void) {
 
 					if (crc == cnthdrcrc) {
 						adblogwrite("Restoring TWSTREAMHDR\n");
-						if (write(adb_control_twrp_fd, result, sizeof(result)) < 0) {
+						if (write(adb_control_twrp_fd, readAdbStream, sizeof(readAdbStream)) < 0) {
 							std::stringstream str;
 							str << strerror(errno);
 							adblogwrite("Cannot write to adb_control_twrp_fd: " + str.str() + "\n");
 							close_restore_fds();
-							return -1;
+							return false;
 						}
 					}
 					else {
 						adblogwrite("ADB TWSTREAMHDR crc header doesn't match\n");
 						close_restore_fds();
-						return -1;
+						return false;
 					}
 				}
 				//Tell TWRP we are sending a partition image
@@ -602,10 +628,9 @@ int twrpback::restore(void) {
 					uint32_t crc, twimghdrcrc;
 
 					digest.init();
-					totalbytes -= sizeof(result);
 					adblogwrite("Restoring TWIMG\n");
 					memset(&twimghdr, 0, sizeof(twimghdr));
-					memcpy(&twimghdr, result, sizeof(result));
+					memcpy(&twimghdr, readAdbStream, sizeof(readAdbStream));
 					md5fnsize = twimghdr.size;
 					twimghdrcrc = twimghdr.crc;
 					memset(&twimghdr.crc, 0, sizeof(twimghdr.crc));
@@ -613,18 +638,18 @@ int twrpback::restore(void) {
 					crc = crc32(0L, Z_NULL, 0);
 					crc = crc32(crc, (const unsigned char*) &twimghdr, sizeof(twimghdr));
 					if (crc == twimghdrcrc) {
-						if (write(adb_control_twrp_fd, result, sizeof(result)) < 1) {
+						if (write(adb_control_twrp_fd, readAdbStream, sizeof(readAdbStream)) < 1) {
 							std::stringstream str;
 							str << strerror(errno);
 							adblogwrite("Cannot write to adb_control_twrp_fd: " + str.str() + "\n");
 							close_restore_fds();
-							return -1;
+							return false;
 						}
 					}
 					else {
 						adblogwrite("ADB TWIMG crc header doesn't match\n");
 						close_restore_fds();
-						return -1;
+						return false;
 					}
 					adblogwrite("opening TW_ADB_RESTORE\n");
 					adb_write_fd = open(TW_ADB_RESTORE, O_WRONLY);
@@ -633,12 +658,11 @@ int twrpback::restore(void) {
 				else if (cmdtype == TWFN) {
 					struct twfilehdr twfilehdr;
 					uint32_t crc, twfilehdrcrc;
-					digest.init();
 
-					totalbytes -= sizeof(result);
+					digest.init();
 					adblogwrite("Restoring TWFN\n");
 					memset(&twfilehdr, 0, sizeof(twfilehdr));
-					memcpy(&twfilehdr, result, sizeof(result));
+					memcpy(&twfilehdr, readAdbStream, sizeof(readAdbStream));
 					md5fnsize = twfilehdr.size;
 					twfilehdrcrc = twfilehdr.crc;
 					memset(&twfilehdr.crc, 0, sizeof(twfilehdr.crc));
@@ -647,111 +671,82 @@ int twrpback::restore(void) {
 					crc = crc32(crc, (const unsigned char*) &twfilehdr, sizeof(twfilehdr));
 
 					if (crc == twfilehdrcrc) {
-						if (write(adb_control_twrp_fd, result, sizeof(result)) < 1) {
+						if (write(adb_control_twrp_fd, readAdbStream, sizeof(readAdbStream)) < 1) {
 							std::stringstream str;
 							str << strerror(errno);
 							adblogwrite("Cannot write to adb_control_twrp_fd: " + str.str() + "\n");
 							close_restore_fds();
-							return -1;
+							return false;
 						}
 					}
 					else {
 						adblogwrite("ADB TWFN crc header doesn't match\n");
 						close_restore_fds();
-						return -1;
+						return false;
 					}
 
 					adblogwrite("opening TW_ADB_RESTORE\n");
 					adb_write_fd = open(TW_ADB_RESTORE, O_WRONLY);
 				}
+				else if (cmdtype == MD5TRAILER) {
+					read_from_adb = false; //don't read from adb until TWRP sends TWEOF
+					close(adb_write_fd);
+					md5TrailerReceived = true;
+					if (!checkMD5Trailer(readAdbStream, md5fnsize, &digest)) {
+						close_restore_fds();
+						return false;
+					}
+				}
 				//Send the tar or partition image md5 to TWRP
 				else if (cmdtype == TWDATA) {
-					totalbytes -= sizeof(result);
+					dataChunkBytes += sizeof(readAdbStream);
 					while (1) {
-						if ((readbytes = fread(result, 1, sizeof(result), adbd_fp)) != sizeof(result)) {
+						if ((readbytes = fread(readAdbStream, 1, sizeof(readAdbStream), adbd_fp)) != sizeof(readAdbStream)) {
 							close_restore_fds();
-							return -1;
+							return false;
 						}
-						totalbytes += readbytes;
-						memcpy(&structcmd, result, sizeof(result));
-						std::string cmdtype = structcmd.get_type();
 
-						if (cmdtype.substr(0, sizeof(MD5TRAILER) - 1) == MD5TRAILER) {
-							struct AdbBackupFileTrailer md5tr;
-							uint32_t crc, md5trcrc, md5ident, md5identmatch;
-
-							ADBSTRUCT_STATIC_ASSERT(sizeof(md5tr) == MAX_ADB_READ);
-							memset(&md5tr, 0, sizeof(md5tr));
-							memcpy(&md5tr, result, sizeof(result));
-							md5ident = md5tr.ident;
-
-							memset(&md5tr.ident, 0, sizeof(md5tr.ident));
-
-							md5identmatch = crc32(0L, Z_NULL, 0);
-							md5identmatch = crc32(md5identmatch, (const unsigned char*) &md5tr, sizeof(md5tr));
-							md5identmatch = crc32(md5identmatch, (const unsigned char*) &md5fnsize, sizeof(md5fnsize));
-
-							if (md5identmatch == md5ident) {
-								totalbytes -= sizeof(result);
-								close(adb_write_fd);
-								adblogwrite("Restoring MD5TRAILER\n");
-								md5trcrc = md5tr.crc;
-								memset(&md5tr.crc, 0, sizeof(md5tr.crc));
-								crc = crc32(0L, Z_NULL, 0);
-								crc = crc32(crc, (const unsigned char*) &md5tr, sizeof(md5tr));
-								if (crc == md5trcrc) {
-									if (write(adb_control_twrp_fd, result, sizeof(result)) < 1) {
-										std::stringstream str;
-										str << strerror(errno);
-										adblogwrite("Cannot write to adb_control_twrp_fd: " + str.str() + "\n");
-										close_restore_fds();
-										return -1;
-									}
-								}
-								else {
-									adblogwrite("ADB MD5TRAILER crc header doesn't match\n");
-									close_restore_fds();
-									return -1;
-								}
-
-								AdbBackupFileTrailer md5;
-
-								memset(&md5, 0, sizeof(md5));
-								strncpy(md5.start_of_trailer, TWRP, sizeof(md5.start_of_trailer));
-								strncpy(md5.type, TWMD5, sizeof(md5.type));
-								std::string md5string = digest.return_digest_string();
-								strncpy(md5.md5, md5string.c_str(), sizeof(md5.md5));
-
-								adblogwrite("Sending MD5Check\n");
-								if (write(adb_control_twrp_fd, &md5, sizeof(md5)) < 1) {
-									std::stringstream str;
-									str << strerror(errno);
-									adblogwrite("Cannot write to adb_control_twrp_fd: " + str.str() + "\n");
-									close_restore_fds();
-									return -1;
-								}
-								read_from_adb = false; //don't read from adb until TWRP sends TWEOF
-								break;
+						memcpy(&structcmd, readAdbStream, sizeof(readAdbStream));
+
+						char *readAdbReadStream = new char [readbytes];
+						memcpy(readAdbReadStream, readAdbStream, readbytes);
+						std::string cmdtype = structcmd.get_type();
+						dataChunkBytes += readbytes;
+						delete [] readAdbReadStream;
+						totalbytes += readbytes;
+						digest.update((unsigned char*)readAdbReadStream, readbytes);
+
+						if (cmdtype == MD5TRAILER) {
+							read_from_adb = false; //don't read from adb until TWRP sends TWEOF
+							close(adb_write_fd);
+							if (!checkMD5Trailer(readAdbStream, md5fnsize, &digest)) {
+								close_restore_fds();
+								return false;
 							}
+							break;
 						}
-						digest.update((unsigned char*)result, sizeof(result));
-						dataChunkBytes += readbytes;
 
-						if (write(adb_write_fd, result, sizeof(result)) < 0) {
-							std::stringstream str;
-							str << strerror(errno);
-							adblogwrite("Cannot write to adb_write_fd\n" + str.str() + ". Retrying.\n");
-							while(write(adb_write_fd, result, sizeof(result)) < 0) {
-								adblogwrite("Cannot write to adb_write_fd\n" + str.str() + ". Retrying.\n");
-								continue;
-							}
+
+						if (write(adb_write_fd, readAdbStream, sizeof(readAdbStream)) < 0) {
+							adblogwrite("end of stream reached.\n");
+							break;
 						}
-						if (dataChunkBytes == ((DATA_MAX_CHUNK_SIZE) - sizeof(result))) {
+						if (dataChunkBytes == DATA_MAX_CHUNK_SIZE) {
 							dataChunkBytes = 0;
 							break;
 						}
-						memset(&result, 0, sizeof(result));
+						memset(&readAdbStream, 0, sizeof(readAdbStream));
+					}
+				}
+				else {
+					if (!md5TrailerReceived) {
+						char *readAdbReadStream = new char [readbytes];
+						memcpy(readAdbReadStream, readAdbStream, readbytes);
+						digest.update((unsigned char*)readAdbReadStream, readbytes);
+						totalbytes += readbytes;
+						delete [] readAdbReadStream;
 					}
+
 				}
 			}
 		}
@@ -760,7 +755,7 @@ int twrpback::restore(void) {
 	std::stringstream str;
 	str << totalbytes;
 	adblogwrite(str.str() + " bytes restored from adbbackup\n");
-	return 0;
+	return true;
 }
 
 void twrpback::streamFileForTWRP(void) {
@@ -785,3 +780,62 @@ void twrpback::threadStream(void) {
 	pthread_create(&thread, NULL, p, this);
 	pthread_join(thread, NULL);
 }
+
+bool twrpback::checkMD5Trailer(char readAdbStream[], uint64_t md5fnsize, twrpMD5 *digest) {
+	struct AdbBackupFileTrailer md5tr;
+	uint32_t crc, md5trcrc, md5ident, md5identmatch;
+
+	ADBSTRUCT_STATIC_ASSERT(sizeof(md5tr) == MAX_ADB_READ);
+	memset(&md5tr, 0, sizeof(md5tr));
+	memcpy(&md5tr, readAdbStream, MAX_ADB_READ);
+	md5ident = md5tr.ident;
+
+	memset(&md5tr.ident, 0, sizeof(md5tr.ident));
+
+	md5identmatch = crc32(0L, Z_NULL, 0);
+	md5identmatch = crc32(md5identmatch, (const unsigned char*) &md5tr, sizeof(md5tr));
+	md5identmatch = crc32(md5identmatch, (const unsigned char*) &md5fnsize, sizeof(md5fnsize));
+
+	if (md5identmatch == md5ident) {
+		adblogwrite("checking MD5TRAILER\n");
+		md5trcrc = md5tr.crc;
+		memset(&md5tr.crc, 0, sizeof(md5tr.crc));
+		crc = crc32(0L, Z_NULL, 0);
+		crc = crc32(crc, (const unsigned char*) &md5tr, sizeof(md5tr));
+		if (crc == md5trcrc) {
+			if (write(adb_control_twrp_fd, &md5tr, sizeof(md5tr)) < 1) {
+				std::stringstream str;
+				str << strerror(errno);
+				adblogwrite("Cannot write to adb_control_twrp_fd: " + str.str() + "\n");
+				close_restore_fds();
+				return false;
+			}
+		}
+		else {
+			adblogwrite("ADB MD5TRAILER crc header doesn't match\n");
+			close_restore_fds();
+			return false;
+		}
+
+		AdbBackupFileTrailer md5;
+
+		memset(&md5, 0, sizeof(md5));
+		strncpy(md5.start_of_trailer, TWRP, sizeof(md5.start_of_trailer));
+		strncpy(md5.type, TWMD5, sizeof(md5.type));
+		std::string md5string = digest->return_digest_string();
+		strncpy(md5.md5, md5string.c_str(), sizeof(md5.md5));
+
+		adblogwrite("sending MD5 verification\n");
+		std::stringstream dstr;
+		dstr << adb_control_twrp_fd;
+		if (write(adb_control_twrp_fd, &md5, sizeof(md5)) < 1) {
+			std::stringstream str;
+			str << strerror(errno);
+			adblogwrite("Cannot write to adb_control_twrp_fd: " + str.str() + "\n");
+			close_restore_fds();
+			return false;
+		}
+		return true;
+	}
+	return false;
+}
diff --git a/adbbu/twrpback.hpp b/adbbu/twrpback.hpp
index 1c6b09f15..c52da3ea7 100644
--- a/adbbu/twrpback.hpp
+++ b/adbbu/twrpback.hpp
@@ -18,14 +18,15 @@
 #define _TWRPBACK_HPP
 
 #include <fstream>
+#include "../twrpDigest/twrpMD5.hpp"
 
 class twrpback {
 public:
 	int adbd_fd;                                                             // adbd data stream
 	twrpback(void);
 	virtual ~twrpback(void);
-	int backup(std::string command);                                         // adb backup stream
-	int restore(void);                                                       // adb restore stream
+	bool backup(std::string command);                                        // adb backup stream
+	bool restore(void);                                                      // adb restore stream
 	void adblogwrite(std::string writemsg);                                  // adb debugging log function
 	void createFifos(void);                                                  // create fifos needed for adb backup
 	void closeFifos(void);                                                   // close created fifos
@@ -52,6 +53,7 @@ private:
 	void adbloginit(void);                                                   // setup adb log stream file
 	void close_backup_fds();                                                 // close backup resources
 	void close_restore_fds();                                                // close restore resources
+	bool checkMD5Trailer(char adbReadStream[], uint64_t md5fnsize, twrpMD5* digest); // Check MD5 Trailer
 };
 
 #endif // _TWRPBACK_HPP
diff --git a/twrpAdbBuFifo.cpp b/twrpAdbBuFifo.cpp
index db34c5ac2..0c7dd1524 100644
--- a/twrpAdbBuFifo.cpp
+++ b/twrpAdbBuFifo.cpp
@@ -39,8 +39,9 @@ twrpAdbBuFifo::twrpAdbBuFifo(void) {
 	unlink(TW_ADB_FIFO);
 }
 
-bool twrpAdbBuFifo::Check_Adb_Fifo_For_Events(void) {
+void twrpAdbBuFifo::Check_Adb_Fifo_For_Events(void) {
 	char cmd[512];
+	int ret;
 
 	memset(&cmd, 0, sizeof(cmd));
 
@@ -51,13 +52,11 @@ bool twrpAdbBuFifo::Check_Adb_Fifo_For_Events(void) {
 		std::string Options(cmd);
 		Options = Options.substr(strlen(ADB_BACKUP_OP) + 1, strlen(cmd));
 		if (cmdcheck == ADB_BACKUP_OP)
-			return Backup_ADB_Command(Options);
+			Backup_ADB_Command(Options);
 		else {
-			return Restore_ADB_Backup();
+			Restore_ADB_Backup();
 		}
 	}
-
-	return true;
 }
 
 bool twrpAdbBuFifo::start(void) {
@@ -195,8 +194,7 @@ bool twrpAdbBuFifo::Restore_ADB_Backup(void) {
 
 			memset(&cmdstruct, 0, sizeof(cmdstruct));
 			memcpy(&cmdstruct, cmd, sizeof(cmdstruct));
-			std::string cmdstr(cmdstruct.type);
-			std::string cmdtype = cmdstr.substr(0, sizeof(cmdstruct.type) - 1);
+			std::string cmdtype = cmdstruct.get_type();
 			if (cmdtype == TWSTREAMHDR) {
 				struct AdbBackupStreamHeader twhdr;
 				memcpy(&twhdr, cmd, sizeof(cmd));
@@ -229,6 +227,8 @@ bool twrpAdbBuFifo::Restore_ADB_Backup(void) {
 					LOGINFO("adbrestore md5 matches\n");
 					LOGINFO("adbmd5.md5: %s\n", adbmd5.md5);
 					LOGINFO("md5check.md5: %s\n", md5check.md5);
+					ret = true;
+					break;
 				}
 			}
 			else if (cmdtype == TWENDADB) {
@@ -269,7 +269,7 @@ bool twrpAdbBuFifo::Restore_ADB_Backup(void) {
 					part_settings.progress = &progress;
 					if (!PartitionManager.Restore_Partition(&part_settings)) {
 						LOGERR("ADB Restore failed.\n");
-						return false;
+						ret = false;
 					}
 				}
 				else if (cmdtype == TWFN) {
@@ -319,18 +319,24 @@ bool twrpAdbBuFifo::Restore_ADB_Backup(void) {
 					part_settings.progress = &progress;
 					if (!PartitionManager.Restore_Partition(&part_settings)) {
 						LOGERR("ADB Restore failed.\n");
-						return false;
+						ret = false;
 					}
 				}
 			}
 		}
 	}
-	gui_msg("restore_complete=Restore Complete");
+
+	if (ret != false)
+		gui_msg("restore_complete=Restore Complete");
+	else
+		gui_err("restore_error=Error during restore process.");
 
 	if (!twadbbu::Write_TWENDADB())
 		ret = false;
 	sleep(2); //give time for user to see messages on console
 	DataManager::SetValue("ui_progress", 100);
 	gui_changePage("main");
+	close(adb_control_bu_fd);
+	close(adb_control_twrp_fd);
 	return ret;
 }
diff --git a/twrpAdbBuFifo.hpp b/twrpAdbBuFifo.hpp
index e709b9671..b34c77b27 100644
--- a/twrpAdbBuFifo.hpp
+++ b/twrpAdbBuFifo.hpp
@@ -31,7 +31,7 @@ class twrpAdbBuFifo {
 	private:
 		bool start(void);
 		bool Backup_ADB_Command(std::string Options);
-		bool Check_Adb_Fifo_For_Events(void);
+		void Check_Adb_Fifo_For_Events(void);
 		bool Restore_ADB_Backup(void);
 		typedef bool (twrpAdbBuFifo::*ThreadPtr)(void);
 		typedef void* (*PThreadPtr)(void *);
diff --git a/twrpTar.cpp b/twrpTar.cpp
index aa00044d4..d15eea6b9 100644
--- a/twrpTar.cpp
+++ b/twrpTar.cpp
@@ -1276,7 +1276,7 @@ int twrpTar::openTar() {
 	} else if (current_archive_type == COMPRESSED) {
 		int pigzfd[2];
 
-		LOGINFO("Opening as a gzip...\n");
+		LOGINFO("Opening gzip compressed tar...\n");
 		if (part_settings->adbbackup)  {
 			LOGINFO("opening TW_ADB_RESTORE compressed stream\n");
 			input_fd = open(TW_ADB_RESTORE, O_RDONLY | O_LARGEFILE);