-
Notifications
You must be signed in to change notification settings - Fork 0
/
Client.hs
151 lines (124 loc) · 4.39 KB
/
Client.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_HADDOCK prune not-home #-}
{- |
Module : Client
Copyright : (c) 2022 Tim Emiola
Maintainer : Tim Emiola <adetokunbo@emio.la>
SPDX-License-Identifier: BSD3
A demo client that contacts the demo server via a framed protocol.
* it generates a sequence of @'Header's@ that represent server requests
* each request specifies the number of frames the server should send in response
* it connects to the server and begins sending the generated requests
* between each request, it consumes the frames sent by the server in response
* it tracks them and confirms that the requested number of frames are received
* finally, it sends a special @Header@ that signals to the server that it should close the connection
-}
module Main (main) where
import Attoparsec.ToyFrame (
FullFrame,
Header (..),
Payload (..),
buildFrameHeader,
parser,
someTriggers,
)
import Data.Attoparsec.Framer (
Framer,
mkFramer,
runFramer,
setOnBadParse,
setOnClosed,
)
import qualified Data.ByteString as BS
import Data.ByteString.Builder (toLazyByteString)
import qualified Data.ByteString.Lazy as LBS
import Data.IORef (
IORef,
newIORef,
readIORef,
writeIORef,
)
import Data.Text (Text)
import qualified Data.Text.IO as Text
import Network.Run.TCP (runTCPClient)
import Network.Socket (Socket)
import Network.Socket.ByteString (recv, sendAll)
main :: IO ()
main = runTCPClient "127.0.0.1" "3927" $ \s -> do
-- generate the sequence of headers to send to the server, save in the tracking IORef
trackingRef <- someTriggers 1024 >>= newTrackingRef
-- trigger an initial response from the server
trackFrames trackingRef (socketSink s) Nothing
-- get the Framer; mkClientFramer uses trackFrames as its FrameHandler
runFramer $ mkClientFramer trackingRef s
-- print a summary after completing
tracking <- readIORef trackingRef
putStrLn $ "Received " ++ (show $ trackingFrames tracking) ++ " of total size " ++ (show $ trackingBytes tracking)
mkClientFramer :: IORef Tracking -> Socket -> Framer IO FullFrame
mkClientFramer ref s =
let sink = socketSink s
onFullFrame' f = trackFrames ref sink $ Just f
in setOnClosed onClosed $
setOnBadParse (onFailedParse' sink) $
mkFramer parser onFullFrame' (recv s . fromIntegral)
data Tracking = Tracking
{ trackingLeft :: ![Header]
, trackingBytes :: !Int
, trackingFrames :: !Int
, trackingCountdown :: !(Int, Int)
}
newTrackingRef :: [Header] -> IO (IORef Tracking)
newTrackingRef xs = newIORef $ Tracking xs 0 0 (0, 0)
trackFrames :: IORef Tracking -> ByteSink -> Maybe FullFrame -> IO ()
trackFrames trackingRef sink frameMb = do
t <- readIORef trackingRef
let (target, lastCount) = trackingCountdown t
nextCount = lastCount + 1
nextFrames = trackingFrames t + 1
incrWithPayload p =
t
{ trackingCountdown = (target, nextCount)
, trackingFrames = nextFrames
, trackingBytes = trackingBytes t + BS.length p
}
countedUp = nextCount == target
incrOr p' action =
if not countedUp
then writeIORef trackingRef $ incrWithPayload p'
else action
case (frameMb, trackingLeft t) of
(Just (_, Payload p'), []) -> incrOr p' $ do
writeIORef trackingRef $ incrWithPayload p'
sink bye
(Just (_, Payload p'), x : xs) -> incrOr p' $ do
let updatedTracking =
(incrWithPayload p')
{ trackingCountdown = (fromIntegral $ hResponseSize x, 0)
, trackingLeft = xs
}
writeIORef trackingRef updatedTracking
sink $ asBytes x
(Nothing, x : xs) -> do
writeIORef
trackingRef
t
{ trackingCountdown = (fromIntegral $ hResponseSize x, 0)
, trackingLeft = xs
}
sink $ asBytes x
(Nothing, []) -> sink bye
type ByteSink = BS.ByteString -> IO ()
socketSink :: Socket -> ByteSink
socketSink = sendAll
bye :: BS.ByteString
bye = asBytes $ Header 0 0
onFailedParse' :: ByteSink -> Text -> IO ()
onFailedParse' sink cause = do
-- if does not parse as a full frame immediately terminate the connection
Text.putStrLn $ "parse error ended a connection to a toy server: " <> cause
sink bye
onClosed :: IO ()
onClosed = Text.putStrLn "finished at the server too!"
asBytes :: Header -> BS.ByteString
asBytes = LBS.toStrict . toLazyByteString . buildFrameHeader