1 files changed, 60 insertions, 13 deletions

diff --git a/Biz/Que/Host.hs b/Biz/Que/Host.hs
index 834ce0e..8d826b4 100755
--- a/Biz/Que/Host.hs
+++ b/Biz/Que/Host.hs
@@ -33,6 +33,7 @@ import qualified Control.Exception as Exception
 import Data.HashMap.Lazy (HashMap)
 import qualified Data.HashMap.Lazy as HashMap
 import Network.HTTP.Media ((//), (/:))
+import Network.Socket (SockAddr (..))
 import qualified Network.Wai.Handler.Warp as Warp
 import qualified Omni.Cli as Cli
 import qualified Omni.Log as Log
@@ -75,7 +76,30 @@ Usage:
 |]
 
 test :: Test.Tree
-test = Test.group "Biz.Que.Host" [Test.unit "id" <| 1 @=? (1 :: Integer)]
+test =
+  Test.group
+    "Biz.Que.Host"
+    [ Test.unit "id" <| 1 @=? (1 :: Integer),
+      Test.unit "putQue requires auth for '_'" <| do
+        st <- atomically <| STM.newTVar mempty
+        let cfg = Envy.defConfig
+        let handlers = paths cfg
+
+        -- Case 1: No auth, should fail
+        let nonLocalHost = SockAddrInet 0 0
+        let handler1 = putQue handlers nonLocalHost Nothing "_" "testq" "body"
+        res1 <- Servant.runHandler (runReaderT handler1 st)
+        case res1 of
+          Left err -> if errHTTPCode err == 401 then pure () else Test.assertFailure ("Expected 401, got " <> show err)
+          Right _ -> Test.assertFailure "Expected failure, got success"
+
+        -- Case 2: Correct auth, should succeed
+        let handler2 = putQue handlers nonLocalHost (Just "admin-key") "_" "testq" "body"
+        res2 <- Servant.runHandler (runReaderT handler2 st)
+        case res2 of
+          Left err -> Test.assertFailure (show err)
+          Right _ -> pure ()
+    ]
 
 type App = ReaderT AppState Servant.Handler
 
@@ -125,23 +149,31 @@ data Paths path = Paths
         :- Get '[JSON] NoContent,
     dash ::
       path
-        :- "_"
+        :- RemoteHost
+          :> Header "Authorization" Text
+          :> "_"
           :> "dash"
           :> Get '[JSON] Ques,
     getQue ::
       path
-        :- Capture "ns" Text
+        :- RemoteHost
+          :> Header "Authorization" Text
+          :> Capture "ns" Text
           :> Capture "quename" Text
           :> Get '[PlainText, HTML, OctetStream] Message,
     getStream ::
       path
-        :- Capture "ns" Text
+        :- RemoteHost
+          :> Header "Authorization" Text
+          :> Capture "ns" Text
           :> Capture "quename" Text
           :> "stream"
           :> StreamGet NoFraming OctetStream (SourceIO Message),
     putQue ::
       path
-        :- Capture "ns" Text
+        :- RemoteHost
+          :> Header "Authorization" Text
+          :> Capture "ns" Text
           :> Capture "quepath" Text
           :> ReqBody '[PlainText, HTML, OctetStream] Text
           :> Post '[PlainText, HTML, OctetStream] NoContent
@@ -149,15 +181,15 @@ data Paths path = Paths
   deriving (Generic)
 
 paths :: Config -> Paths (AsServerT App)
-paths _ =
-  -- TODO revive authkey stuff
-  -- - read Authorization header, compare with queSkey
-  -- - Only allow my IP or localhost to publish to '_' namespace
+paths Config {..} =
   Paths
     { home =
         throwError <| err301 {errHeaders = [("Location", "/_/index")]},
-      dash = gets,
-      getQue = \ns qn -> do
+      dash = \rh mAuth -> do
+        checkAuth queSkey rh mAuth "_"
+        gets,
+      getQue = \rh mAuth ns qn -> do
+        checkAuth queSkey rh mAuth ns
         guardNs ns ["pub", "_"]
         modify <| upsertNamespace ns
         q <- que ns qn
@@ -165,7 +197,8 @@ paths _ =
           |> liftIO
           +> Go.tap
           |> liftIO,
-      getStream = \ns qn -> do
+      getStream = \rh mAuth ns qn -> do
+        checkAuth queSkey rh mAuth ns
         guardNs ns ["pub", "_"]
         modify <| upsertNamespace ns
         q <- que ns qn
@@ -174,7 +207,8 @@ paths _ =
           +> Go.tap
           |> Source.fromAction (const False) -- peek chan instead of False?
           |> pure,
-      putQue = \ns qp body -> do
+      putQue = \rh mAuth ns qp body -> do
+        checkAuth queSkey rh mAuth ns
         guardNs ns ["pub", "_"]
         modify <| upsertNamespace ns
         q <- que ns qp
@@ -188,6 +222,19 @@ paths _ =
           >> pure NoContent
     }
 
+checkAuth :: Text -> SockAddr -> Maybe Text -> Text -> App ()
+checkAuth skey rh mAuth ns = do
+  let authorized = mAuth == Just skey
+  let isLocal = isLocalhost rh
+  when (ns == "_" && not (authorized || isLocal)) <| do
+    throwError err401 {errBody = "Authorized access only for '_' namespace"}
+
+isLocalhost :: SockAddr -> Bool
+isLocalhost (SockAddrInet _ h) = h == 0x0100007f -- 127.0.0.1
+isLocalhost (SockAddrInet6 _ _ (0, 0, 0, 1) _) = True -- ::1
+isLocalhost (SockAddrUnix _) = True
+isLocalhost _ = False
+
 -- | Given `guardNs ns whitelist`, if `ns` is not in the `whitelist`
 -- list, return a 405 error.
 guardNs :: (Applicative a, MonadError ServerError a) => Text -> [Text] -> a ()