diff --git a/Remote/Helper/Chunked.hs b/Remote/Helper/Chunked.hs
index b7522aa89c..00c089e802 100644
--- a/Remote/Helper/Chunked.hs
+++ b/Remote/Helper/Chunked.hs
@@ -221,7 +221,7 @@ retrieveChunks
-> Key
-> FilePath
-> MeterUpdate
- -> (Handle -> Maybe MeterUpdate -> L.ByteString -> IO ())
+ -> (Maybe Handle -> Maybe MeterUpdate -> ContentSource -> Annex Bool)
-> Annex Bool
retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
| noChunks chunkconfig =
@@ -244,34 +244,37 @@ retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
firstavail _ [] = return False
firstavail currsize ([]:ls) = firstavail currsize ls
- firstavail currsize ((k:ks):ls) = do
- let offset = resumeOffset currsize k
- let p = maybe basep
- (offsetMeterUpdate basep . toBytesProcessed)
- offset
- v <- tryNonAsyncAnnex $
- retriever (encryptor k) p $ \content ->
- bracketIO (maybe opennew openresume offset) hClose $ \h -> do
- tosink h p content
- let sz = toBytesProcessed $
- fromMaybe 0 $ keyChunkSize k
- getrest p h sz sz ks
- `catchNonAsyncAnnex` giveup
- case v of
- Left e
- | null ls -> giveup e
- | otherwise -> firstavail currsize ls
- Right r -> return r
+ firstavail currsize ((k:ks):ls)
+ | k == basek = getunchunked
+ `catchNonAsyncAnnex` (const $ firstavail currsize ls)
+ | otherwise = do
+ let offset = resumeOffset currsize k
+ let p = maybe basep
+ (offsetMeterUpdate basep . toBytesProcessed)
+ offset
+ v <- tryNonAsyncAnnex $
+ retriever (encryptor k) p $ \content ->
+ bracketIO (maybe opennew openresume offset) hClose $ \h -> do
+ void $ tosink (Just h) p content
+ let sz = toBytesProcessed $
+ fromMaybe 0 $ keyChunkSize k
+ getrest p h sz sz ks
+ `catchNonAsyncAnnex` giveup
+ case v of
+ Left e
+ | null ls -> giveup e
+ | otherwise -> firstavail currsize ls
+ Right r -> return r
getrest _ _ _ _ [] = return True
getrest p h sz bytesprocessed (k:ks) = do
let p' = offsetMeterUpdate p bytesprocessed
- retriever (encryptor k) p' $ tosink h p'
- getrest p h sz (addBytesProcessed bytesprocessed sz) ks
+ ifM (retriever (encryptor k) p' $ tosink (Just h) p')
+ ( getrest p h sz (addBytesProcessed bytesprocessed sz) ks
+ , giveup "chunk retrieval failed"
+ )
- getunchunked = bracketIO opennew hClose $ \h -> do
- retriever (encryptor basek) basep $ tosink h basep
- return True
+ getunchunked = retriever (encryptor basek) basep $ tosink Nothing basep
opennew = openBinaryFile dest WriteMode
@@ -290,13 +293,11 @@ retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
- it is not responsible for updating progress (often it cannot).
- Instead, the sink is passed a meter to update as it consumes
- the ByteString. -}
- tosink h p (ByteContent b) = liftIO $ do
- sink h (Just p) b
- return True
- tosink h _ (FileContent f) = liftIO $ do
- sink h Nothing =<< L.readFile f
- nukeFile f
- return True
+ tosink h p content = sink h p' content
+ where
+ p'
+ | isByteContent content = Just p
+ | otherwise = Nothing
{- Can resume when the chunk's offset is at or before the end of
- the dest file. -}
diff --git a/Remote/Helper/ChunkedEncryptable.hs b/Remote/Helper/ChunkedEncryptable.hs
index 550a6934bd..ca73802bae 100644
--- a/Remote/Helper/ChunkedEncryptable.hs
+++ b/Remote/Helper/ChunkedEncryptable.hs
@@ -18,9 +18,9 @@ module Remote.Helper.ChunkedEncryptable (
byteStorer,
fileRetriever,
byteRetriever,
- chunkedEncryptableRemote,
storeKeyDummy,
retreiveKeyFileDummy,
+ chunkedEncryptableRemote,
module X
) where
@@ -36,6 +36,7 @@ import Annex.Content
import Annex.Exception
import qualified Data.ByteString.Lazy as L
+import Control.Exception (bracket)
-- Use when nothing needs to be done to prepare a helper.
simplyPrepare :: helper -> Preparer helper
@@ -78,6 +79,16 @@ withBytes :: ContentSource -> (L.ByteString -> Annex a) -> Annex a
withBytes (ByteContent b) a = a b
withBytes (FileContent f) a = a =<< liftIO (L.readFile f)
+{- The base Remote that is provided to chunkedEncryptableRemote
+ - needs to have storeKey and retreiveKeyFile methods, but they are
+ - never actually used (since chunkedEncryptableRemote replaces
+ - them). Here are some dummy ones.
+ -}
+storeKeyDummy :: Key -> AssociatedFile -> MeterUpdate -> Annex Bool
+storeKeyDummy _ _ _ = return False
+retreiveKeyFileDummy :: Key -> AssociatedFile -> FilePath -> MeterUpdate -> Annex Bool
+retreiveKeyFileDummy _ _ _ _ = return False
+
-- Modifies a base Remote to support both chunking and encryption.
chunkedEncryptableRemote
:: RemoteConfig
@@ -131,17 +142,8 @@ chunkedEncryptableRemote c preparestorer prepareretriever baser = encr
where
go (Just retriever) = metered (Just p) k $ \p' ->
retrieveChunks retriever (uuid baser) chunkconfig
- enck k dest p' sink
+ enck k dest p' (sink dest enc)
go Nothing = return False
- sink h mp b = do
- let write = case mp of
- Just p' -> meteredWrite p' h
- Nothing -> L.hPut h
- case enc of
- Nothing -> write b
- Just (cipher, _) ->
- decrypt cipher (feedBytes b) $
- readBytes write
enck = maybe id snd enc
removeKeyGen k enc = removeChunks remover (uuid baser) chunkconfig enck k
@@ -154,12 +156,39 @@ chunkedEncryptableRemote c preparestorer prepareretriever baser = encr
enck = maybe id snd enc
checker = hasKey baser
-{- The base Remote that is provided to chunkedEncryptableRemote
- - needs to have storeKey and retreiveKeyFile methods, but they are
- - never actually used (since chunkedEncryptableRemote replaces
- - them). Here are some dummy ones.
+{- Sink callback for retrieveChunks. Stores the file content into the
+ - provided Handle, decrypting it first if necessary.
+ -
+ - If the remote did not store the content using chunks, no Handle
+ - will be provided, and it's up to us to open the destination file.
+ -
+ - Note that when neither chunking nor encryption is used, and the remote
+ - provides FileContent, that file only needs to be renamed
+ - into place. (And it may even already be in the right place..)
-}
-storeKeyDummy :: Key -> AssociatedFile -> MeterUpdate -> Annex Bool
-storeKeyDummy _ _ _ = return False
-retreiveKeyFileDummy :: Key -> AssociatedFile -> FilePath -> MeterUpdate -> Annex Bool
-retreiveKeyFileDummy _ _ _ _ = return False
+sink
+ :: FilePath
+ -> Maybe (Cipher, EncKey)
+ -> Maybe Handle
+ -> Maybe MeterUpdate
+ -> ContentSource
+ -> Annex Bool
+sink dest enc mh mp content = do
+ case (enc, mh, content) of
+ (Nothing, Nothing, FileContent f)
+ | f == dest -> noop
+ | otherwise -> liftIO $ moveFile f dest
+ (Just (cipher, _), _, _) ->
+ withBytes content $ \b ->
+ decrypt cipher (feedBytes b) $
+ readBytes write
+ (Nothing, _, _) -> withBytes content write
+ return True
+ where
+ write b = case mh of
+ Just h -> liftIO $ b `streamto` h
+ Nothing -> liftIO $ bracket opendest hClose (b `streamto`)
+ streamto b h = case mp of
+ Just p -> meteredWrite p h b
+ Nothing -> L.hPut h b
+ opendest = openBinaryFile dest WriteMode
diff --git a/Types/StoreRetrieve.hs b/Types/StoreRetrieve.hs
index 0ee2fd5019..33f66efb19 100644
--- a/Types/StoreRetrieve.hs
+++ b/Types/StoreRetrieve.hs
@@ -23,6 +23,10 @@ data ContentSource
= FileContent FilePath
| ByteContent L.ByteString
+isByteContent :: ContentSource -> Bool
+isByteContent (ByteContent _) = True
+isByteContent (FileContent _) = False
+
-- Action that stores a Key's content on a remote.
-- Can throw exceptions.
type Storer = Key -> ContentSource -> MeterUpdate -> Annex Bool