1--
2-- Licensed to the Apache Software Foundation (ASF) under one
3-- or more contributor license agreements. See the NOTICE file
4-- distributed with this work for additional information
5-- regarding copyright ownership. The ASF licenses this file
6-- to you under the Apache License, Version 2.0 (the
7-- "License"); you may not use this file except in compliance
8-- with the License. You may obtain a copy of the License at
9--
10--   http://www.apache.org/licenses/LICENSE-2.0
11--
12-- Unless required by applicable law or agreed to in writing,
13-- software distributed under the License is distributed on an
14-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15-- KIND, either express or implied. See the License for the
16-- specific language governing permissions and limitations
17-- under the License.
18--
19
20{-# LANGUAGE OverloadedStrings, RecordWildCards, ScopedTypeVariables #-}
21module Main where
22
23import Control.Exception
24import Control.Monad
25import Data.Functor
26import Data.List.Split
27import Data.String
28import Network
29import Network.URI
30import System.Environment
31import System.Exit
32import qualified Data.ByteString.Lazy as LBS
33import qualified Data.HashMap.Strict as Map
34import qualified Data.HashSet as Set
35import qualified Data.Vector as Vector
36import qualified System.IO as IO
37
38import ThriftTest_Iface
39import ThriftTest_Types
40import qualified ThriftTest_Client as Client
41
42import Thrift.Transport
43import Thrift.Transport.Framed
44import Thrift.Transport.Handle
45import Thrift.Transport.HttpClient
46import Thrift.Protocol
47import Thrift.Protocol.Binary
48import Thrift.Protocol.Compact
49import Thrift.Protocol.Header
50import Thrift.Protocol.JSON
51
52data Options = Options
53  { host         :: String
54  , port         :: Int
55  , domainSocket :: String
56  , transport    :: String
57  , protocol     :: ProtocolType
58  -- TODO: Haskell lib does not have SSL support
59  , ssl          :: Bool
60  , testLoops    :: Int
61  }
62  deriving (Show, Eq)
63
64data TransportType = Buffered IO.Handle
65                   | Framed (FramedTransport IO.Handle)
66                   | Http HttpClient
67                   | NoTransport String
68
69getTransport :: String -> String -> Int -> (IO TransportType)
70getTransport "buffered" host port = do
71  h <- hOpen (host, PortNumber $ fromIntegral port)
72  IO.hSetBuffering h $ IO.BlockBuffering Nothing
73  return $ Buffered h
74getTransport "framed" host port = do
75  h <- hOpen (host, PortNumber $ fromIntegral port)
76  t <- openFramedTransport h
77  return $ Framed t
78getTransport "http" host port = let uriStr = "http://" ++ host ++ ":" ++ show port in
79                                case parseURI uriStr of
80                                  Nothing -> do return (NoTransport $ "Failed to parse URI: " ++ uriStr)
81                                  Just(uri) -> do
82                                    t <- openHttpClient uri
83                                    return $ Http t
84getTransport t host port = do return (NoTransport $ "Unsupported transport: " ++ t)
85
86data ProtocolType = Binary
87                  | Compact
88                  | JSON
89                  | Header
90                  deriving (Show, Eq)
91
92getProtocol :: String -> ProtocolType
93getProtocol "binary"  = Binary
94getProtocol "compact" = Compact
95getProtocol "json"    = JSON
96getProtocol "header" = Header
97getProtocol p = error $ "Unsupported Protocol: " ++ p
98
99defaultOptions :: Options
100defaultOptions = Options
101  { port         = 9090
102  , domainSocket = ""
103  , host         = "localhost"
104  , transport    = "buffered"
105  , protocol     = Binary
106  , ssl          = False
107  , testLoops    = 1
108  }
109
110runClient :: Protocol p => p -> IO ()
111runClient p = do
112  let prot = (p,p)
113  putStrLn "Starting Tests"
114
115  -- VOID Test
116  putStrLn "testVoid"
117  Client.testVoid prot
118
119  -- String Test
120  putStrLn "testString"
121  s <- Client.testString prot "Test"
122  when (s /= "Test") exitFailure
123
124  -- Bool Test
125  putStrLn "testBool"
126  bool <- Client.testBool prot True
127  when (not bool) exitFailure
128  putStrLn "testBool"
129  bool <- Client.testBool prot False
130  when (bool) exitFailure
131
132  -- Byte Test
133  putStrLn "testByte"
134  byte <- Client.testByte prot 1
135  when (byte /= 1) exitFailure
136
137  -- I32 Test
138  putStrLn "testI32"
139  i32 <- Client.testI32 prot (-1)
140  when (i32 /= -1) exitFailure
141
142  -- I64 Test
143  putStrLn "testI64"
144  i64 <- Client.testI64 prot (-34359738368)
145  when (i64 /= -34359738368) exitFailure
146
147  -- Double Test
148  putStrLn "testDouble"
149  dub <- Client.testDouble prot (-5.2098523)
150  when (abs (dub + 5.2098523) > 0.001) exitFailure
151
152  -- Binary Test
153  putStrLn "testBinary"
154  bin <- Client.testBinary prot (LBS.pack . reverse $ [-128..127])
155  when ((reverse [-128..127]) /= LBS.unpack bin) exitFailure
156
157  -- Struct Test
158  let structIn = Xtruct{ xtruct_string_thing = "Zero"
159                       , xtruct_byte_thing   = 1
160                       , xtruct_i32_thing    = -3
161                       , xtruct_i64_thing    = -5
162                       }
163  putStrLn "testStruct"
164  structOut <- Client.testStruct prot structIn
165  when (structIn /= structOut) exitFailure
166
167  -- Nested Struct Test
168  let nestIn = Xtruct2{ xtruct2_byte_thing   = 1
169                      , xtruct2_struct_thing = structIn
170                      , xtruct2_i32_thing    = 5
171                      }
172  putStrLn "testNest"
173  nestOut <- Client.testNest prot nestIn
174  when (nestIn /= nestOut) exitFailure
175
176  -- Map Test
177  let mapIn = Map.fromList $ map (\i -> (i, i-10)) [1..5]
178  putStrLn "testMap"
179  mapOut <- Client.testMap prot mapIn
180  when (mapIn /= mapOut) exitFailure
181
182  -- Set Test
183  let setIn = Set.fromList [-2..3]
184  putStrLn "testSet"
185  setOut <- Client.testSet prot setIn
186  when (setIn /= setOut) exitFailure
187
188  -- List Test
189  let listIn = Vector.fromList [-2..3]
190  putStrLn "testList"
191  listOut <- Client.testList prot listIn
192  when (listIn /= listOut) exitFailure
193
194  -- Enum Test
195  putStrLn "testEnum"
196  numz1 <- Client.testEnum prot ONE
197  when (numz1 /= ONE) exitFailure
198
199  putStrLn "testEnum"
200  numz2 <- Client.testEnum prot TWO
201  when (numz2 /= TWO) exitFailure
202
203  putStrLn "testEnum"
204  numz5 <- Client.testEnum prot FIVE
205  when (numz5 /= FIVE) exitFailure
206
207  -- Typedef Test
208  putStrLn "testTypedef"
209  uid <- Client.testTypedef prot 309858235082523
210  when (uid /= 309858235082523) exitFailure
211
212  -- Nested Map Test
213  putStrLn "testMapMap"
214  _ <- Client.testMapMap prot 1
215
216  -- Exception Test
217  putStrLn "testException"
218  exn1 <- try $ Client.testException prot "Xception"
219  case exn1 of
220    Left (Xception _ _) -> return ()
221    _ -> putStrLn (show exn1) >> exitFailure
222
223  putStrLn "testException"
224  exn2 <- try $ Client.testException prot "TException"
225  case exn2 of
226    Left (_ :: SomeException) -> return ()
227    Right _ -> exitFailure
228
229  putStrLn "testException"
230  exn3 <- try $ Client.testException prot "success"
231  case exn3 of
232    Left (_ :: SomeException) -> exitFailure
233    Right _ -> return ()
234
235  -- Multi Exception Test
236  putStrLn "testMultiException"
237  multi1 <- try $ Client.testMultiException prot "Xception" "test 1"
238  case multi1 of
239    Left (Xception _ _) -> return ()
240    _ -> exitFailure
241
242  putStrLn "testMultiException"
243  multi2 <- try $ Client.testMultiException prot "Xception2" "test 2"
244  case multi2 of
245    Left (Xception2 _ _) -> return ()
246    _ -> exitFailure
247
248  putStrLn "testMultiException"
249  multi3 <- try $ Client.testMultiException prot "success" "test 3"
250  case multi3 of
251    Left (_ :: SomeException) -> exitFailure
252    Right _ -> return ()
253
254
255main :: IO ()
256main = do
257  options <- flip parseFlags defaultOptions <$> getArgs
258  case options of
259    Nothing -> showHelp
260    Just Options{..} -> do
261      trans <- Main.getTransport transport host port
262      case trans of
263        Buffered t -> runTest testLoops protocol t
264        Framed t   -> runTest testLoops protocol t
265        Http t     -> runTest testLoops protocol t
266        NoTransport err -> putStrLn err
267  where
268    makeClient p t = case p of
269                       Binary  -> runClient $ BinaryProtocol t
270                       Compact -> runClient $ CompactProtocol t
271                       JSON    -> runClient $ JSONProtocol t
272                       Header  -> createHeaderProtocol t t >>= runClient
273    runTest loops p t = do
274      let client = makeClient p t
275      replicateM_ loops client
276      putStrLn "COMPLETED SUCCESSFULLY"
277
278parseFlags :: [String] -> Options -> Maybe Options
279parseFlags (flag : flags) opts = do
280  let pieces = splitOn "=" flag
281  case pieces of
282    "--port" : arg : _ -> parseFlags flags opts{ port = read arg }
283    "--domain-socket" : arg : _ -> parseFlags flags opts{ domainSocket = read arg }
284    "--host" : arg : _ -> parseFlags flags opts{ host = arg }
285    "--transport" : arg : _ -> parseFlags flags opts{ transport = arg }
286    "--protocol" : arg : _ -> parseFlags flags opts{ protocol = getProtocol arg }
287    "-n" : arg : _ -> parseFlags flags opts{ testLoops = read arg }
288    "--h" : _ -> Nothing
289    "--help" : _ -> Nothing
290    "--ssl" : _ -> parseFlags flags opts{ ssl = True }
291    "--processor-events" : _ -> parseFlags flags opts
292    _ -> Nothing
293parseFlags [] opts = Just opts
294
295showHelp :: IO ()
296showHelp = putStrLn
297  "Allowed options:\n\
298  \  -h [ --help ]               produce help message\n\
299  \  --host arg (=localhost)     Host to connect\n\
300  \  --port arg (=9090)          Port number to connect\n\
301  \  --domain-socket arg         Domain Socket (e.g. /tmp/ThriftTest.thrift),\n\
302  \                              instead of host and port\n\
303  \  --transport arg (=buffered) Transport: buffered, framed, http\n\
304  \  --protocol arg (=binary)    Protocol: binary, compact, json\n\
305  \  --ssl                       Encrypted Transport using SSL\n\
306  \  -n [ --testloops ] arg (=1) Number of Tests"
307