1 {-# LANGUAGE CPP               #-}
    2 {-# LANGUAGE OverloadedStrings #-}
    3 
    4 -- | This module provides facilities for patching incoming 'Requests' to
    5 -- correct the value of 'rqClientAddr' if the snap server is running behind a
    6 -- proxy.
    7 --
    8 -- Example usage:
    9 --
   10 -- @
   11 -- m :: Snap ()
   12 -- m = undefined  -- code goes here
   13 --
   14 -- applicationHandler :: Snap ()
   15 -- applicationHandler = behindProxy X_Forwarded_For m
   16 -- @
   17 --
   18 module Snap.Util.Proxy
   19   ( ProxyType(..)
   20   , behindProxy
   21   ) where
   22 
   23 ------------------------------------------------------------------------------
   24 import           Control.Applicative   (Alternative ((<|>)))
   25 import           Control.Arrow         (second)
   26 import qualified Data.ByteString.Char8 as S (break, breakEnd, drop, dropWhile, readInt, spanEnd)
   27 import           Data.Char             (isSpace)
   28 import           Data.Maybe            (fromJust)
   29 import           Snap.Core             (MonadSnap, Request (rqClientAddr, rqClientPort), getHeader, modifyRequest)
   30 #if !MIN_VERSION_base(4,8,0)
   31 import           Control.Applicative   ((<$>))
   32 #endif
   33 ------------------------------------------------------------------------------
   34 
   35 
   36 ------------------------------------------------------------------------------
   37 -- | What kind of proxy is this? Affects which headers 'behindProxy' pulls the
   38 -- original remote address from.
   39 --
   40 -- Currently only proxy servers that send @X-Forwarded-For@ or @Forwarded-For@
   41 -- are supported.
   42 data ProxyType = NoProxy          -- ^ no proxy, leave the request alone
   43                | X_Forwarded_For  -- ^ Use the @Forwarded-For@ or
   44                                   --   @X-Forwarded-For@ header
   45   deriving (Read, Show, Eq, Ord)
   46 
   47 
   48 ------------------------------------------------------------------------------
   49 -- | Rewrite 'rqClientAddr' if we're behind a proxy.
   50 --
   51 -- Example:
   52 --
   53 -- @
   54 -- ghci> :set -XOverloadedStrings
   55 -- ghci> import qualified "Data.Map" as M
   56 -- ghci> import qualified "Snap.Test" as T
   57 -- ghci> let r = T.get \"\/foo\" M.empty >> T.addHeader \"X-Forwarded-For\" \"1.2.3.4\"
   58 -- ghci> let h = 'Snap.Core.getsRequest' 'rqClientAddr' >>= 'Snap.Core.writeBS')
   59 -- ghci> T.runHandler r h
   60 -- HTTP\/1.1 200 OK
   61 -- server: Snap\/test
   62 -- date: Fri, 08 Aug 2014 14:32:29 GMT
   63 --
   64 -- 127.0.0.1
   65 -- ghci> T.runHandler r ('behindProxy' 'X_Forwarded_For' h)
   66 -- HTTP\/1.1 200 OK
   67 -- server: Snap\/test
   68 -- date: Fri, 08 Aug 2014 14:33:02 GMT
   69 --
   70 -- 1.2.3.4
   71 -- @
   72 behindProxy :: MonadSnap m => ProxyType -> m a -> m a
   73 behindProxy NoProxy         = id
   74 behindProxy X_Forwarded_For = ((modifyRequest xForwardedFor) >>)
   75 {-# INLINE behindProxy #-}
   76 
   77 
   78 ------------------------------------------------------------------------------
   79 xForwardedFor :: Request -> Request
   80 xForwardedFor req = req { rqClientAddr = ip
   81                         , rqClientPort = port
   82                         }
   83   where
   84     proxyString  = getHeader "Forwarded-For"   req <|>
   85                    getHeader "X-Forwarded-For" req <|>
   86                    Just (rqClientAddr req)
   87 
   88     proxyAddr    = trim . snd . S.breakEnd (== ',') . fromJust $ proxyString
   89 
   90     trim         = fst . S.spanEnd isSpace . S.dropWhile isSpace
   91 
   92     (ip,portStr) = second (S.drop 1) . S.break (== ':') $ proxyAddr
   93 
   94     port         = fromJust (fst <$> S.readInt portStr <|>
   95                              Just (rqClientPort req))
   96 {-# INLINE xForwardedFor #-}