Server.hs

{- P2P protocol over HTTP, server
 -
 - https://git-annex.branchable.com/design/p2p_protocol_over_http/
 -
 - Copyright 2024-2026 Joey Hess <id@joeyh.name>
 -
 - Licensed under the GNU AGPL version 3 or higher.
 -}

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE BangPatterns #-}

module P2P.Http.Server (
	module P2P.Http,
	module P2P.Http.Server,
	module P2P.Http.Types,
	module P2P.Http.State,
) where

import Annex.Common
import P2P.Http
import P2P.Http.Types
import P2P.Http.State
import P2P.Protocol hiding (Offset, Bypass, auth)
import qualified P2P.Protocol
import P2P.IO
import P2P.Annex
import Annex.WorkerPool
import Types.WorkerPool
import Remote.List
import Types.Direction
import Utility.Metered

import Servant
import qualified Servant.Types.SourceT as S
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Lazy.Internal as LI
import Control.Concurrent.Async
import Control.Concurrent.STM
import Control.Concurrent
import System.IO.Unsafe

p2pHttpApp :: TMVar P2PHttpServerState -> Application
p2pHttpApp = serve p2pHttpAPI . serveP2pHttp

serveP2pHttp :: TMVar P2PHttpServerState -> Server P2PHttpAPI
serveP2pHttp st
	=    serveGet st
	:<|> serveGet st
	:<|> serveGet st
	:<|> serveGet st
	:<|> serveGet st
	:<|> serveCheckPresent st
	:<|> serveCheckPresent st
	:<|> serveCheckPresent st
	:<|> serveCheckPresent st
	:<|> serveCheckPresent st
	:<|> serveRemove st id
	:<|> serveRemove st id
	:<|> serveRemove st id
	:<|> serveRemove st dePlus
	:<|> serveRemove st dePlus
	:<|> serveRemoveBefore st
	:<|> serveRemoveBefore st
	:<|> serveGetTimestamp st
	:<|> serveGetTimestamp st
	:<|> servePut st id
	:<|> servePut' st id
	:<|> servePut' st id
	:<|> servePut' st dePlus
	:<|> servePut' st dePlus
	:<|> servePutOffset st id
	:<|> servePutOffset st id
	:<|> servePutOffset st id
	:<|> servePutOffset st dePlus
	:<|> serveLockContent st
	:<|> serveLockContent st
	:<|> serveLockContent st
	:<|> serveLockContent st
	:<|> serveLockContent st
	:<|> serveKeepLocked st
	:<|> serveKeepLocked st
	:<|> serveKeepLocked st
	:<|> serveKeepLocked st
	:<|> serveKeepLocked st
	:<|> serveGetGeneric st

serveGetGeneric
	:: TMVar P2PHttpServerState
	-> B64UUID ServerSide
	-> B64Key
	-> Maybe (B64UUID ClientSide)
	-> [B64UUID Bypass]
	-> IsSecure
	-> Maybe Auth
	-> Handler (Headers '[DataLengthHeader] (S.SourceT IO B.ByteString))
serveGetGeneric st su@(B64UUID u) k mcu bypass =
	-- Use V0 because it does not alter the returned data to indicate
	-- Invalid content.
	serveGet st su V0 k (fromMaybe scu mcu) bypass Nothing Nothing
  where
	-- Reuse server UUID as client UUID.
	scu = B64UUID u :: B64UUID ClientSide

serveGet
	:: APIVersion v
	=> TMVar P2PHttpServerState
	-> B64UUID ServerSide
	-> v
	-> B64Key
	-> B64UUID ClientSide
	-> [B64UUID Bypass]
	-> Maybe B64FilePath
	-> Maybe Offset
	-> IsSecure
	-> Maybe Auth
	-> Handler (Headers '[DataLengthHeader] (S.SourceT IO B.ByteString))
serveGet mst su apiver (B64Key k) cu bypass baf startat sec auth = do
	(conn, st) <- getP2PConnection apiver WorkerPoolRunner mst cu su bypass sec auth ReadAction id
	bsv <- liftIO newEmptyTMVarIO
	endv <- liftIO newEmptyTMVarIO
	validityv <- liftIO newEmptyTMVarIO
	finalv <- liftIO newEmptyTMVarIO
	annexworker <- liftIO $ async $ handleRequestAnnex st $ do
		let storer _offset len = sendContentWith $ \bs -> liftIO $ do
			atomically $ putTMVar bsv (len, bs)
			atomically $ takeTMVar endv
			signalFullyConsumedByteString $
				connOhdl $ serverP2PConnection conn
			return $ \v -> do
				liftIO $ atomically $ putTMVar validityv v
				return True
		let noothermessages = const Nothing
		enteringStage (TransferStage Upload) $
			runFullProto (clientRunState conn) remoteList (clientP2PConnection conn) $
				void $ receiveContent Nothing nullMeterUpdate
					sizer storer noothermessages getreq
	void $ liftIO $ forkIO $ waitfinal endv finalv conn annexworker
	(Len len, bs) <- liftIO $ atomically $ takeTMVar bsv
	bv <- liftIO $ newMVar (filter (not . B.null) (L.toChunks bs))
	szv <- liftIO $ newMVar 0
	let streamer = S.SourceT $ do
		\s -> s (stream (bv, szv, len, endv, validityv, finalv))
			`onException` streamexception finalv
 	return $ addHeader (DataLength len) streamer
  where
	stream (bv, szv, len, endv, validityv, finalv) =
		S.fromActionStep B.null $
			modifyMVar bv $ nextchunk szv $
				checkvalidity szv len endv validityv finalv
	
	nextchunk szv checkvalid (b:[]) = do
		updateszv szv b
		ifM checkvalid
			( return ([], b)
			-- The key's content is invalid, but
			-- the amount of data is the same as the
			-- DataLengthHeader indicated. Truncate
			-- the response by one byte to indicate
			-- to the client that it's not valid.
			, return ([], B.take (B.length b - 1) b)
			)
	nextchunk szv _checkvalid (b:bs) = do
		updateszv szv b
		return (bs, b)
	nextchunk _szv checkvalid [] = do
		void checkvalid
		-- Result ignored because 0 bytes of data are sent,
		-- so even if the key is invalid, if that's the
		-- amount of data that the DataLengthHeader indicates,
		-- we've successfully served an empty key.
		return ([], mempty)
	
	updateszv szv b = modifyMVar szv $ \sz ->
		let !sz' = sz + fromIntegral (B.length b)
		in return (sz', ())

	-- Returns False when the key's content is invalid, but the
	-- amount of data sent was the same as indicated by the
	-- DataLengthHeader.
	checkvalidity szv len endv validityv finalv =
		ifM (atomically $ isEmptyTMVar endv)
			( do
				atomically $ putTMVar endv ()
				validity <- atomically $ takeTMVar validityv
				sz <- takeMVar szv
				atomically $ putTMVar finalv True
				void $ atomically $ tryPutTMVar endv ()
				return $ case validity of
					Nothing -> True
					Just Valid -> True
					Just Invalid -> sz /= len
			, pure True
			)

	waitfinal endv finalv conn annexworker = do
		-- Wait for everything to be transferred before
		-- stopping the annexworker. The finalv will usually
		-- be written to at the end. If the client disconnects
		-- early that does not happen, so catch STM exceptions.
		alltransferred <- 
			either (const False) id 
				<$> liftIO (tryNonAsync $ atomically $ takeTMVar finalv)
		-- Make sure the annexworker is not left blocked on endv
		-- if the client disconnected early.
		void $ liftIO $ atomically $ tryPutTMVar endv ()
		void $ tryNonAsync $ if alltransferred
			then releaseP2PConnection conn
			else closeP2PConnection conn
		void $ tryNonAsync $ wait annexworker
	
	-- Slowloris attack prevention can cancel the streamer. Be sure to
	-- close the P2P connection when that happens.
	streamexception finalv =
		liftIO $ atomically $ putTMVar finalv False
	
	sizer = pure $ Len $ case startat of
		Just (Offset o) -> fromIntegral o
		Nothing -> 0
	
	getreq offset = P2P.Protocol.GET offset af k
	
	af = ProtoAssociatedFile $ b64FilePathToAssociatedFile baf

serveCheckPresent
	:: APIVersion v
	=> TMVar P2PHttpServerState
	-> B64UUID ServerSide
	-> v
	-> B64Key
	-> B64UUID ClientSide
	-> [B64UUID Bypass]
	-> IsSecure
	-> Maybe Auth
	-> Handler CheckPresentResult
serveCheckPresent st su apiver (B64Key k) cu bypass sec auth = do
	res <- withP2PConnection apiver WorkerPoolRunner st cu su bypass sec auth ReadAction id
		$ \(conn, _) -> liftIO $ proxyClientNetProto conn $ checkPresent k
	case res of
		Right b -> return (CheckPresentResult b)
		Left err -> throwError $ err500 { errBody = encodeBL err }

serveRemove
	:: APIVersion v
	=> TMVar P2PHttpServerState
	-> (RemoveResultPlus -> t)
	-> B64UUID ServerSide
	-> v
	-> B64Key
	-> B64UUID ClientSide
	-> [B64UUID Bypass]
	-> IsSecure
	-> Maybe Auth
	-> Handler t
serveRemove st resultmangle su apiver (B64Key k) cu bypass sec auth = changesBranch st su $ do
	res <- withP2PConnection apiver WorkerPoolRunner st cu su bypass sec auth RemoveAction id
		$ \(conn, _) ->
			liftIO $ proxyClientNetProto conn $ remove Nothing k
	case res of
		(Right b, plusuuids) -> return $ resultmangle $ 
			RemoveResultPlus b (map B64UUID (fromMaybe [] plusuuids))
		(Left err, _) -> throwError $
			err500 { errBody = encodeBL err }

serveRemoveBefore
	:: APIVersion v
	=> TMVar P2PHttpServerState
	-> B64UUID ServerSide
	-> v
	-> B64Key
	-> B64UUID ClientSide
	-> [B64UUID Bypass]
	-> Timestamp
	-> IsSecure
	-> Maybe Auth
	-> Handler RemoveResultPlus
serveRemoveBefore st su apiver (B64Key k) cu bypass (Timestamp ts) sec auth = changesBranch st su $ do
	res <- withP2PConnection apiver WorkerPoolRunner st cu su bypass sec auth RemoveAction id
		$ \(conn, _) ->
			liftIO $ proxyClientNetProto conn $
				removeBeforeRemoteEndTime ts k
	case res of
		(Right b, plusuuids) -> return $ 
			RemoveResultPlus b (map B64UUID (fromMaybe [] plusuuids))
		(Left err, _) -> throwError $
			err500 { errBody = encodeBL err }

serveGetTimestamp
	:: APIVersion v
	=> TMVar P2PHttpServerState
	-> B64UUID ServerSide
	-> v
	-> B64UUID ClientSide
	-> [B64UUID Bypass]
	-> IsSecure
	-> Maybe Auth
	-> Handler GetTimestampResult
serveGetTimestamp st su apiver cu bypass sec auth = do
	res <- withP2PConnection apiver WorkerPoolRunner st cu su bypass sec auth ReadAction id
		$ \(conn, _) ->
			liftIO $ proxyClientNetProto conn getTimestamp
	case res of
		Right ts -> return $ GetTimestampResult (Timestamp ts)
		Left err -> throwError $
			err500 { errBody = encodeBL err }

servePut
	:: APIVersion v
	=> TMVar P2PHttpServerState
	-> (PutResultPlus -> t)
	-> B64UUID ServerSide
	-> v
	-> Maybe Bool
	-> DataLength
	-> B64Key
	-> B64UUID ClientSide
	-> [B64UUID Bypass]
	-> Maybe B64FilePath
	-> Maybe Offset
	-> S.SourceT IO B.ByteString
	-> IsSecure
	-> Maybe Auth
	-> Handler t
servePut mst resultmangle su apiver (Just True) _ k cu bypass baf _ _ sec auth = changesBranch mst su $ do
	res <- withP2PConnection' apiver WorkerPoolRunner mst cu su bypass sec auth WriteAction
		(\cst -> cst { connectionWaitVar = False }) (liftIO . protoaction)
	servePutResult resultmangle res
  where
	protoaction conn = servePutAction conn k baf $ \_offset -> do
		net $ sendMessage DATA_PRESENT
		checkSuccessPlus
servePut mst resultmangle su apiver _datapresent (DataLength len) k cu bypass baf moffset stream sec auth = changesBranch mst su $ do
	validityv <- liftIO newEmptyTMVarIO
	let validitycheck = local $ runValidityCheck $
		liftIO $ atomically $ readTMVar validityv
	tooshortv <- liftIO newEmptyTMVarIO
	content <- liftIO $ S.unSourceT stream (gather validityv tooshortv)
	res <- withP2PConnection' apiver WorkerPoolRunner mst cu su bypass sec auth WriteAction
		(\cst -> cst { connectionWaitVar = False }) $ \(conn, st) -> do
			liftIO $ void $ async $ checktooshort conn tooshortv
			liftIO (protoaction conn st content validitycheck)
	servePutResult resultmangle res
  where
	protoaction conn st content validitycheck = 
		servePutAction (conn, st) k baf $ \offset' ->
			let offsetdelta = offset' - offset
			in case compare offset' offset of
				EQ -> sendContent' nullMeterUpdate (Len len)
					content validitycheck
				GT -> sendContent' nullMeterUpdate
					(Len (len - fromIntegral offsetdelta))
					(L.drop (fromIntegral offsetdelta) content)
					validitycheck
				LT -> sendContent' nullMeterUpdate
					(Len len)
					content
					(validitycheck >>= \_ -> return Invalid)
	
	offset = case moffset of
		Just (Offset o) -> o
		Nothing -> 0

	-- Streams the ByteString from the client. Avoids returning a longer
	-- than expected ByteString by truncating to the expected length. 
	-- Returns a shorter than expected ByteString when the data is not
	-- valid.
	gather validityv tooshortv = unsafeInterleaveIO . go 0
	  where
		go n S.Stop = do
			atomically $ do
				writeTMVar validityv $
					if n == len then Valid else Invalid
				writeTMVar tooshortv (n /= len)
			return LI.Empty
		go n (S.Error _err) = do
			atomically $ do
				writeTMVar validityv Invalid
				writeTMVar tooshortv (n /= len)
			return LI.Empty
		go n (S.Skip s) = go n s
		go n (S.Effect ms) = ms >>= go n
		go n (S.Yield v s) =
			let !n' = n + fromIntegral (B.length v)
			in if n' > len
				then do
					atomically $ do
						writeTMVar validityv Invalid
						writeTMVar tooshortv True
					return $ LI.Chunk
						(B.take (fromIntegral (len - n')) v)
						LI.Empty
				else LI.Chunk v <$> unsafeInterleaveIO (go n' s)
			
	-- The connection can no longer be used when too short a DATA has
	-- been written to it.
	checktooshort conn tooshortv = do
		liftIO $ whenM (atomically $ takeTMVar tooshortv) $
			closeP2PConnection conn

servePutAction
	:: (P2PConnectionPair, PerRepoServerState)
	-> B64Key
	-> Maybe B64FilePath
	-> (P2P.Protocol.Offset -> Proto (Maybe [UUID]))
	-> IO (Either SomeException (Either ProtoFailure (Maybe [UUID])))
servePutAction (conn, st) (B64Key k) baf a = handleRequestAnnex st $
	enteringStage (TransferStage Download) $
		runFullProto (clientRunState conn) remoteList (clientP2PConnection conn) $
			put' k af a
  where
	af = b64FilePathToAssociatedFile baf

servePutResult :: (PutResultPlus -> t) -> Either SomeException (Either ProtoFailure (Maybe [UUID])) -> Handler t
servePutResult resultmangle res = case res of
	Right (Right (Just plusuuids)) -> return $ resultmangle $
		PutResultPlus True (map B64UUID plusuuids)
	Right (Right Nothing) -> return $ resultmangle $
		PutResultPlus False []
	Right (Left protofail) -> throwError $
		err500 { errBody = encodeBL (describeProtoFailure protofail) }
	Left err -> throwError $
		err500 { errBody = encodeBL (show err) }

servePut'
	:: APIVersion v
	=> TMVar P2PHttpServerState
	-> (PutResultPlus -> t)
	-> B64UUID ServerSide
	-> v
	-> DataLength
	-> B64Key
	-> B64UUID ClientSide
	-> [B64UUID Bypass]
	-> Maybe B64FilePath
	-> Maybe Offset
	-> S.SourceT IO B.ByteString
	-> IsSecure
	-> Maybe Auth
	-> Handler t
servePut' st resultmangle su v = servePut st resultmangle su v Nothing

servePutOffset
	:: APIVersion v
	=> TMVar P2PHttpServerState
	-> (PutOffsetResultPlus -> t)
	-> B64UUID ServerSide
	-> v
	-> B64Key
	-> B64UUID ClientSide
	-> [B64UUID Bypass]
	-> IsSecure
	-> Maybe Auth
	-> Handler t
servePutOffset st resultmangle su apiver (B64Key k) cu bypass sec auth = do
	res <- withP2PConnection apiver WorkerPoolRunner st cu su bypass sec auth WriteAction
		(\cst -> cst { connectionWaitVar = False }) $ \(conn, _) ->
			liftIO $ proxyClientNetProto conn $ getPutOffset k af
	case res of
		Right offset -> return $ resultmangle $
			PutOffsetResultPlus (Offset offset)
		Left plusuuids -> return $ resultmangle $
			PutOffsetResultAlreadyHavePlus (map B64UUID plusuuids)
  where
	af = AssociatedFile Nothing

serveLockContent
	:: APIVersion v
	=> TMVar P2PHttpServerState
	-> B64UUID ServerSide
	-> v
	-> B64Key
	-> B64UUID ClientSide
	-> [B64UUID Bypass]
	-> IsSecure
	-> Maybe Auth
	-> Handler LockResult
serveLockContent mst su apiver (B64Key k) cu bypass sec auth =
	checkAuthActionClass mst su sec auth LockAction $ \st mode ->
		ifM (liftIO $ consumeLockedFilesQSem st)
			( go st mode
			, return $ LockResult False Nothing
			)
  where
	go st mode = do
		-- Uses RequestRunner to avoid using a worker
		-- pool slot for a potentially long-running lock.
		(conn, st') <- getP2PConnection' apiver RequestRunner cu su bypass id st mode
		liftIO $ mkLocker (lock conn st') (unlock conn st') >>= \case
			Just (locker, lockid) -> do
				liftIO $ storeLock lockid locker st
				return $ LockResult True (Just lockid)
			Nothing -> return $ LockResult False Nothing
	
	lock conn st = do
		lockresv <- newEmptyTMVarIO
		unlockv <- newEmptyTMVarIO
		-- A thread takes the lock, and keeps running
		-- until unlock in order to keep the lock held.
		lockthread <- async $ handleRequestAnnex st $ do
			lockres <- runFullProto (clientRunState conn) remoteList (clientP2PConnection conn) $ do
				net $ sendMessage (LOCKCONTENT k)
				checkSuccess
			liftIO $ atomically $ putTMVar lockresv lockres
			case lockres of
				Right True -> do
					liftIO $ atomically $ takeTMVar unlockv
					void $ runFullProto (clientRunState conn) remoteList (clientP2PConnection conn) $ do
						net $ sendMessage UNLOCKCONTENT
				_ -> return ()
		atomically (takeTMVar lockresv) >>= \case
			Right True -> return (Just (lockthread, unlockv))
			_ -> return Nothing
	
	unlock conn st (lockthread, unlockv) = do
		atomically $ putTMVar unlockv ()
		void $ wait lockthread
		releaseP2PConnection conn
		liftIO $ releaseLockedFilesQSem st

serveKeepLocked
	:: APIVersion v
	=> TMVar P2PHttpServerState
	-> B64UUID ServerSide
	-> v
	-> LockID
	-> Maybe (B64UUID ClientSide)
	-> [B64UUID Bypass]
	-> IsSecure
	-> Maybe Auth
	-> Maybe ConnectionKeepAlive
	-> Maybe KeepAlive
	-> S.SourceT IO UnlockRequest
	-> Handler LockResult
serveKeepLocked mst su _apiver lckid _cu _bypass sec auth _ _ unlockrequeststream = do
	checkAuthActionClass mst su sec auth LockAction $ \st _ -> do
		liftIO $ keepingLocked lckid st
		_ <- liftIO $ S.unSourceT unlockrequeststream (go st)
		return (LockResult False Nothing)
  where
	go st S.Stop = dropLock lckid st
	go st (S.Error _err) = dropLock lckid st
	go st (S.Skip s)    = go st s
	go st (S.Effect ms) = ms >>= go st
	go st (S.Yield (UnlockRequest False) s) = go st s
	go st (S.Yield (UnlockRequest True) _) = dropLock lckid st