Do not leak uncaught exceptions from serivce handelr

Summary:
Do not leak uncaught exceptions from service handler which may later crash
the whole process. We need reduce the response to a normal form to ensure
no exception is leaked due to laziness.

Reviewed By: josefs

Differential Revision: D63715052

fbshipit-source-id: 19c41c3052098b923aca475f3ad95751188ec747

lib/Thrift/Channel/SocketChannel/Server.hs

@@ -160,13 +160,13 @@ withServerIO p mport maxQueuedConns handler postProcess client  = do
               seqNum' <- counter
               case mexc of
                 Just (_exc, _blame) -> do
-                  _ <- sendBS sock (toStrict $ toLazyByteString response)
+                  _ <- sendBS sock response
                   return (Nothing, seqNum')
                 Nothing -> do
                   let info = Map.lookup (reqName cmd) (methodsInfo (Proxy :: Proxy c))
                       isOneway = maybe False methodIsOneway info
                   unless isOneway $
-                    void $ sendBS sock (toStrict $ toLazyByteString response)
+                    void $ sendBS sock response
                   if BS.null leftover
                     then return (Nothing, seqNum')
                     else processInput seqNum' counter leftover sock

lib/Thrift/Processor.hs

@@ -85,13 +85,14 @@ process proxy seqNum handler postProcess input = do
       let ex = ApplicationException (Text.pack err)
             ApplicationExceptionType_ProtocolError
       return
-        ( genMsgBegin proxy "" 3 seqNum
-          <> buildStruct proxy ex
-          <> genMsgEnd proxy
+        ( toStrict $ toLazyByteString $
+            genMsgBegin proxy "" 3 seqNum
+            <> buildStruct proxy ex
+            <> genMsgEnd proxy
         , Just (toException ex, ClientError)
         , [] )
     Right (Some cmd) -> processCommand proxy seqNum handler postProcess cmd
-  return (toStrict (toLazyByteString response), exc, headers)
+  return (response, exc, headers)
 
 processCommand
   :: (Processor s, Protocol p)
@@ -100,14 +101,27 @@ processCommand
   -> (forall r . s r -> IO r) -- ^ Handler for user-code
   -> (forall r . s r -> Either SomeException r -> Header)
   -> s r                      -- ^ input command
-  -> IO (Builder, Maybe (SomeException, Blame), Header)
-processCommand proxy seqNum handler postProcess cmd = do
-  -- Run the handler and generate its return struct, forcing evaluation
-  res <- try (handler cmd)
-  let (builder, exc) = respWriter proxy seqNum cmd res
-      headers = postProcess cmd res
-  builder' <- evaluate builder
-  return (builder', exc, headers)
+  -> IO (ByteString, Maybe (SomeException, Blame), Header)
+processCommand proxy seqNum handler postProcess cmd =
+  -- in case we cannot serialize the uncaught exception itself (e.g. due to
+  -- another exception being thrown when serializing it), we fallback to an
+  -- predefined one, assuming @respWriter@ and @postProcess@ never fail on
+  -- it.
+  handle (\(_ :: SomeException) -> buildResp $ Left nonserializableError) $
+  -- catch all uncaught exception from handler to prevent it from being leaked.
+  handle (\e -> buildResp $ Left e) $
+  buildResp . Right =<< handler cmd
+  where
+    nonserializableError = toException $ ApplicationException
+      "processCommand: uncaught non-serializable error"
+      ApplicationExceptionType_InternalError
+    buildResp res =
+      -- force evaluation of the response
+      evaluate $ bs `seq` exc `seq` headers `seq` (bs, exc, headers)
+      where
+        (builder, exc) = respWriter proxy seqNum cmd res
+        headers = postProcess cmd res
+        bs = toStrict $ toLazyByteString builder
 
 msgParser
   :: (Processor s, Protocol p)

lib/test/if/math.thrift

@@ -12,11 +12,18 @@ service Adder {
   i64 add(1: i64 x, 2: i64 y);
 }
 
+struct QuotRemResponse {
+  1: i64 quot;
+  2: i64 rem;
+}
+
 service Calculator extends Adder {
   double divide(1: double dividend, 2: double divisor) throws (
     1: DivideByZero divisionError,
   );
 
+  QuotRemResponse quotRem(1: i64 dividend, 2: i64 divisor);
+
   oneway void put(1: i64 val);
 
   oneway void putMany(1: list<i64> val) (haxl.batched);

lib/test/lib/TestCommon.hs

@@ -37,6 +37,7 @@ processCommand _ (SuperAdder (Add x y)) = pure $ x + y
 processCommand _ (Divide x y)
   | y == 0    = throwIO DivideByZero
   | otherwise = pure $ x / y
+processCommand _ (QuotRem x y) = pure $ QuotRemResponse (x `quot` y) (x `rem` y)
 processCommand state (Put x) = writeIORef state x
 processCommand state (PutMany xs) = mapM_ (writeIORef state) xs
 processCommand state Get = readIORef state

server/test/HandlerTest.hs

@@ -11,9 +11,12 @@ module HandlerTest where
 import Control.Concurrent
 import Control.Exception
 import Control.Monad.Trans.Class
+import Control.Monad.Trans.Reader (mapReaderT)
 import Data.Proxy
+import qualified Data.Text as Text
 import Thrift.Monad
 import Thrift.Processor
+import Thrift.Protocol.ApplicationException.Types
 import Thrift.Protocol.Binary
 
 import TestRunner
@@ -22,6 +25,8 @@ import Test.HUnit hiding (State)
 import EchoHandler
 import Echoer.Echoer.Client
 import Math.Adder.Client
+import Math.Calculator.Client as Calculator
+import Math.Types (QuotRemResponse(..))
 import TestChannel
 
 mkTest :: String -> ThriftM Binary TestChannel Echoer () -> Test
@@ -46,8 +51,26 @@ addTest = mkTest "Add" $ do
   res <- add 4 5
   lift $ assertEqual "Add added" (4 + 5) res
 
+quotRemTest :: Test
+quotRemTest = mkTest "QuotRem" $ do
+  res <- Calculator.quotRem 7 3
+  lift $ assertEqual "QuotRem" (QuotRemResponse 2 1) res
+
+quotRemByZeroTest :: Test
+quotRemByZeroTest = mkTest "QuotRemByZero" $ do
+  res <- mapReaderT try $ Calculator.quotRem 1 0
+  lift $ case res of
+    Left e | Just ApplicationException{..} <- fromException e -> do
+      assertEqual "message" (Text.pack $ show DivideByZero)
+        applicationException_message
+      assertEqual "type" ApplicationExceptionType_InternalError
+        applicationException_type
+    _ -> assertFailure $ "An ApplicationException is expected, got " ++ show res
+
 main :: IO ()
 main = testRunner $ TestList
   [ echoTest
   , addTest
+  , quotRemTest
+  , quotRemByZeroTest
   ]

server/test/common/CalculatorHandler.hs

@@ -32,6 +32,9 @@ calculatorHandler _ (SuperAdder (Add x y)) = return $ x + y
 calculatorHandler _ (Divide x y)
   | y == 0    = throw DivideByZero
   | otherwise = return $ x / y
+calculatorHandler _ (QuotRem x y) =
+  -- Intentionally leak an unhandled error
+  return $ QuotRemResponse (x `quot` y) (x `rem` y)
 calculatorHandler (CalculatorState ref) (Put v) =
   writeIORef ref (fromIntegral v)
 calculatorHandler (CalculatorState ref) (PutMany v) =