~ubuntu-branches/ubuntu/utopic/haskell-wai-extra/utopic : revision 21

43

import System.IO (hClose, openBinaryTempFile)

44

import Network.Wai

45

import Data.Conduit

46

import Data.Conduit.Internal (sinkToPipe)

46

import Data.Conduit.Internal ()

47

import qualified Data.Conduit.List as CL

48

import qualified Data.Conduit.Binary as CB

49

import Control.Monad.IO.Class (liftIO)

51

import Data.Either (partitionEithers)

52

import Control.Monad (when, unless)

53

import Control.Monad.Trans.Class (lift)

54

import Control.Monad.Trans.Resource (allocate, release, register)

55

#if MIN_VERSION_conduit(1, 0, 0)

56

import Data.Conduit.Internal (Pipe (NeedInput, HaveOutput), (>+>), withUpstream, Sink (..), injectLeftovers, ConduitM (..))

54

import Control.Monad.Trans.Resource (allocate, release, register, InternalState, runInternalState)

55

import Data.Conduit.Internal (Pipe (NeedInput, HaveOutput), (>+>), withUpstream, injectLeftovers, ConduitM (..))

57

56

import Data.Void (Void)

58

#endif

59

57

60

58

breakDiscard :: Word8 -> S.ByteString -> (S.ByteString, S.ByteString)

61

59

breakDiscard w s =

91

89

lbsBackEnd _ _ = fmap L.fromChunks CL.consume

92

90

93

91

-- | Save uploaded files on disk as temporary files

94

tempFileBackEnd :: MonadResource m => ignored1 -> ignored2 -> Sink S.ByteString m FilePath

92

--

93

-- Note: starting with version 2.0, it is the responsibility of the caller to

94

-- remove any temp files created by using this backend.

95

tempFileBackEnd :: InternalState -> ignored1 -> ignored2 -> Sink S.ByteString IO FilePath

95

96

tempFileBackEnd = tempFileBackEndOpts getTemporaryDirectory "webenc.buf"

96

97

98

-- | Same as 'tempFileSink', but use configurable temp folders and patterns.

98

tempFileBackEndOpts :: MonadResource m

99

=> IO FilePath -- ^ get temporary directory

99

tempFileBackEndOpts :: IO FilePath -- ^ get temporary directory

100

-> String -- ^ filename pattern

101

-> InternalState

101

102

-> ignored1

102

103

-> ignored2

103

-> Sink S.ByteString m FilePath

104

tempFileBackEndOpts getTmpDir pattern _ _ = do

105

(key, (fp, h)) <- lift $ allocate (do

104

-> Sink S.ByteString IO FilePath

105

tempFileBackEndOpts getTmpDir pattern internalState _ _ = do

106

(key, (fp, h)) <- flip runInternalState internalState $ allocate (do

106

107

tempDir <- getTmpDir

107

108

openBinaryTempFile tempDir pattern) (\(_, h) -> hClose h)

108

_ <- lift $ register $ removeFile fp

109

_ <- runInternalState (register $ removeFile fp) internalState

109

110

CB.sinkHandle h

110

111

lift $ release key

111

112

return fp

128

129

-- type, and returns a `Sink` for storing the contents.

129

130

type BackEnd a = S.ByteString -- ^ parameter name

130

131

-> FileInfo ()

131

-> Sink S.ByteString (ResourceT IO) a

132

-> Sink S.ByteString IO a

132

133

134

data RequestBodyType = UrlEncoded | Multipart S.ByteString

134

135

167

168

169

parseRequestBody :: BackEnd y

169

170

-> Request

170

-> ResourceT IO ([Param], [File y])

171

-> IO ([Param], [File y])

171

172

parseRequestBody s r =

172

173

case getRequestBodyType r of

173

174

Nothing -> return ([], [])

175

176

177

sinkRequestBody :: BackEnd y

177

178

-> RequestBodyType

178

-> Sink S.ByteString (ResourceT IO) ([Param], [File y])

179

-> Sink S.ByteString IO ([Param], [File y])

179

180

sinkRequestBody s r = fmap partitionEithers $ conduitRequestBody s r =$ CL.consume

180

181

182

conduitRequestBody :: BackEnd y

182

183

-> RequestBodyType

183

-> Conduit S.ByteString (ResourceT IO) (Either Param (File y))

184

-> Conduit S.ByteString IO (Either Param (File y))

184

185

conduitRequestBody _ UrlEncoded = do

185

186

-- NOTE: in general, url-encoded data will be in a single chunk.

186

187

-- Therefore, I'm optimizing for the usual case by sticking with

190

191

conduitRequestBody backend (Multipart bound) =

191

192

parsePieces backend $ S8.pack "--" `S.append` bound

192

193

#if MIN_VERSION_conduit(1, 0, 0)

194

takeLine :: Monad m => Consumer S.ByteString m (Maybe S.ByteString)

195

#else

196

takeLine :: Monad m => Pipe S.ByteString S.ByteString o u m (Maybe S.ByteString)

197

#endif

198

195

takeLine =

199

196

go id

200

197

where

209

206

when (S.length y > 1) $ leftover $ S.drop 1 y

210

207

return $ Just $ killCR x

211

208

212

#if MIN_VERSION_conduit(1, 0, 0)

213

takeLines :: Consumer S.ByteString (ResourceT IO) [S.ByteString]

214

#else

215

takeLines :: Pipe S.ByteString S.ByteString o u (ResourceT IO) [S.ByteString]

216

#endif

209

takeLines :: Consumer S.ByteString IO [S.ByteString]

217

210

takeLines = do

218

211

res <- takeLine

219

212

case res of

226

219

227

220

parsePieces :: BackEnd y

228

221

-> S.ByteString

229

#if MIN_VERSION_conduit(1, 0, 0)

230

-> ConduitM S.ByteString (Either Param (File y)) (ResourceT IO) ()

231

#else

232

-> Pipe S.ByteString S.ByteString (Either Param (File y)) u (ResourceT IO) ()

233

#endif

222

-> ConduitM S.ByteString (Either Param (File y)) IO ()

234

223

parsePieces sink bound =

235

224

loop

236

225

where

305

294

-> S.ByteString

306

295

-> FileInfo ()

307

296

-> BackEnd y

308

#if MIN_VERSION_conduit(1, 0, 0)

309

-> ConduitM S.ByteString o (ResourceT IO) (Bool, y)

310

#else

311

-> Pipe S.ByteString S.ByteString o u (ResourceT IO) (Bool, y)

312

#endif

297

-> ConduitM S.ByteString o IO (Bool, y)

313

298

sinkTillBound' bound name fi sink =

314

#if MIN_VERSION_conduit(1, 0, 0)

315

299

ConduitM $ anyOutput $

316

#endif

317

300

conduitTillBound bound >+> withUpstream (fix $ sink name fi)

318

301

where

319

#if MIN_VERSION_conduit(1, 0, 0)

320

fix :: Sink S8.ByteString (ResourceT IO) y -> Pipe Void S8.ByteString Void Bool (ResourceT IO) y

321

fix (ConduitM p) = ignoreTerm >+> injectLeftovers p

302

fix :: Sink S8.ByteString IO y -> Pipe Void S8.ByteString Void Bool IO y

303

fix p = ignoreTerm >+> injectLeftovers (unConduitM p)

322

304

ignoreTerm = await' >>= maybe (return ()) (\x -> yield' x >> ignoreTerm)

323

305

await' = NeedInput (return . Just) (const $ return Nothing)

324

306

yield' = HaveOutput (return ()) (return ())

325

307

326

308

anyOutput p = p >+> dropInput

327

309

dropInput = NeedInput (const dropInput) return

328

#else

329

fix = sinkToPipe

330

#endif

331

310

332

311

conduitTillBound :: Monad m

333

312

=> S.ByteString -- bound

334

#if MIN_VERSION_conduit(1, 0, 0)

335

313

-> Pipe S.ByteString S.ByteString S.ByteString () m Bool

336

#else

337

-> Pipe S.ByteString S.ByteString S.ByteString u m Bool

338

#endif

339

314

conduitTillBound bound =

340

#if MIN_VERSION_conduit(1, 0, 0)

341

315

unConduitM $

342

#endif

343

316

go id

344

317

where

345

318

go front = await >>= maybe (close front) (push front)

369

342

sinkTillBound :: S.ByteString

370

343

-> (x -> S.ByteString -> IO x)

371

344

-> x

372

#if MIN_VERSION_conduit(1, 0, 0)

373

-> Consumer S.ByteString (ResourceT IO) (Bool, x)

374

#else

375

-> Pipe S.ByteString S.ByteString o u (ResourceT IO) (Bool, x)

376

#endif

345

-> Consumer S.ByteString IO (Bool, x)

377

346

sinkTillBound bound iter seed0 =

378

#if MIN_VERSION_conduit(1, 0, 0)

379

347

ConduitM $

380

#endif

381

348

(conduitTillBound bound >+> (withUpstream $ ij $ CL.foldM iter' seed0))

382

349

where

383

350

iter' a b = liftIO $ iter a b

384

#if MIN_VERSION_conduit(1, 0, 0)

385

351

ij (ConduitM p) = ignoreTerm >+> injectLeftovers p

386

352

ignoreTerm = await' >>= maybe (return ()) (\x -> yield' x >> ignoreTerm)

387

353

await' = NeedInput (return . Just) (const $ return Nothing)

388

354

yield' = HaveOutput (return ()) (return ())

389

#else

390

ij = id

391

#endif

392

355

393

356

parseAttrs :: S.ByteString -> [(S.ByteString, S.ByteString)]

394

357

parseAttrs = map go . S.split 59 -- semicolon