pull/push over tor working now

Still a couple bugs:

* Closing the connection to the server leaves git upload-pack /
  receive-pack running, which could be used to DOS.

* Sometimes the data is transferred, but it fails at the end, sometimes
  with:

  git-remote-tor-annex: <socket: 10>: commitBuffer: resource vanished (Broken pipe)

  Must be a race condition around shutdown.

2 files changed, 109 insertions, 113 deletions

diff --git a/Remote/Helper/P2P.hs b/Remote/Helper/P2P.hs
index 7e49968ee..0e604a50d 100644
--- a/Remote/Helper/P2P.hs
+++ b/Remote/Helper/P2P.hs
@@ -141,30 +141,18 @@ data NetF c
 	| SendBytes Len L.ByteString c
 	| ReceiveBytes Len (L.ByteString -> c)
 	| CheckAuthToken UUID AuthToken (Bool -> c)
-	| Relay RelayHandle
-		(RelayData -> Net (Maybe ExitCode))
-		(ExitCode -> c)
-	-- ^ Waits for data to be written to the RelayHandle, and for messages
-	-- to be received from the peer, and passes the data to the
-	-- callback, continuing until it returns an ExitCode.
-	| RelayService Service
-		(RelayHandle -> RelayData -> Net (Maybe ExitCode))
-		(ExitCode -> c)
-	-- ^ Runs a service, and waits for it to output to stdout,
-	-- and for messages to be received from the peer, and passes
-	-- the data to the callback (which is also passed the service's
-	-- stdin RelayHandle), continuing uniil the service exits.
-	| WriteRelay RelayHandle L.ByteString c
-	-- ^ Write data to a relay's handle, flushing it immediately.
+	| RelayService Service c
+	-- ^ Runs a service, relays its output to the peer, and data
+	-- from the peer to it.
+	| Relay RelayHandle RelayHandle (ExitCode -> c)
+	-- ^ Reads from the first RelayHandle, and sends the data to a
+	-- peer, while at the same time accepting input from the peer
+	-- which is sent the the second RelayHandle. Continues until 
+	-- the peer sends an ExitCode.
 	deriving (Functor)
 
 type Net = Free NetF
 
-data RelayData
-	= RelayData L.ByteString
-	| RelayMessage Message
-	deriving (Show)
-
 newtype RelayHandle = RelayHandle Handle
 
 data LocalF c
@@ -212,8 +200,7 @@ runNet (SendBytes _ _ next) ms = ("> bytes", Nothing):runPure next ms
 runNet (ReceiveBytes _ next) ms = ("< bytes", Nothing):runPure (next L.empty) ms
 runNet (CheckAuthToken _ _ next) ms = runPure (next True) ms
 runNet (Relay _ _ next) ms = runPure (next ExitSuccess) ms
-runNet (RelayService _ _ next) ms = runPure (next ExitSuccess) ms
-runNet (WriteRelay _ _ next) ms = runPure next ms
+runNet (RelayService _ next) ms = runPure next ms
 
 runLocal :: Show r => LocalF (Proto r) -> [Message] -> [(String, Maybe Message)]
 runLocal (KeyFileSize _ next) ms = runPure (next (Len 100)) ms
@@ -341,9 +328,7 @@ serve myuuid = go Nothing
 		-- setPresent not called because the peer may have
 		-- requested the data but not permanatly stored it.
 		GET offset key -> void $ sendContent key offset
-		CONNECT service -> do
-			exitcode <- net $ relayService service relayCallback
-			net $ sendMessage (CONNECTDONE exitcode)
+		CONNECT service -> net $ relayService service
 		_ -> net $ sendMessage (ERROR "unexpected command")
 
 sendContent :: Key -> Offset -> Proto Bool
@@ -399,19 +384,28 @@ readKeyFileLen key (Offset offset) = do
 connect :: Service -> Handle -> Handle -> Proto ExitCode
 connect service hin hout = do
 	net $ sendMessage (CONNECT service)
-	net $ relay (RelayHandle hin) (relayCallback (RelayHandle hout))
-
-relayCallback :: RelayHandle -> RelayData -> Net (Maybe ExitCode)
-relayCallback hout (RelayMessage (DATA len)) = do
-	writeRelay hout =<< receiveBytes len
-	return Nothing
-relayCallback _ (RelayMessage (CONNECTDONE exitcode)) =
-	return (Just exitcode)
-relayCallback _ (RelayMessage m) = do
-	sendMessage $ ERROR $ "expected DATA or CONNECTDONE not " ++ unwords (Proto.formatMessage m)
-	return (Just (ExitFailure 1))
-relayCallback _ (RelayData b) = do
+	net $ relay (RelayHandle hin) (RelayHandle hout)
+
+data RelayData
+	= RelayToPeer L.ByteString
+	| RelayFromPeer L.ByteString
+	| RelayDone ExitCode
+	deriving (Show)
+
+relayFromPeer :: Net RelayData
+relayFromPeer = do
+	r <- receiveMessage
+	case r of
+		CONNECTDONE exitcode -> return $ RelayDone exitcode
+		DATA len -> RelayFromPeer <$> receiveBytes len
+		_ -> do
+			sendMessage $ ERROR "expected DATA or CONNECTDONE"
+			return $ RelayDone $ ExitFailure 1
+
+relayToPeer :: RelayData -> Net ()
+relayToPeer (RelayDone exitcode) = sendMessage (CONNECTDONE exitcode)
+relayToPeer (RelayToPeer b) = do
 	let len = Len $ fromIntegral $ L.length b
 	sendMessage (DATA len)
 	sendBytes len b
-	return Nothing
+relayToPeer (RelayFromPeer _) = return ()
diff --git a/Remote/Helper/P2P/IO.hs b/Remote/Helper/P2P/IO.hs
index c6a80cdbf..b3597abf9 100644
--- a/Remote/Helper/P2P/IO.hs
+++ b/Remote/Helper/P2P/IO.hs
@@ -14,7 +14,6 @@ module Remote.Helper.P2P.IO
 
 import Remote.Helper.P2P
 import Utility.Process
-import Types.UUID
 import Git
 import Git.Command
 import Utility.SafeCommand
@@ -24,10 +23,10 @@ import Utility.Exception
 import Control.Monad
 import Control.Monad.Free
 import Control.Monad.IO.Class
-import Data.Maybe
 import System.Exit (ExitCode(..))
 import System.IO
 import Control.Concurrent
+import Control.Concurrent.Async
 import qualified Data.ByteString as B
 import qualified Data.ByteString.Lazy as L
 
@@ -58,7 +57,6 @@ runNetHandle s runner f = case f of
 		runner next
 	ReceiveMessage next -> do
 		l <- liftIO $ hGetLine (ihdl s)
-		-- liftIO $ hPutStrLn stderr ("< " ++ show l)
 		case parseMessage l of
 			Just m -> runner (next m)
 			Nothing -> runner $ do
@@ -72,64 +70,43 @@ runNetHandle s runner f = case f of
 		runner next
 	ReceiveBytes (Len n) next -> do
 		b <- liftIO $ L.hGet (ihdl s) (fromIntegral n)
-		--liftIO $ hPutStrLn stderr $ "!!!" ++ show (L.length b)
 		runner (next b)
 	CheckAuthToken u t next -> do
 		authed <- return True -- TODO XXX FIXME really check
 		runner (next authed)
-	Relay hout callback next ->
-		runRelay runner hout callback >>= runner . next
-	RelayService service callback next ->
-		runRelayService s runner service callback >>= runner . next
-	WriteRelay (RelayHandle h) b next -> do
-		liftIO $ do
-			-- L.hPut h b
-			hPutStrLn h (show ("relay got:", b, L.length b))
-			hFlush h
-		runner next
+	Relay hin hout next ->
+		runRelay runner hin hout >>= runner . next
+	RelayService service next ->
+		runRelayService s runner service >> runner next
 
 runRelay
 	:: MonadIO m
 	=> RunProto
 	-> RelayHandle
-	-> (RelayData -> Net (Maybe ExitCode))
+	-> RelayHandle
 	-> m ExitCode
-runRelay runner (RelayHandle hout) callback = do
-	v <- liftIO newEmptyMVar
-	_ <- liftIO $ forkIO $ readout v
-	feeder <- liftIO $ forkIO $ feedin v
-	exitcode <- liftIO $ drain v
-	liftIO $ killThread feeder
-	return exitcode
+runRelay runner (RelayHandle hout) (RelayHandle hin) = liftIO $
+	bracket setup cleanup go
   where
-	feedin v = forever $ do
-		m <- runner $ net receiveMessage
-		putMVar v $ RelayMessage m
+	setup = do
+		v <- newEmptyMVar
+		void $ forkIO $ relayFeeder runner v
+		void $ forkIO $ relayReader v hout
+		return v
 	
-	readout v = do
-		b <- B.hGetSome hout 65536
-		if B.null b
-			then hClose hout
-			else do
-				putMVar v $ RelayData (L.fromChunks [b])
-				readout v
-
-	drain v = do
-		d <- takeMVar v
-		liftIO $ hPutStrLn stderr (show d)
-		r <- runner $ net $ callback d
-		case r of
-			Nothing -> drain v
-			Just exitcode -> return exitcode
+	cleanup _ = do
+		hClose hin
+		hClose hout
+	
+	go v = relayHelper runner v hin
 
 runRelayService
-	:: (MonadIO m, MonadMask m)
+	:: MonadIO m
 	=> S
 	-> RunProto
 	-> Service
-	-> (RelayHandle -> RelayData -> Net (Maybe ExitCode))
-	-> m ExitCode
-runRelayService s runner service callback = bracket setup cleanup go
+	-> m ()
+runRelayService s runner service = liftIO $ bracket setup cleanup go
   where
 	cmd = case service of
 		UploadPack -> "upload-pack"
@@ -141,45 +118,70 @@ runRelayService s runner service callback = bracket setup cleanup go
 		] (repo s)
 
 	setup = do
-		v <- liftIO newEmptyMVar
-		(Just hin, Just hout, _, pid) <- liftIO $ 
-			createProcess serviceproc
-				{ std_out = CreatePipe
-				, std_in = CreatePipe
-				}
-		feeder <- liftIO $ forkIO $ feedin v
-		return (v, feeder, hin, hout, pid)
-
-	cleanup (_, feeder, hin, hout, pid) = liftIO $ do
+		(Just hin, Just hout, _, pid) <- createProcess serviceproc
+			{ std_out = CreatePipe
+			, std_in = CreatePipe
+			}
+		v <- newEmptyMVar
+		feeder <- async $ relayFeeder runner v
+		reader <- async $ relayReader v hout
+		waiter <- async $ waitexit v pid
+		return (v, feeder, reader, waiter, hin, hout, pid)
+
+	cleanup (_, feeder, reader, waiter, hin, hout, pid) = do
+		hPutStrLn stderr "!!!!\n\nIN CLEANUP"
+		hFlush stderr
 		hClose hin
 		hClose hout
-		liftIO $ killThread feeder
+		cancel reader
+		cancel waiter
 		void $ waitForProcess pid
 
-	go (v, _, hin, hout, pid) = do
-		_ <- liftIO $ forkIO $ readout v hout
-		_ <- liftIO $ forkIO $ putMVar v . Left =<< waitForProcess pid
-		liftIO $ drain v hin
+	go (v, _, _, _, hin, _, _) = do
+		exitcode <- relayHelper runner v hin
+		runner $ net $ relayToPeer (RelayDone exitcode)
+	
+	waitexit v pid = putMVar v . RelayDone =<< waitForProcess pid
 
-	drain v hin = do
+-- Processes RelayData as it is put into the MVar.
+relayHelper :: RunProto -> MVar RelayData -> Handle -> IO ExitCode
+relayHelper runner v hin = loop
+  where
+	loop = do
 		d <- takeMVar v
 		case d of
-			Left exitcode -> return exitcode
-			Right relaydata -> do
-				liftIO $ hPutStrLn stderr ("> " ++ show relaydata)
-				_ <- runner $ net $
-					callback (RelayHandle hin) relaydata
-				drain v hin
-	
-	readout v hout = do
+			RelayFromPeer b -> do
+				L.hPut hin b
+				hFlush hin
+				loop
+			RelayToPeer b -> do
+				runner $ net $ relayToPeer (RelayToPeer b)
+				loop
+			RelayDone exitcode -> do
+				runner $ net $ relayToPeer (RelayDone exitcode)
+				return exitcode
+
+-- Takes input from the peer, and puts it into the MVar for processing.
+-- Repeats until the peer tells it it's done.
+relayFeeder :: RunProto -> MVar RelayData -> IO ()
+relayFeeder runner v = loop
+  where
+	loop = do
+		rd <- runner $ net relayFromPeer
+		putMVar v rd
+		case rd of
+			RelayDone _ -> return ()
+			_ -> loop
+
+-- Reads input from the Handle and puts it into the MVar for relaying to
+-- the peer. Continues until EOF on the Handle.
+relayReader :: MVar RelayData -> Handle -> IO ()
+relayReader v hout = loop
+  where
+	loop = do
 		b <- B.hGetSome hout 65536
 		if B.null b
 			then return ()
 			else do
-				putMVar v $ Right $ 
-					RelayData (L.fromChunks [b])
-				readout v hout
-
-	feedin v = forever $ do
-		m <- runner $ net receiveMessage
-		putMVar v $ Right $ RelayMessage m
+				putMVar v $ RelayToPeer (L.fromChunks [b])
+				loop