refactor

2 files changed, 85 insertions, 64 deletions

diff --git a/Remote/Directory.hs b/Remote/Directory.hs
index 5ea0a1e6b..2d31d12b2 100644
--- a/Remote/Directory.hs
+++ b/Remote/Directory.hs
@@ -42,17 +42,20 @@ gen :: Git.Repo -> UUID -> Maybe RemoteConfig -> Annex (Remote Annex)
 gen r u c = do
 	dir <- getConfig r "directory" (error "missing directory")
 	cst <- remoteCost r cheapRemoteCost
-	return $ Remote {
-		uuid = u,
-		cost = cst,
-		name = Git.repoDescribe r,
- 		storeKey = storeKeyEncrypted c $ store dir,
-		retrieveKeyFile = retrieveKeyFileEncrypted c $ retrieve dir,
-		removeKey = removeKeyEncrypted c $ remove dir,
-		hasKey = hasKeyEncrypted c $ checkPresent dir,
-		hasKeyCheap = True,
-		config = Nothing
-	}
+	return $ encryptedRemote c
+		(storeEncrypted dir)
+		(retrieveEncrypted dir)
+		Remote {
+			uuid = u,
+			cost = cst,
+			name = Git.repoDescribe r,
+ 			storeKey = store dir,
+			retrieveKeyFile = retrieve dir,
+			removeKey = remove dir,
+			hasKey = checkPresent dir,
+			hasKeyCheap = True,
+			config = Nothing
+		}
 
 directorySetup :: UUID -> RemoteConfig -> Annex RemoteConfig
 directorySetup u c = do
@@ -74,43 +77,47 @@ dirKey d k = d </> hashDirMixed k </> f </> f
 	where
 		f = keyFile k
 
-store :: FilePath -> Key -> Maybe (Cipher, Key) -> Annex Bool
-store d k c = do
+store :: FilePath -> Key -> Annex Bool
+store d k = do
 	g <- Annex.gitRepo
-	let src = gitAnnexLocation g k	
-	liftIO $ catch (copy src) (const $ return False)
+	let src = gitAnnexLocation g k
+	let dest = dirKey d k
+	liftIO $ catch (storeHelper dest $ copyFile src dest) (const $ return False)
+
+storeEncrypted :: FilePath -> (Cipher, Key) -> Key -> Annex Bool
+storeEncrypted d (cipher, enck) k = do
+	g <- Annex.gitRepo
+	let src = gitAnnexLocation g k
+	let dest = dirKey d enck
+	liftIO $ catch (storeHelper dest $ encrypt src dest) (const $ return False)
 	where
-		copy src = case c of
-			Just (cipher, enckey) -> do
-				content <- L.readFile src
-				let dest = dirKey d enckey
-				prep dest
-				withEncryptedContent cipher content $ \s -> do
-					L.writeFile dest s
-					cleanup True dest
-			_ -> do
-				let dest = dirKey d k
-				prep dest
-				ok <- copyFile src dest
-				cleanup ok dest
-		prep dest = liftIO $ do
-			let dir = parentDir dest
-			createDirectoryIfMissing True dir
-			allowWrite dir
-		cleanup ok dest = do
-			when ok $ do
-				let dir = parentDir dest
-				preventWrite dest
-				preventWrite dir
-			return ok
+		encrypt src dest = do
+			content <- L.readFile src
+			withEncryptedContent cipher content $ L.writeFile dest
+			return True
 
-retrieve :: FilePath -> Key -> FilePath -> Maybe (Cipher, Key) -> Annex Bool
-retrieve d k f Nothing = liftIO $ copyFile (dirKey d k) f
-retrieve d k f (Just (cipher, enckey)) = 
-	liftIO $ flip catch (const $ return False) $ do
-		content <- L.readFile (dirKey d enckey)
-		withDecryptedContent cipher content $ L.writeFile f
-		return True
+storeHelper :: FilePath -> IO Bool -> IO Bool
+storeHelper dest a = do
+	let dir = parentDir dest
+	createDirectoryIfMissing True dir
+	allowWrite dir	
+	ok <- a
+	when ok $ do
+		preventWrite dest
+		preventWrite dir
+	return ok
+
+retrieve :: FilePath -> Key -> FilePath -> Annex Bool
+retrieve d k f = liftIO $ copyFile (dirKey d k) f
+
+retrieveEncrypted :: FilePath -> (Cipher, Key) -> FilePath -> Annex Bool
+retrieveEncrypted d (cipher, enck) f =
+	liftIO $ catch decrypt (const $ return False)
+	where
+		decrypt = do
+			content <- L.readFile (dirKey d enck)
+			withDecryptedContent cipher content $ L.writeFile f
+			return True
 
 remove :: FilePath -> Key -> Annex Bool
 remove d k = liftIO $ catch del (const $ return False)
diff --git a/Remote/Encrypted.hs b/Remote/Encrypted.hs
index 0ff2833b3..255b41d73 100644
--- a/Remote/Encrypted.hs
+++ b/Remote/Encrypted.hs
@@ -33,16 +33,39 @@ encryptionSetup c =
 			cipher <- liftIO a
 			return $ M.delete "encryption" $ storeCipher c cipher
 
-{- Helpers that can be applied to a Remote's normal actions to
- - add crypto support. -}
-storeKeyEncrypted :: Maybe RemoteConfig -> (Key -> Maybe (Cipher, Key) -> Annex a) -> Key -> Annex a
-storeKeyEncrypted c a k = a k =<< cipherKey c k
-retrieveKeyFileEncrypted :: Maybe RemoteConfig -> (Key -> FilePath -> Maybe (Cipher, Key) -> Annex a) -> Key -> FilePath -> Annex a
-retrieveKeyFileEncrypted c a k f = a k f =<< cipherKey c k
-removeKeyEncrypted :: Maybe RemoteConfig -> (Key -> Annex a) -> Key -> Annex a
-removeKeyEncrypted = withEncryptedKey
-hasKeyEncrypted :: Maybe RemoteConfig -> (Key -> Annex a) -> Key -> Annex a
-hasKeyEncrypted = withEncryptedKey
+{- Modifies a Remote to support encryption.
+ -
+ - Two additional functions must be provided by the remote,
+ - to support storing and retrieving encrypted content. -}
+encryptedRemote
+	:: Maybe RemoteConfig
+	-> ((Cipher, Key) -> Key -> Annex Bool)
+	-> ((Cipher, Key) -> FilePath -> Annex Bool)
+	-> Remote Annex 
+	-> Remote Annex
+encryptedRemote c storeKeyEncrypted retrieveKeyFileEncrypted r = 
+	r {
+		storeKey = store,
+		retrieveKeyFile = retrieve,
+		removeKey = withkey $ removeKey r,
+		hasKey = withkey $ hasKey r
+	}
+	where
+		store k = do
+			v <- cipherKey c k
+			case v of
+				Nothing -> (storeKey r) k
+				Just x -> storeKeyEncrypted x k
+		retrieve k f = do
+			v <- cipherKey c k
+			case v of
+				Nothing -> (retrieveKeyFile r) k f
+				Just x -> retrieveKeyFileEncrypted x f
+		withkey a k = do
+			v <- cipherKey c k
+			case v of
+				Nothing -> a k
+				Just (_, k') -> a k'
 
 {- Gets encryption Cipher, and encrypted version of Key. 
  -
@@ -64,12 +87,3 @@ cipherKey (Just c) k = do
 		ret cipher = do
 			k' <- liftIO $ encryptKey cipher k
 			return $ Just (cipher, k')
-
-{- Passes the encrypted version of the key to the action when encryption
- - is enabled, and the non-encrypted version otherwise. -}
-withEncryptedKey :: Maybe RemoteConfig -> (Key -> Annex a) -> Key -> Annex a
-withEncryptedKey c a k = do
-	v <- cipherKey c k
-	case v of
-		Nothing -> a k
-		Just (_, k') -> a k'