From fb027dee588d709662556874ba22af44af2183bd Mon Sep 17 00:00:00 2001
From: Gregor Kleen <kleen@cip.ifi.lmu.de>
Date: Tue, 6 Aug 2019 08:31:22 +0200
Subject: [PATCH] feat(audit): take IP from header

---
 src/Audit.hs | 35 ++++++++++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/Audit.hs b/src/Audit.hs
index 8a058485c..a148da1c5 100644
--- a/src/Audit.hs
+++ b/src/Audit.hs
@@ -8,10 +8,13 @@ module Audit
 
 
 import Import.NoModel
+import Settings
 import Model
 import Database.Persist.Sql
 import Audit.Types
 
+import qualified Data.Text.Encoding as Text
+
 import Utils.Lens
 import qualified Network.Wai as Wai
 import qualified Network.Socket as Wai
@@ -27,13 +30,33 @@ data AuditRemoteException
 instance Exception AuditRemoteException
 
 
-getRemote :: (MonadHandler m, MonadThrow m) => m IP
+getRemote :: (MonadHandler m, MonadThrow m, HasAppSettings (HandlerSite m)) => m IP
 getRemote = do
+  ipFromHeader <- getsYesod $ view _appIpFromHeader
   wai <- waiRequest
-  case Wai.remoteHost wai of
-    Wai.SockAddrInet _ hAddr -> let (b1, b2, b3, b4) = Wai.hostAddressToTuple hAddr in return $ IP.ipv4 b1 b2 b3 b4
-    Wai.SockAddrInet6 _ _ hAddr _ -> let (w1, w2, w3, w4, w5, w6, w7, w8) = Wai.hostAddress6ToTuple hAddr in return $ IP.ipv6 w1 w2 w3 w4 w5 w6 w7 w8
-    _other -> throwM ARUnsupportedSocketKind
+
+  if
+    | ipFromHeader
+    , Just ip <- byHeader wai
+      -> return ip
+    | otherwise
+      -> byRemoteHost wai
+  
+  where
+    byHeader wai = listToMaybe $ do
+      (h, v) <- Wai.requestHeaders wai
+      guard $ h `elem` ["x-real-ip", "x-forwarded-for"]
+      v' <- either (const mzero) return $ Text.decodeUtf8' v
+      maybeToList $ IP.decode v'
+    
+    byRemoteHost wai = case Wai.remoteHost wai of
+      Wai.SockAddrInet _ hAddr
+        -> let (b1, b2, b3, b4) = Wai.hostAddressToTuple hAddr
+            in return $ IP.ipv4 b1 b2 b3 b4
+      Wai.SockAddrInet6 _ _ hAddr _
+        -> let (w1, w2, w3, w4, w5, w6, w7, w8) = Wai.hostAddress6ToTuple hAddr
+            in return $ IP.ipv6 w1 w2 w3 w4 w5 w6 w7 w8
+      _other -> throwM ARUnsupportedSocketKind
 
 
 data AuditException
@@ -51,6 +74,7 @@ audit :: ( AuthId (HandlerSite m) ~ Key User
          , YesodAuthPersist (HandlerSite m)
          , MonadHandler m
          , MonadCatch m
+         , HasAppSettings (HandlerSite m)
          )
       => Transaction -- ^ Transaction to record
       -> [UserId] -- ^ Affected users
@@ -81,6 +105,7 @@ audit' :: ( AuthId (HandlerSite m) ~ Key User
           , YesodAuthPersist (HandlerSite m)
           , MonadHandler m
           , MonadCatch m
+          , HasAppSettings (HandlerSite m)
           )
        => Transaction -- ^ Transaction to record
        -> ReaderT (YesodPersistBackend (HandlerSite m)) m ()