Telegram bot: user whitelist access control

- Add tgAllowedUserIds field to TelegramConfig
- Load ALLOWED_TELEGRAM_USER_IDS from environment (comma-separated)
- Check isUserAllowed before processing messages
- Reject unauthorized users with friendly message
- Empty whitelist or '*' allows all users
- Add tests for whitelist behavior

1 files changed, 67 insertions, 12 deletions

diff --git a/Omni/Agent/Telegram.hs b/Omni/Agent/Telegram.hs
index 0c3a870..566377e 100644
--- a/Omni/Agent/Telegram.hs
+++ b/Omni/Agent/Telegram.hs
@@ -71,11 +71,22 @@ test =
               TelegramConfig
                 { tgBotToken = "test-token",
                   tgPollingTimeout = 30,
-                  tgApiBaseUrl = "https://api.telegram.org"
+                  tgApiBaseUrl = "https://api.telegram.org",
+                  tgAllowedUserIds = [123, 456]
                 }
         case Aeson.decode (Aeson.encode cfg) of
           Nothing -> Test.assertFailure "Failed to decode TelegramConfig"
-          Just decoded -> tgBotToken decoded Test.@=? "test-token",
+          Just decoded -> do
+            tgBotToken decoded Test.@=? "test-token"
+            tgAllowedUserIds decoded Test.@=? [123, 456],
+      Test.unit "isUserAllowed checks whitelist" <| do
+        let cfg = defaultTelegramConfig "token" [100, 200, 300]
+        isUserAllowed cfg 100 Test.@=? True
+        isUserAllowed cfg 200 Test.@=? True
+        isUserAllowed cfg 999 Test.@=? False,
+      Test.unit "isUserAllowed allows all when empty" <| do
+        let cfg = defaultTelegramConfig "token" []
+        isUserAllowed cfg 12345 Test.@=? True,
       Test.unit "TelegramMessage JSON roundtrip" <| do
         let msg =
               TelegramMessage
@@ -122,7 +133,8 @@ test =
 data TelegramConfig = TelegramConfig
   { tgBotToken :: Text,
     tgPollingTimeout :: Int,
-    tgApiBaseUrl :: Text
+    tgApiBaseUrl :: Text,
+    tgAllowedUserIds :: [Int]
   }
   deriving (Show, Eq, Generic)
 
@@ -131,7 +143,8 @@ instance Aeson.ToJSON TelegramConfig where
     Aeson.object
       [ "bot_token" .= tgBotToken c,
         "polling_timeout" .= tgPollingTimeout c,
-        "api_base_url" .= tgApiBaseUrl c
+        "api_base_url" .= tgApiBaseUrl c,
+        "allowed_user_ids" .= tgAllowedUserIds c
       ]
 
 instance Aeson.FromJSON TelegramConfig where
@@ -140,16 +153,23 @@ instance Aeson.FromJSON TelegramConfig where
       (TelegramConfig </ (v .: "bot_token"))
         <*> (v .:? "polling_timeout" .!= 30)
         <*> (v .:? "api_base_url" .!= "https://api.telegram.org")
+        <*> (v .:? "allowed_user_ids" .!= [])
 
 -- | Default Telegram configuration (requires token from env).
-defaultTelegramConfig :: Text -> TelegramConfig
-defaultTelegramConfig token =
+defaultTelegramConfig :: Text -> [Int] -> TelegramConfig
+defaultTelegramConfig token allowedIds =
   TelegramConfig
     { tgBotToken = token,
       tgPollingTimeout = 30,
-      tgApiBaseUrl = "https://api.telegram.org"
+      tgApiBaseUrl = "https://api.telegram.org",
+      tgAllowedUserIds = allowedIds
     }
 
+-- | Check if a user is allowed to use the bot.
+isUserAllowed :: TelegramConfig -> Int -> Bool
+isUserAllowed cfg usrId =
+  null (tgAllowedUserIds cfg) || usrId `elem` tgAllowedUserIds cfg
+
 -- | A parsed Telegram message from a user.
 data TelegramMessage = TelegramMessage
   { tmUpdateId :: Int,
@@ -371,16 +391,35 @@ handleMessage ::
   TelegramMessage ->
   IO ()
 handleMessage tgConfig provider engineCfg msg = do
-  sendTypingAction tgConfig (tmChatId msg)
-
   let userName =
         tmUserFirstName msg
           <> maybe "" (" " <>) (tmUserLastName msg)
       chatId = tmChatId msg
+      usrId = tmUserId msg
+
+  unless (isUserAllowed tgConfig usrId) <| do
+    putText <| "Unauthorized user: " <> tshow usrId <> " (" <> userName <> ")"
+    sendMessage tgConfig chatId "sorry, you're not authorized to use this bot."
+    pure ()
 
-  user <- Memory.getOrCreateUserByTelegramId (tmUserId msg) userName
-  let uid = Memory.userId user
+  when (isUserAllowed tgConfig usrId) <| do
+    sendTypingAction tgConfig chatId
 
+    user <- Memory.getOrCreateUserByTelegramId usrId userName
+    let uid = Memory.userId user
+
+    handleAuthorizedMessage tgConfig provider engineCfg msg uid userName chatId
+
+handleAuthorizedMessage ::
+  TelegramConfig ->
+  Provider.Provider ->
+  Engine.EngineConfig ->
+  TelegramMessage ->
+  Text ->
+  Text ->
+  Int ->
+  IO ()
+handleAuthorizedMessage tgConfig provider engineCfg msg uid userName chatId = do
   _ <- Memory.saveMessage uid chatId Memory.UserRole (tmText msg)
 
   (conversationContext, contextTokens) <- Memory.getConversationContext uid chatId maxConversationTokens
@@ -483,12 +522,28 @@ startBot maybeToken = do
           putText "Error: TELEGRAM_BOT_TOKEN not set and no --token provided"
           exitFailure
 
+  allowedIds <- loadAllowedUserIds
+
   apiKey <- lookupEnv "OPENROUTER_API_KEY"
   case apiKey of
     Nothing -> do
       putText "Error: OPENROUTER_API_KEY not set"
       exitFailure
     Just key -> do
-      let tgConfig = defaultTelegramConfig token
+      let tgConfig = defaultTelegramConfig token allowedIds
           provider = Provider.defaultOpenRouter (Text.pack key) "anthropic/claude-sonnet-4"
+      putText <| "Allowed user IDs: " <> tshow allowedIds
       runTelegramBot tgConfig provider
+
+-- | Load allowed user IDs from environment variable.
+-- Format: comma-separated integers, e.g. "123,456,789"
+-- Empty list means allow all users.
+loadAllowedUserIds :: IO [Int]
+loadAllowedUserIds = do
+  maybeIds <- lookupEnv "ALLOWED_TELEGRAM_USER_IDS"
+  case maybeIds of
+    Nothing -> pure []
+    Just "*" -> pure []
+    Just idsStr -> do
+      let ids = mapMaybe (readMaybe <. Text.unpack <. Text.strip) (Text.splitOn "," (Text.pack idsStr))
+      pure ids