safe recv-key in direct mode

Checks the key's size and checksum. This is sorta expensive, but it avoids
needing to add another round-trip to the protocol.

1 files changed, 30 insertions, 1 deletions

diff --git a/Command/RecvKey.hs b/Command/RecvKey.hs
index 62ab9a7ee..11a5fd5ca 100644
--- a/Command/RecvKey.hs
+++ b/Command/RecvKey.hs
@@ -14,6 +14,10 @@ import Annex.Content
 import Utility.Rsync
 import Logs.Transfer
 import Command.SendKey (fieldTransfer)
+import qualified Fields
+import qualified Types.Key
+import qualified Types.Backend
+import qualified Backend
 
 def :: [Command]
 def = [noCommit $ command "recvkey" paramKey seek
@@ -26,7 +30,7 @@ start :: Key -> CommandStart
 start key = ifM (inAnnex key)
 	( error "key is already present in annex"
 	, fieldTransfer Download key $ \_p -> do
-		ifM (getViaTmp key $ liftIO . rsyncServerReceive)
+		ifM (getViaTmp key go)
 			( do
 				-- forcibly quit after receiving one key,
 				-- and shutdown cleanly
@@ -35,3 +39,28 @@ start key = ifM (inAnnex key)
 			, return False
 			)
 	)
+  where
+	go tmp = ifM (liftIO $ rsyncServerReceive tmp)
+		( ifM (isJust <$> Fields.getField Fields.direct)
+			( directcheck tmp
+			, return True
+			)
+		, return False
+		)
+	{- If the sending repository uses direct mode, the file
+	 - it sends could be modified as it's sending it. So check
+	 - that the right size file was received, and that the key/value
+	 - Backend is happy with it. -}
+	directcheck tmp = do
+		oksize <- case Types.Key.keySize key of
+		        Nothing -> return True
+		        Just size -> do
+				size' <- fromIntegral . fileSize
+       	        	        	<$> liftIO (getFileStatus tmp)
+				return $ size == size'
+		if oksize
+			then case Backend.maybeLookupBackendName (Types.Key.keyBackendName key) of
+				Nothing -> return False
+				Just backend -> maybe (return True) (\a -> a key tmp)
+					(Types.Backend.fsckKey backend)
+			else return False