[CRIU] [PATCH 2/3] p.haul: workaround bug in tarfile module

[Nikita Spiridonov]

tarfile module always align data on source side according to
RECORDSIZE constant (10240 bytes), but it don't read aligning bytes
on target side in some cases depending on received buffer size. Read
aligning manually and discard. This is needed since tarfile share
socket with criu and it must stay in consistent state.

Signed-off-by: Nikita Spiridonov <nspiridonov at virtuozzo.com>
---
 phaul/images.py |    6 +++---
 phaul/util.py   |   29 ++++++++++++++++-------------
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/phaul/images.py b/phaul/images.py
index 1b004d6..656ab7b 100644
--- a/phaul/images.py
+++ b/phaul/images.py
@@ -41,11 +41,11 @@ class untar_thread(threading.Thread):
 
 	def run(self):
 		try:
-			tf = tarfile.open(mode="r|", fileobj=util.fileobj_wrap(self.__sk))
+			sk_wrap = util.fileobj_wrap(self.__sk)
+			tf = tarfile.open(mode="r|", fileobj=sk_wrap)
 			tf.extractall(self.__dir)
 			tf.close()
-			# Discard all remaining data in socket after tarfile closed
-			util.discard_sk_input(self.__sk)
+			sk_wrap.discard_unread_input()
 		except:
 			logging.exception("Exception in untar_thread")
 
diff --git a/phaul/util.py b/phaul/util.py
index 20954ea..420169f 100644
--- a/phaul/util.py
+++ b/phaul/util.py
@@ -3,6 +3,7 @@ import fcntl
 import errno
 import logging
 import socket
+import tarfile
 
 
 class fileobj_wrap:
@@ -16,28 +17,30 @@ class fileobj_wrap:
 
 	def __init__(self, sk):
 		self.__sk = sk
+		self.__nread = 0
 
 	def read(self, size=0x10000):
-		return self.__sk.recv(size)
+		data = self.__sk.recv(size)
+		self.__nread += len(data)
+		return data
 
 	def write(self, data):
 		self.__sk.sendall(data)
 		return len(data)
 
+	def discard_unread_input(self):
+		"""
+		Cleanup socket after tarfile
 
-def discard_sk_input(sk):
-	"""Read all data from socket and discard it
+		tarfile module always align data on source side according to RECORDSIZE
+		constant, but it don't read aligning bytes on target side in some cases
+		depending on received buffer size. Read aligning manually and discard.
+		"""
 
-	Current helper function needed to workaround tarfile library bug that
-	leads to ownerless garbage zero blocks in socket when tarfile constructed
-	with socket as file object to transfer tarballs over network.
-	"""
-	try:
-		while True:
-			sk.recv(0x10000, socket.MSG_DONTWAIT)
-	except socket.error as e:
-		if e.errno != errno.EWOULDBLOCK:
-			raise e
+		remainder = self.__nread % tarfile.RECORDSIZE
+		if remainder > 0:
+			self.__sk.recv(tarfile.RECORDSIZE - remainder, socket.MSG_WAITALL)
+			self.__nread += tarfile.RECORDSIZE - remainder
 
 
 class net_dev:
-- 
1.7.1