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