From 726c815a7fcf7831a1288f959e410ef23a501174 Mon Sep 17 00:00:00 2001
From: Joey Hess <joeyh@joeyh.name>
Date: Mon, 22 Jul 2024 11:26:22 -0400
Subject: [PATCH] fix releasing of p2p connection

---
 P2P/Http.hs       | 20 ++++++++++----------
 P2P/Http/State.hs | 25 +++++++++++++++++++++----
 2 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/P2P/Http.hs b/P2P/Http.hs
index 433c51cc52..d120afba49 100644
--- a/P2P/Http.hs
+++ b/P2P/Http.hs
@@ -322,7 +322,7 @@ serveCheckPresent
 	-> Maybe Auth
 	-> Handler CheckPresentResult
 serveCheckPresent st su apiver (B64Key k) cu bypass sec auth = do
-	res <- withP2PConnection apiver st cu su bypass sec auth ReadAction
+	res <- withP2PConnection apiver st cu su bypass sec auth ReadAction id
 		$ \conn -> liftIO $ proxyClientNetProto conn $ checkPresent k
 	case res of
 		Right b -> return (CheckPresentResult b)
@@ -373,7 +373,7 @@ serveRemove
 	-> Maybe Auth
 	-> Handler t
 serveRemove st resultmangle su apiver (B64Key k) cu bypass sec auth = do
-	res <- withP2PConnection apiver st cu su bypass sec auth RemoveAction
+	res <- withP2PConnection apiver st cu su bypass sec auth RemoveAction id
 		$ \conn ->
 			liftIO $ proxyClientNetProto conn $ remove Nothing k
 	case res of
@@ -429,7 +429,7 @@ serveRemoveBefore
 	-> Maybe Auth
 	-> Handler RemoveResultPlus
 serveRemoveBefore st su apiver (B64Key k) cu bypass (Timestamp ts) sec auth = do
-	res <- withP2PConnection apiver st cu su bypass sec auth RemoveAction
+	res <- withP2PConnection apiver st cu su bypass sec auth RemoveAction id
 		$ \conn ->
 			liftIO $ proxyClientNetProto conn $
 				removeBeforeRemoteEndTime ts k
@@ -481,7 +481,7 @@ serveGetTimestamp
 	-> Maybe Auth
 	-> Handler GetTimestampResult
 serveGetTimestamp st su apiver cu bypass sec auth = do
-	res <- withP2PConnection apiver st cu su bypass sec auth ReadAction
+	res <- withP2PConnection apiver st cu su bypass sec auth ReadAction id
 		$ \conn ->
 			liftIO $ proxyClientNetProto conn getTimestamp
 	case res of
@@ -545,12 +545,12 @@ servePut st resultmangle su apiver (DataLength len) (B64Key k) cu bypass baf mof
 	let validitycheck = local $ runValidityCheck $
 		liftIO $ atomically $ readTMVar validityv
 	content <- liftIO $ S.unSourceT stream (gather validityv)
-	conn <- getP2PConnection apiver st cu su bypass sec auth WriteAction $
-		\st -> st { connectionWaitVar = False }
-	res <- liftIO $ inAnnexWorker st $
-		enteringStage (TransferStage Download) $
-			runFullProto (clientRunState conn) (clientP2PConnection conn) $
-				protoaction content validitycheck
+	res <- withP2PConnection' apiver st cu su bypass sec auth WriteAction
+		(\st -> st { connectionWaitVar = False }) $ \conn ->
+			liftIO $ inAnnexWorker st $
+				enteringStage (TransferStage Download) $
+					runFullProto (clientRunState conn) (clientP2PConnection conn) $
+						protoaction content validitycheck
 	case res of
 		Right (Right (Just plusuuids)) -> return $ resultmangle $
 			PutResultPlus True (map B64UUID plusuuids)
diff --git a/P2P/Http/State.hs b/P2P/Http/State.hs
index bbd68003aa..559a9048d9 100644
--- a/P2P/Http/State.hs
+++ b/P2P/Http/State.hs
@@ -62,18 +62,35 @@ withP2PConnection
 	-> IsSecure
 	-> Maybe Auth
 	-> ActionClass
+	-> (ConnectionParams -> ConnectionParams)
 	-> (P2PConnectionPair -> Handler (Either ProtoFailure a))
 	-> Handler a
-withP2PConnection apiver st cu su bypass sec auth actionclass connaction = do
-	conn <- getP2PConnection apiver st cu su bypass sec auth actionclass id
-	connaction' conn
-		`finally` liftIO (releaseP2PConnection conn)
+withP2PConnection apiver st cu su bypass sec auth actionclass fconnparams connaction =
+	withP2PConnection' apiver st cu su bypass sec auth actionclass fconnparams connaction'
   where
 	connaction' conn = connaction conn >>= \case
 		Right r -> return r
 		Left err -> throwError $
 			err500 { errBody = encodeBL (describeProtoFailure err) }
 
+withP2PConnection'
+	:: APIVersion v
+	=> v
+	-> P2PHttpServerState
+	-> B64UUID ClientSide
+	-> B64UUID ServerSide
+	-> [B64UUID Bypass]
+	-> IsSecure
+	-> Maybe Auth
+	-> ActionClass
+	-> (ConnectionParams -> ConnectionParams)
+	-> (P2PConnectionPair -> Handler a)
+	-> Handler a
+withP2PConnection' apiver st cu su bypass sec auth actionclass fconnparams connaction = do
+	conn <- getP2PConnection apiver st cu su bypass sec auth actionclass fconnparams
+	connaction conn
+		`finally` liftIO (releaseP2PConnection conn)
+
 getP2PConnection
 	:: APIVersion v
 	=> v