1{- 2 Copyright (c) Facebook, Inc. and its affiliates. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15-} 16 17{-# LANGUAGE CPP, ScopedTypeVariables, DeriveDataTypeable #-} 18module Thrift.FuzzerSupport 19where 20 21import Control.Exception (catches, Handler(..), throw, IOException, Exception) 22import Data.Maybe (isNothing) 23import Data.Typeable (Typeable) 24import System.Console.GetOpt (getOpt, usageInfo, ArgOrder(..), OptDescr(..), ArgDescr(..)) 25import System.IO (Handle, hFlush, stdout) 26import System.Random (split) 27import System.Timeout (timeout) 28import Test.QuickCheck.Gen (Gen(..)) 29import Test.QuickCheck.Random (newQCGen) 30import Thrift (AppExn) 31import Thrift.Protocol.Binary (BinaryProtocol(..)) 32import Thrift.Transport (Transport, TransportExn, tClose) 33import Thrift.Transport.Handle (Port, hOpen) 34import Thrift.Transport.Framed (openFramedTransport, FramedTransport) 35 36-- Configuration via command-line parsing 37data Options = Options { 38 opt_host :: String, 39 opt_port :: Port, 40 opt_service :: String, 41 opt_timeout :: Int, 42 opt_framed :: Bool, 43 opt_verbose :: Bool 44} 45 46defaultOptions :: Options 47defaultOptions = Options { 48 opt_host = "localhost" 49 , opt_port = "9090" 50 , opt_service = "ERROR" 51 , opt_timeout = 1 52 , opt_framed = False 53 , opt_verbose = False 54} 55optionsDescriptions :: [OptDescr (Options -> Options)] 56optionsDescriptions = [ 57 Option ['h'] ["host"] (ReqArg getHost "HOST") "hostname of service server" 58 , Option ['p'] ["port"] (ReqArg getPort "PORT") "port of service" 59 , Option ['t'] ["timeout"] (ReqArg getTimeout "TIMEOUT") "timeout (s) to determine if a service died" 60 , Option ['f'] ["framed"] (NoArg getFramed) "use a framed transport" 61 , Option ['v'] ["verbose"] (NoArg getVerbose) "print information for application exceptions" 62 ] 63getHost, getPort, getTimeout :: String -> Options -> Options 64getHost newHost oldOpts = oldOpts { opt_host = newHost } 65getPort newPort oldOpts = oldOpts { opt_port = newPort } 66getTimeout newTimeout oldOpts = oldOpts { opt_timeout = fromIntegral $ (read newTimeout :: Int) } 67 68getFramed, getVerbose :: Options -> Options 69getFramed oldOpts = oldOpts { opt_framed = True } 70getVerbose oldOpts = oldOpts { opt_verbose = True } 71 72getOptions :: [String] -> ([Options -> Options], [String], [String]) 73getOptions = getOpt RequireOrder optionsDescriptions 74 75usage :: String 76usage = usageInfo header optionsDescriptions 77 where 78 header = "[OPTIONS ...] serviceName" 79 80-- timeout 81data Timeout = Timeout deriving (Show, Typeable) 82instance Exception Timeout 83 84-- Generic random data generation 85infexamples :: Gen a -> IO [a] 86infexamples (MkGen m) = 87 do rand <- newQCGen 88 let rnds rnd = rnd1 : rnds rnd2 where (rnd1, rnd2) = split rnd 89 return [(m r n) | (r, n) <- rnds rand `zip` [0,2..] ] 90 91-- Thrift setup 92withHandle :: Options -> (Handle -> IO a) -> IO a 93withHandle opts action = do 94 transport <- getHandle opts 95 result <- action transport 96 tClose transport 97 return result 98 where 99 getHandle (Options host port _service _timeout _framed _verbose) = 100 hOpen (host, port) 101 102withFramedTransport :: Options -> (FramedTransport Handle -> IO a) -> IO a 103withFramedTransport opts action = withHandle opts $ \h -> do 104 transport <- openFramedTransport h 105 result <- action transport 106 tClose transport 107 return result 108 109getClient :: Transport a => a -> (BinaryProtocol a, BinaryProtocol a) 110getClient transport = (BinaryProtocol transport, BinaryProtocol transport) 111 112secsToMicrosecs :: Int 113secsToMicrosecs = 10 ^ (6 :: Int) 114 115withThriftDo :: Transport t => Options -> ((t -> IO ()) -> IO ()) -> ((BinaryProtocol t, BinaryProtocol t) -> IO ()) -> IO () -> IO () 116withThriftDo opts withTransport action onException = do 117 hFlush stdout 118 timed <- timeout (opt_timeout opts * secsToMicrosecs) (withTransport (action . getClient)) 119 if isNothing timed then throw Timeout else return () 120 putStr "." 121 `catches` 122 [ Handler (\ (ex :: TransportExn) -> do { putStrLn "Crashed it. I win. "; onException; throw ex } ) 123 , Handler (\ (_ :: AppExn) -> if opt_verbose opts then putStr "\n" >> onException else putStr "*" ) 124 , Handler (\ (ex :: IOException) -> do { putStrLn "Service is down. I win."; onException; throw ex } ) 125 , Handler (\ (ex :: Timeout) -> do { putStrLn "Timeout. I win."; onException; throw ex } ) 126 ] 127