@@ -0,0 +1,183 @@
1+{-# LANGUAGE BangPatterns, OverloadedStrings #-}
2+
3+module Network.Wai.Middleware.Push.Referer (
4+ pushOnReferer
5+ , defaultMakePushPromise
6+ ) where
7+
8+import Control.Monad (when)
9+import Control.Reaper
10+import Data.ByteString (ByteString)
11+import qualified Data.ByteString as BS
12+import Data.ByteString.Internal (ByteString(..), memchr)
13+import Data.Map (Map)
14+import qualified Data.Map.Strict as M
15+import Data.Maybe (isNothing)
16+import Data.Set (Set)
17+import qualified Data.Set as S
18+import Data.Word (Word8)
19+import Data.Word8
20+import Foreign.ForeignPtr (withForeignPtr, ForeignPtr)
21+import Foreign.Ptr (Ptr, plusPtr, minusPtr, nullPtr)
22+import Foreign.Storable (peek)
23+import Network.Wai
24+import Network.Wai.Handler.Warp
25+import Network.Wai.Internal (Response(..))
26+import System.IO.Unsafe (unsafePerformIO)
27+
28+-- $setup
29+-- >>> :set -XOverloadedStrings
30+
31+type URLPath = ByteString
32+
33+type Cache = Map URLPath (Set PushPromise)
34+
35+emptyCache :: Cache
36+emptyCache = M.empty
37+
38+cacheReaper :: Reaper Cache (URLPath,PushPromise)
39+cacheReaper = unsafePerformIO $ mkReaper settings
40+{-# NOINLINE cacheReaper #-}
41+
42+settings :: ReaperSettings Cache (URLPath,PushPromise)
43+settings = defaultReaperSettings {
44+ reaperAction = \_ -> return (\_ -> emptyCache)
45+ , reaperCons = insert
46+ , reaperNull = M.null
47+ , reaperEmpty = emptyCache
48+ }
49+
50+insert :: (URLPath,PushPromise) -> Cache -> Cache
51+insert (path,pp) m = M.alter ins path m
52+ where
53+ ins Nothing = Just $ S.singleton pp
54+ ins (Just set) = Just $ S.insert pp set
55+
56+pushOnReferer :: (URLPath -> URLPath -> FilePath -> IO (Maybe PushPromise))
57+ -> Middleware
58+pushOnReferer func app req sendResponse = app req $ \res -> do
59+ let !path = rawPathInfo req
60+ m <- reaperRead cacheReaper
61+ case M.lookup path m of
62+ Nothing -> case requestHeaderReferer req of
63+ Nothing -> return ()
64+ Just referer -> case res of
65+ ResponseFile _ _ file Nothing -> do
66+ (mauth,refPath) <- parseUrl referer
67+ when (isNothing mauth
68+ || requestHeaderHost req == mauth) $ do
69+ when (path /= refPath) $ do -- just in case
70+ mpp <- func refPath path file
71+ case mpp of
72+ Nothing -> return ()
73+ Just pp -> reaperAdd cacheReaper (refPath,pp)
74+ _ -> return ()
75+ Just pset -> do
76+ let !ps = S.toList pset
77+ !h2d = defaultHTTP2Data { http2dataPushPromise = ps}
78+ setHTTP2Data req (Just h2d)
79+ sendResponse res
80+
81+defaultMakePushPromise :: URLPath -- ^ path in referer
82+ -> URLPath -- ^ path to be pushed
83+ -> FilePath -- ^ file to be pushed
84+ -> IO (Maybe PushPromise)
85+defaultMakePushPromise refPath path file
86+ | isHTML refPath = case getCT path of
87+ Nothing -> return Nothing
88+ Just ct -> do
89+ let pp = defaultPushPromise {
90+ promisedPath = path
91+ , promisedFile = file
92+ , promisedResponseHeaders = [("content-type", ct)
93+ ,("x-http2-push", refPath)]
94+ }
95+ return $ Just pp
96+ | otherwise = return Nothing
97+
98+getCT :: URLPath -> Maybe ByteString
99+getCT p
100+ | ".js" `BS.isSuffixOf` p = Just "application/javascript"
101+ | ".css" `BS.isSuffixOf` p = Just "text/css"
102+ | otherwise = Nothing
103+
104+isHTML :: URLPath -> Bool
105+isHTML p = ("/" `BS.isSuffixOf` p)
106+ || (".html" `BS.isSuffixOf` p)
107+ || (".htm" `BS.isSuffixOf` p)
108+
109+-- |
110+--
111+-- >>> parseUrl ""
112+-- (Nothing,"")
113+-- >>> parseUrl "/"
114+-- (Nothing,"/")
115+-- >>> parseUrl "ht"
116+-- (Nothing,"")
117+-- >>> parseUrl "http://example.com/foo/bar/"
118+-- (Just "example.com","/foo/bar/")
119+-- >>> parseUrl "https://www.example.com/path/to/dir/"
120+-- (Just "www.example.com","/path/to/dir/")
121+-- >>> parseUrl "http://www.example.com:8080/path/to/dir/"
122+-- (Just "www.example.com:8080","/path/to/dir/")
123+-- >>> parseUrl "//www.example.com:8080/path/to/dir/"
124+-- (Just "www.example.com:8080","/path/to/dir/")
125+-- >>> parseUrl "/path/to/dir/"
126+-- (Nothing,"/path/to/dir/")
127+
128+parseUrl :: ByteString -> IO (Maybe ByteString, URLPath)
129+parseUrl bs@(PS fptr0 off len)
130+ | len == 0 = return (Nothing, "")
131+ | len == 1 = return (Nothing, bs)
132+ | otherwise = withForeignPtr fptr0 $ \ptr0 -> do
133+ let begptr = ptr0 `plusPtr` off
134+ limptr = begptr `plusPtr` len
135+ parseUrl' fptr0 ptr0 begptr limptr len
136+
137+parseUrl' :: ForeignPtr Word8 -> Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> Int
138+ -> IO (Maybe ByteString, URLPath)
139+parseUrl' fptr0 ptr0 begptr limptr len0 = do
140+ w0 <- peek begptr
141+ if w0 == _slash then do
142+ w1 <- peek $ begptr `plusPtr` 1
143+ if w1 == _slash then
144+ doubleSlashed begptr len0
145+ else
146+ slashed begptr len0 Nothing
147+ else do
148+ colonptr <- memchr begptr _colon $ fromIntegral len0
149+ if colonptr == nullPtr then
150+ return (Nothing, "")
151+ else do
152+ let !authptr = colonptr `plusPtr` 1
153+ doubleSlashed authptr (limptr `minusPtr` authptr)
154+ where
155+ -- // / ?
156+ doubleSlashed :: Ptr Word8 -> Int -> IO (Maybe ByteString, URLPath)
157+ doubleSlashed ptr len
158+ | len < 2 = return (Nothing, "")
159+ | otherwise = do
160+ let ptr1 = ptr `plusPtr` 2
161+ pathptr <- memchr ptr1 _slash $ fromIntegral len
162+ if pathptr == nullPtr then
163+ return (Nothing, "")
164+ else do
165+ let !auth = bs ptr0 ptr1 pathptr
166+ slashed pathptr (limptr `minusPtr` pathptr) (Just auth)
167+
168+ -- / ?
169+ slashed :: Ptr Word8 -> Int -> Maybe ByteString -> IO (Maybe ByteString, URLPath)
170+ slashed ptr len mauth = do
171+ questionptr <- memchr ptr _question $ fromIntegral len
172+ if questionptr == nullPtr then do
173+ let !path = bs ptr0 ptr limptr
174+ return (mauth, path)
175+ else do
176+ let !path = bs ptr0 ptr questionptr
177+ return (mauth, path)
178+ bs p0 p1 p2 = path
179+ where
180+ !off = p1 `minusPtr` p0
181+ !siz = p2 `minusPtr` p1
182+ !path = PS fptr0 off siz
183+