-
Notifications
You must be signed in to change notification settings - Fork 0
/
NetSpec.hs
196 lines (166 loc) · 5.15 KB
/
NetSpec.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
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
-- | Simplify static Networking tasks.
module Network.NetSpec (
-- * Types and Constructors
NetSpec (..)
, SpecState (..)
-- * Functions
-- ** Running a NetSpec
, runSpec
-- ** Continue and Stop Combinators
, continue
, continue_
, continueIf
, continueIf'
, continueIf_
, stop
, stop_
, stopIf
, stopIf'
, stopIf_
-- * Convenience
-- ** Composition
, (.:)
-- ** IO and Networking
, module I
, module N
-- ** Functors
, module A
-- ** State
, module S
, stateT
) where
import System.IO as I (Handle)
import Network as N (PortID (..))
import Control.Monad.State as S
(StateT (..), execStateT, evalStateT, get, put)
import Control.Applicative as A ((<$>))
import Control.Monad
import Control.Exception
import Data.Traversable as T
import Data.Foldable as F
import Network
import System.IO (hClose)
fst' :: (a,b,c) -> a
fst' (a,_,_) = a
-- | Lift a state function into a 'S.StateT' monad stack
stateT :: Monad m => (s -> (a, s)) -> StateT s m a
stateT = StateT . fmap return
-- | Compose two functions, similar to @.@ from "Prelude".
-- If @h = f .: g@ then @h x y = f (g x y)@.
(.:) :: (a -> b) -> (c -> d -> a) -> c -> d -> b
(.:) f g x y = f (g x y)
-- | Indicate whether to @Continue@ or @Stop@
-- with a given state
data SpecState s = Continue s | Stop s
-- | Continue with a given state
continue :: Monad m => s -> m (SpecState s)
continue = return . Continue
-- | Continue (statless)
continue_ :: Monad m => m (SpecState ())
continue_ = continue ()
-- | Stop with a given state
stop :: Monad m => s -> m (SpecState s)
stop = return . Stop
-- | Stop (stateless)
stop_ :: Monad m => m (SpecState ())
stop_ = stop ()
-- | Conditionally continue with a given state,
-- based on that state and additional given information.
--
-- Recommended usage:
--
-- > _loop = \handles -> continueIf f .: runStateT $ do ...
continueIf :: Monad m => (a -> s -> Bool) -> m (a,s) -> m (SpecState s)
continueIf f ms = do
(a,s) <- ms
if f a s then continue s else stop s
-- | Conditionally continue statelessly,
-- based on given information.
--
-- Recommended usage
--
-- > _loop = \handles () -> continueIf_ f $ do ...
continueIf_ :: Monad m => (a -> Bool) -> m a -> m (SpecState ())
continueIf_ f ms = continueIf (\a () -> f a) (liftM (\x -> (x,())) ms)
-- | Conditionally continue with a given state,
-- based solely on that state.
--
-- Recommended usage:
--
-- > _loop = \handles -> continueIf' f .: execStateT $ do ...
continueIf' :: Monad m => (s -> Bool) -> m s -> m (SpecState s)
continueIf' f ms = continueIf (\() s -> f s) (liftM ((,) ()) ms)
-- | Conditionally stop with a given state,
-- based on that state and additional given information.
stopIf :: Monad m => (a -> s -> Bool) -> m (a,s) -> m (SpecState s)
stopIf f = continueIf (not .: f)
-- | Conditionally stop with a given state,
-- based solely on that state.
stopIf' :: Monad m => (s -> Bool) -> m s -> m (SpecState s)
stopIf' f = continueIf' (not . f)
-- | Conditionally stop statlessly,
-- based on given information.
stopIf_ :: Monad m => (a -> Bool) -> m a -> m (SpecState ())
stopIf_ f = continueIf_ (not . f)
instance Functor SpecState where
fmap f (Continue s) = Continue $ f s
fmap f (Stop s) = Stop $ f s
-- | Define the specification of your networking task
-- as a begin, loop, and end proceedure. Run your NetSpec
-- with 'runSpec'.
--
-- @t@ indicates the 'T.Traversable' structure used.
-- @[]@ is recommended for simplicity, but you are at liberty
-- to use any Traversable you see fit.
--
-- @s@ indicates the type used for state.
-- Use @()@ for a stateless specification.
--
-- A server must specify which ports to listen on,
-- while a client instead specifies tuples of (hostname, port)
-- to connect to.
data NetSpec t s = ServerSpec
{ _ports :: t PortID
, _begin :: t Handle -> IO s
, _loop :: t Handle -> s -> IO (SpecState s)
, _end :: t Handle -> s -> IO ()
}
| ClientSpec
{ _conns :: t (String, PortID)
, _begin :: t Handle -> IO s
, _loop :: t Handle -> s -> IO (SpecState s)
, _end :: t Handle -> s -> IO ()
}
-- | Run a 'NetSpec'.
--
-- Running a spec will step through your 'T.Traversable'
-- of connection descriptions, and replace each one with a 'I.Handle',
-- preserving the structure of the Traversable otherwise.
--
-- Regardless of exceptions, all Handles and Sockets
-- opened by the spec will be closed at the end of the run;
-- you should not need to close any of the Handles given to you
-- by the spec.
--
-- (Note @runSpec@ calls 'N.withSocketsDo' for you)
runSpec :: Traversable t => NetSpec t s -> IO ()
runSpec spec = withSocketsDo $ case spec of
ServerSpec{} -> bracket a c b
ClientSpec{} -> bracket a' c' b'
where
a = do
ss <- T.mapM listenOn $ _ports spec
hs <- fmap fst' <$> T.mapM accept ss
return (ss, hs)
b (_, hs) = _begin spec hs >>= go hs
c (ss, hs) = do
F.mapM_ hClose hs
F.mapM_ sClose ss
a' = T.mapM (uncurry connectTo) $ _conns spec
b' hs = _begin spec hs >>= go hs
c' = F.mapM_ hClose
go hs s = do
res <- _loop spec hs s
case res of
Continue s' -> go hs s'
Stop s' -> _end spec hs s'