Referer.hs

{-# LANGUAGE BangPatterns, OverloadedStrings #-}

-- | Middleware for server push learning dependency based on Referer:.
module Network.Wai.Middleware.Push.Referer (
    pushOnReferer
  , MakePushPromise
  , defaultMakePushPromise
  , URLPath
  ) where

import Control.Monad (when)
import Control.Reaper
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.ByteString.Internal (ByteString(..), memchr)
import Data.Map (Map)
import qualified Data.Map.Strict as M
import Data.Maybe (isNothing)
import Data.Set (Set)
import qualified Data.Set as S
import Data.Word (Word8)
import Data.Word8
import Foreign.ForeignPtr (withForeignPtr, ForeignPtr)
import Foreign.Ptr (Ptr, plusPtr, minusPtr, nullPtr)
import Foreign.Storable (peek)
import Network.Wai
import Network.Wai.Handler.Warp
import Network.Wai.Internal (Response(..))
import System.IO.Unsafe (unsafePerformIO)

-- $setup
-- >>> :set -XOverloadedStrings

-- | Making a push promise based on Referer:,
--   path to be pushed and file to be pushed.
--   If the middleware should push this file in the next time when
--   the page of Referer: is accessed,
--   this function should return 'Just'.
--   If 'Nothing' is returned,
--   the middleware learns nothing.
type MakePushPromise = URLPath  -- ^ path in referer
                    -> URLPath  -- ^ path to be pushed
                    -> FilePath -- ^ file to be pushed
                    -> IO (Maybe PushPromise)

-- | Type for URL path.
type URLPath = ByteString

type Cache = Map URLPath (Set PushPromise)

emptyCache :: Cache
emptyCache = M.empty

cacheReaper :: Reaper Cache (URLPath,PushPromise)
cacheReaper = unsafePerformIO $ mkReaper settings
{-# NOINLINE cacheReaper #-}

settings :: ReaperSettings Cache (URLPath,PushPromise)
settings = defaultReaperSettings {
      reaperAction = \_ -> return (\_ -> emptyCache)
    , reaperCons   = insert
    , reaperNull   = M.null
    , reaperEmpty  = emptyCache
--  , reaperDelay  = 30000000 -- FIXME hard-coding
    }

insert :: (URLPath,PushPromise) -> Cache -> Cache
insert (path,pp) m = M.alter ins path m
  where
    ins Nothing    = Just $ S.singleton pp
    ins (Just set) = Just $ S.insert pp set

-- | The middleware to push files based on Referer:.
--   Learning strategy is implemented in the first argument.
--   Learning information is kept for 30 seconds.
pushOnReferer :: MakePushPromise -> Middleware
pushOnReferer func app req sendResponse = app req $ \res -> do
    let !path = rawPathInfo req
    m <- reaperRead cacheReaper
    case M.lookup path m of
        Nothing -> case requestHeaderReferer req of
            Nothing      -> return ()
            Just referer -> case res of
                ResponseFile _ _ file Nothing -> do
                    (mauth,refPath) <- parseUrl referer
                    when (isNothing mauth
                       || requestHeaderHost req == mauth) $ do
                        when (path /= refPath) $ do -- just in case
                            mpp <- func refPath path file
                            case mpp of
                                Nothing -> return ()
                                Just pp -> reaperAdd cacheReaper (refPath,pp)
                _ -> return ()
        Just pset -> do
            let !ps = S.toList pset
                !h2d = defaultHTTP2Data { http2dataPushPromise = ps}
            setHTTP2Data req (Just h2d)
    sendResponse res

-- | Learn if the file to be pushed is CSS (.css) or JavaScript (.js) file
--   AND the Referer: ends with \"/\" or \".html\" or \".htm\".
defaultMakePushPromise :: MakePushPromise
defaultMakePushPromise refPath path file
  | isHTML refPath = case getCT path of
      Nothing -> return Nothing
      Just ct -> do
          let pp = defaultPushPromise {
                       promisedPath = path
                     , promisedFile = file
                     , promisedResponseHeaders = [("content-type", ct)
                                                 ,("x-http2-push", refPath)]
                     }
          return $ Just pp
  | otherwise = return Nothing

getCT :: URLPath -> Maybe ByteString
getCT p
  | ".js"  `BS.isSuffixOf` p = Just "application/javascript"
  | ".css" `BS.isSuffixOf` p = Just "text/css"
  | otherwise                = Nothing

isHTML :: URLPath -> Bool
isHTML p = ("/" `BS.isSuffixOf` p)
        || (".html" `BS.isSuffixOf` p)
        || (".htm" `BS.isSuffixOf` p)

-- |
--
-- >>> parseUrl ""
-- (Nothing,"")
-- >>> parseUrl "/"
-- (Nothing,"/")
-- >>> parseUrl "ht"
-- (Nothing,"")
-- >>> parseUrl "http://example.com/foo/bar/"
-- (Just "example.com","/foo/bar/")
-- >>> parseUrl "https://www.example.com/path/to/dir/"
-- (Just "www.example.com","/path/to/dir/")
-- >>> parseUrl "http://www.example.com:8080/path/to/dir/"
-- (Just "www.example.com:8080","/path/to/dir/")
-- >>> parseUrl "//www.example.com:8080/path/to/dir/"
-- (Just "www.example.com:8080","/path/to/dir/")
-- >>> parseUrl "/path/to/dir/"
-- (Nothing,"/path/to/dir/")

parseUrl :: ByteString -> IO (Maybe ByteString, URLPath)
parseUrl bs@(PS fptr0 off len)
  | len == 0 = return (Nothing, "")
  | len == 1 = return (Nothing, bs)
  | otherwise = withForeignPtr fptr0 $ \ptr0 -> do
      let begptr = ptr0 `plusPtr` off
          limptr = begptr `plusPtr` len
      parseUrl' fptr0 ptr0 begptr limptr len

parseUrl' :: ForeignPtr Word8 -> Ptr Word8 -> Ptr Word8 -> Ptr Word8 -> Int
          -> IO (Maybe ByteString, URLPath)
parseUrl' fptr0 ptr0 begptr limptr len0 = do
      w0 <- peek begptr
      if w0 == _slash then do
          w1 <- peek $ begptr `plusPtr` 1
          if w1 == _slash  then
              doubleSlashed begptr len0
            else
              slashed begptr len0 Nothing
        else do
          colonptr <- memchr begptr _colon $ fromIntegral len0
          if colonptr == nullPtr then
              return (Nothing, "")
            else do
              let !authptr = colonptr `plusPtr` 1
              doubleSlashed authptr (limptr `minusPtr` authptr)
  where
    -- // / ?
    doubleSlashed :: Ptr Word8 -> Int -> IO (Maybe ByteString, URLPath)
    doubleSlashed ptr len
      | len < 2  = return (Nothing, "")
      | otherwise = do
          let ptr1 = ptr `plusPtr` 2
          pathptr <- memchr ptr1 _slash $ fromIntegral len
          if pathptr == nullPtr then
              return (Nothing, "")
            else do
              let !auth = bs ptr0 ptr1 pathptr
              slashed pathptr (limptr `minusPtr` pathptr) (Just auth)

    -- / ?
    slashed :: Ptr Word8 -> Int -> Maybe ByteString -> IO (Maybe ByteString, URLPath)
    slashed ptr len mauth = do
        questionptr <- memchr ptr _question $ fromIntegral len
        if questionptr == nullPtr then do
            let !path = bs ptr0 ptr limptr
            return (mauth, path)
          else do
            let !path = bs ptr0 ptr questionptr
            return (mauth, path)
    bs p0 p1 p2 = path
      where
        !off = p1 `minusPtr` p0
        !siz = p2 `minusPtr` p1
        !path = PS fptr0 off siz