add ContentSource type, for remotes that act on files rather than ByteStrings

Note that currently nothing cleans up a ContentSource's file, when eg,
retrieving chunks.

5 files changed, 102 insertions, 49 deletions

diff --git a/Remote/Directory.hs b/Remote/Directory.hs
index b107c18e9..5d8a040d4 100644
--- a/Remote/Directory.hs
+++ b/Remote/Directory.hs
@@ -110,9 +110,9 @@ tmpDir d k = addTrailingPathSeparator $ d </> "tmp" </> keyFile k
 prepareStore :: FilePath -> ChunkConfig -> Preparer Storer
 prepareStore d chunkconfig = checkPrepare
 	(\k -> checkDiskSpace (Just d) k 0)
-	(store d chunkconfig)
+	(byteStorer $ store d chunkconfig)
 
-store :: FilePath -> ChunkConfig -> Storer
+store :: FilePath -> ChunkConfig -> Key -> L.ByteString -> MeterUpdate -> IO Bool
 store d chunkconfig k b p = do
 	void $ tryIO $ createDirectoryIfMissing True tmpdir
 	case chunkconfig of
@@ -137,7 +137,8 @@ store d chunkconfig k b p = do
 
 retrieve :: FilePath -> ChunkConfig -> Preparer Retriever
 retrieve d (LegacyChunks _) = Legacy.retrieve locations d
-retrieve d _ = simplyPrepare $ \k -> L.readFile =<< getLocation d k
+retrieve d _ = simplyPrepare $ byteRetriever $ 
+	\k -> L.readFile =<< getLocation d k
 
 retrieveCheap :: FilePath -> ChunkConfig -> Key -> FilePath -> Annex Bool
 -- no cheap retrieval possible for chunks
diff --git a/Remote/Directory/LegacyChunked.hs b/Remote/Directory/LegacyChunked.hs
index c7b8ad52c..af846a2e6 100644
--- a/Remote/Directory/LegacyChunked.hs
+++ b/Remote/Directory/LegacyChunked.hs
@@ -96,7 +96,7 @@ retrieve locations d basek a = do
 	tmpdir <- fromRepo $ gitAnnexTmpMiscDir
 	createAnnexDirectory tmpdir
 	let tmp = tmpdir </> keyFile basek ++ ".directorylegacy.tmp"
-	a $ Just $ \k -> do
+	a $ Just $ byteRetriever $ \k -> do
 		void $ withStoredFiles d locations k $ \fs -> do
 			forM_ fs $
 				S.appendFile tmp <=< S.readFile
diff --git a/Remote/Helper/Chunked.hs b/Remote/Helper/Chunked.hs
index 708d87800..70e541cce 100644
--- a/Remote/Helper/Chunked.hs
+++ b/Remote/Helper/Chunked.hs
@@ -17,6 +17,7 @@ module Remote.Helper.Chunked (
 
 import Common.Annex
 import Utility.DataUnits
+import Types.StoreRetrieve
 import Types.Remote
 import Types.Key
 import Logs.Chunk
@@ -90,29 +91,31 @@ storeChunks
 	-> Key
 	-> FilePath
 	-> MeterUpdate
-	-> (Key -> L.ByteString -> MeterUpdate -> IO Bool)
+	-> (Key -> ContentSource -> MeterUpdate -> IO Bool)
 	-> (Key -> Annex (Either String Bool))
 	-> Annex Bool
-storeChunks u chunkconfig k f p storer checker = bracketIO open close go
+storeChunks u chunkconfig k f p storer checker = 
+	case chunkconfig of
+		(UnpaddedChunks chunksize) -> 
+			bracketIO open close (go chunksize)
+		_ -> showprogress $
+			liftIO . storer k (FileContent f)
   where
+	showprogress = metered (Just p) k
+
 	open = tryIO $ openBinaryFile f ReadMode
 
 	close (Right h) = hClose h
 	close (Left _) = noop
 
-	go (Left e) = do
+	go _ (Left e) = do
 		warning (show e)
 		return False
-	go (Right h) = metered (Just p) k $ \meterupdate ->
-		case chunkconfig of
-			(UnpaddedChunks chunksize) -> do
-				let chunkkeys = chunkKeyStream k chunksize
-				(chunkkeys', startpos) <- seekResume h chunkkeys checker
-				b <- liftIO $ L.hGetContents h
-				gochunks meterupdate startpos chunksize b chunkkeys'
-			_ -> liftIO $ do
-				b <- L.hGetContents h
-				storer k b meterupdate
+	go chunksize (Right h) = showprogress $ \meterupdate -> do
+		let chunkkeys = chunkKeyStream k chunksize
+		(chunkkeys', startpos) <- seekResume h chunkkeys checker
+		b <- liftIO $ L.hGetContents h
+		gochunks meterupdate startpos chunksize b chunkkeys'
 
 	gochunks :: MeterUpdate -> BytesProcessed -> ChunkSize -> L.ByteString -> ChunkKeyStream -> Annex Bool
 	gochunks meterupdate startpos chunksize = loop startpos . splitchunk
@@ -127,7 +130,7 @@ storeChunks u chunkconfig k f p storer checker = bracketIO open close go
 				return True
 			| otherwise = do
 				let (chunkkey, chunkkeys') = nextChunkKeyStream chunkkeys
-				ifM (liftIO $ storer chunkkey chunk meterupdate')
+				ifM (liftIO $ storer chunkkey (ByteContent chunk) meterupdate')
 					( do
 						let bytesprocessed' = addBytesProcessed bytesprocessed (L.length chunk)
 						loop bytesprocessed' (splitchunk bs) chunkkeys'
@@ -197,8 +200,7 @@ removeChunks remover u chunkconfig encryptor k = do
 		forM_ chunksizes $ chunksRemoved u k . FixedSizeChunks . fromIntegral
 	return ok
 
-{- Retrieves a key from a remote, using a retriever action that
- - streams it to a ByteString.
+{- Retrieves a key from a remote, using a retriever action.
  -
  - When the remote is chunked, tries each of the options returned by
  - chunkKeys until it finds one where the retriever successfully
@@ -214,7 +216,7 @@ removeChunks remover u chunkconfig encryptor k = do
  - to resume.
  -}
 retrieveChunks 
-	:: (Key -> IO L.ByteString)
+	:: Retriever
 	-> UUID
 	-> ChunkConfig
 	-> EncKey
@@ -250,13 +252,13 @@ retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
 			Left e
 				| null ls -> giveup e
 				| otherwise -> firstavail currsize ls
-			Right b -> do
+			Right content -> do
 				let offset = resumeOffset currsize k
 				let p = maybe basep
 					(offsetMeterUpdate basep . toBytesProcessed)
 					offset
 				bracket (maybe opennew openresume offset) hClose $ \h -> do
-					sink h p b
+					withBytes content $ sink h p
 					let sz = toBytesProcessed $
 						fromMaybe 0 $ keyChunkSize k
 					getrest p h sz sz ks
@@ -264,11 +266,13 @@ retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
 	getrest _ _ _ _ [] = return True
 	getrest p h sz bytesprocessed (k:ks) = do
 		let p' = offsetMeterUpdate p bytesprocessed
-		sink h p' =<< retriever (encryptor k)
+		content <- retriever (encryptor k)
+		withBytes content $ sink h p'
 		getrest p h sz (addBytesProcessed bytesprocessed sz) ks
 
 	getunchunked = liftIO $ bracket opennew hClose $ \h -> do
-		retriever (encryptor basek) >>= sink h basep
+		content <- retriever (encryptor basek)
+		withBytes content $ sink h basep
 		return True
 
 	opennew = openBinaryFile dest WriteMode
diff --git a/Remote/Helper/ChunkedEncryptable.hs b/Remote/Helper/ChunkedEncryptable.hs
index 402a64891..b851ecd94 100644
--- a/Remote/Helper/ChunkedEncryptable.hs
+++ b/Remote/Helper/ChunkedEncryptable.hs
@@ -5,23 +5,28 @@
  - Licensed under the GNU GPL version 3 or higher.
  -}
 
-{-# LANGUAGE Rank2Types #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ImpredicativeTypes #-}
+{-# LANGUAGE FlexibleContexts #-}
 
 module Remote.Helper.ChunkedEncryptable (
 	Preparer,
-	simplyPrepare,
-	checkPrepare,
 	Storer,
 	Retriever,
+	simplyPrepare,
+	checkPrepare,
+	fileStorer,
+	byteStorer,
+	fileRetriever,
+	byteRetriever,
 	chunkedEncryptableRemote,
 	storeKeyDummy,
 	retreiveKeyFileDummy,
 	module X
 ) where
 
-import qualified Data.ByteString.Lazy as L
-
 import Common.Annex
+import Types.StoreRetrieve
 import Types.Remote
 import Crypto
 import Config.Cost
@@ -31,10 +36,6 @@ import Remote.Helper.Encryptable as X
 import Annex.Content
 import Annex.Exception
 
--- Prepares for and then runs an action that will act on a Key,
--- passing it a helper when the preparation is successful.
-type Preparer helper = forall a. Key -> (Maybe helper -> Annex a) -> Annex a
-
 simplyPrepare :: helper -> Preparer helper
 simplyPrepare helper _ a = a $ Just helper
 
@@ -44,14 +45,6 @@ checkPrepare checker helper k a = ifM (checker k)
 	, a Nothing
 	)
 
--- Stores a Key, which may be encrypted and/or a chunk key.
--- May throw exceptions.
-type Storer = Key -> L.ByteString -> MeterUpdate -> IO Bool
-
--- Retrieves a Key, which may be encrypted and/or a chunk key.
--- Throws exception if key is not present, or remote is not accessible.
-type Retriever = Key -> IO L.ByteString
-
 {- Modifies a base Remote to support both chunking and encryption.
  -}
 chunkedEncryptableRemote
@@ -88,16 +81,17 @@ chunkedEncryptableRemote c preparestorer prepareretriever baser = encr
 		go (Just storer) = sendAnnex k rollback $ \src ->
 			metered (Just p) k $ \p' ->
 				storeChunks (uuid baser) chunkconfig k src p'
-					(storechunk storer)
+					(storechunk enc storer)
 					(hasKey baser)
 		go Nothing = return False
 		rollback = void $ removeKey encr k
-		storechunk storer k' b p' = case enc of
-			Nothing -> storer k' b p'
-			Just (cipher, enck) -> 
-				encrypt gpgopts cipher (feedBytes b) $
-					readBytes $ \encb ->
-						storer (enck k') encb p'
+
+	storechunk Nothing storer k content p = storer k content p
+	storechunk (Just (cipher, enck)) storer k content p =
+		withBytes content $ \b ->
+			encrypt gpgopts cipher (feedBytes b) $
+				readBytes $ \encb ->
+					storer (enck k) (ByteContent encb) p
 
 	-- call retriever to get chunks; decrypt them; stream to dest file
 	retrieveKeyFileGen k dest p enc =
diff --git a/Types/StoreRetrieve.hs b/Types/StoreRetrieve.hs
new file mode 100644
index 000000000..2520d6309
--- /dev/null
+++ b/Types/StoreRetrieve.hs
@@ -0,0 +1,54 @@
+{- Types for Storer and Retriever
+ -
+ - Copyright 2014 Joey Hess <joey@kitenet.net>
+ -
+ - Licensed under the GNU GPL version 3 or higher.
+ -}
+
+{-# LANGUAGE Rank2Types #-}
+
+module Types.StoreRetrieve where
+
+import Common.Annex
+import Utility.Metered
+import Utility.Tmp
+
+import qualified Data.ByteString.Lazy as L
+
+-- Prepares for and then runs an action that will act on a Key's
+-- content, passing it a helper when the preparation is successful.
+type Preparer helper = forall a. Key -> (Maybe helper -> Annex a) -> Annex a
+
+-- A source of a Key's content.
+data ContentSource
+	= FileContent FilePath
+	| ByteContent L.ByteString
+
+-- Action that stores a Key's content on a remote.
+-- Can throw exceptions.
+type Storer = Key -> ContentSource -> MeterUpdate -> IO Bool
+
+-- Action that retrieves a Key's content from a remote.
+-- Throws exception if key is not present, or remote is not accessible.
+type Retriever = Key -> IO ContentSource
+
+fileStorer :: (Key -> FilePath -> MeterUpdate -> IO Bool) -> Storer
+fileStorer a k (FileContent f) m = a k f m
+fileStorer a k (ByteContent b) m = do
+	 withTmpFile "tmpXXXXXX" $ \f h -> do
+		L.hPut h b
+		hClose h
+		a k f m
+
+byteStorer :: (Key -> L.ByteString -> MeterUpdate -> IO Bool) -> Storer
+byteStorer a k c m = withBytes c $ \b -> a k b m
+
+withBytes :: ContentSource -> (L.ByteString -> IO a) -> IO a
+withBytes (ByteContent b) a = a b
+withBytes (FileContent f) a = a =<< L.readFile f
+
+fileRetriever :: (Key -> IO FilePath) -> Retriever
+fileRetriever a k = FileContent <$> a k
+
+byteRetriever :: (Key -> IO L.ByteString) -> Retriever
+byteRetriever a k = ByteContent <$> a k