1use crate::config::BotConfig;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum TalkReason {
5 AllowedUser,
6 Oncall,
7 Peer,
8 AllowedChatGroup,
9 ChatGrant,
10 GlobalGrant,
11 Open,
12 None,
13}
14
15#[derive(Debug, Clone)]
16pub struct TalkEvaluation {
17 pub allowed: bool,
18 pub reason: TalkReason,
19 pub quota_key: Option<String>,
20}
21
22pub fn evaluate_talk(
23 bot: &BotConfig,
24 chat_id: &str,
25 sender_open_id: &str,
26 resolved_allowed_users: &[String],
27 peer_bot_open_ids: &[String],
28) -> TalkEvaluation {
29 let sender_is_allowed = resolved_allowed_users.iter().any(|id| id == sender_open_id);
30
31 if sender_is_allowed {
32 return TalkEvaluation {
33 allowed: true,
34 reason: TalkReason::AllowedUser,
35 quota_key: None,
36 };
37 }
38
39 let oncall_match = bot.oncall_chats.iter().any(|oc| oc.chat_id == chat_id);
40 if oncall_match {
41 return TalkEvaluation {
42 allowed: true,
43 reason: TalkReason::Oncall,
44 quota_key: None,
45 };
46 }
47
48 let peer_match = peer_bot_open_ids.iter().any(|id| id == sender_open_id);
49 if peer_match {
50 return TalkEvaluation {
51 allowed: true,
52 reason: TalkReason::Peer,
53 quota_key: None,
54 };
55 }
56
57 let chat_group_match = bot.allowed_chat_groups.iter().any(|cg| cg == chat_id);
58 if chat_group_match {
59 return TalkEvaluation {
60 allowed: true,
61 reason: TalkReason::AllowedChatGroup,
62 quota_key: None,
63 };
64 }
65
66 if let Some(granted) = bot.chat_grants.get(chat_id)
67 && granted.iter().any(|id| id == sender_open_id)
68 {
69 return TalkEvaluation {
70 allowed: true,
71 reason: TalkReason::ChatGrant,
72 quota_key: Some(format!("chat:{}:{}", chat_id, sender_open_id)),
73 };
74 }
75
76 if bot.global_grants.iter().any(|id| id == sender_open_id) {
77 return TalkEvaluation {
78 allowed: true,
79 reason: TalkReason::GlobalGrant,
80 quota_key: Some(format!("global:{}", sender_open_id)),
81 };
82 }
83
84 if bot.allowed_users.is_empty()
85 && bot.allowed_chat_groups.is_empty()
86 && bot.chat_grants.is_empty()
87 && bot.global_grants.is_empty()
88 && bot.oncall_chats.is_empty()
89 {
90 return TalkEvaluation {
91 allowed: true,
92 reason: TalkReason::Open,
93 quota_key: None,
94 };
95 }
96
97 TalkEvaluation {
98 allowed: false,
99 reason: TalkReason::None,
100 quota_key: None,
101 }
102}
103
104pub fn can_operate(
105 bot: &BotConfig,
106 sender_open_id: &str,
107 resolved_allowed_users: &[String],
108 peer_bot_open_ids: &[String],
109) -> bool {
110 let sender_is_allowed = resolved_allowed_users.iter().any(|id| id == sender_open_id);
111
112 if sender_is_allowed {
113 return true;
114 }
115
116 let peer_match = peer_bot_open_ids.iter().any(|id| id == sender_open_id);
117 if peer_match {
118 return true;
119 }
120
121 let has_allowlist = !bot.allowed_users.is_empty()
122 || !bot.allowed_chat_groups.is_empty()
123 || !bot.chat_grants.is_empty()
124 || !bot.global_grants.is_empty()
125 || !bot.oncall_chats.is_empty();
126
127 if !has_allowlist {
128 return true;
129 }
130
131 false
132}
133
134pub fn is_owner(open_id: &str, resolved_allowed_users: &[String]) -> bool {
135 resolved_allowed_users
136 .first()
137 .map(|owner| owner == open_id)
138 .unwrap_or(false)
139}
140
141pub fn get_owner_open_id(resolved_allowed_users: &[String]) -> Option<String> {
142 resolved_allowed_users.first().cloned()
143}
144
145pub fn grant_restricted(talk_eval: &TalkEvaluation, restrict_grant_commands: bool) -> bool {
146 if !restrict_grant_commands {
147 return false;
148 }
149 matches!(
150 talk_eval.reason,
151 TalkReason::ChatGrant | TalkReason::GlobalGrant
152 )
153}
154
155#[cfg(test)]
156mod tests {
157 use super::*;
158 use crate::config::{BotConfig, OncallChatBinding};
159 use std::collections::HashMap;
160
161 fn default_bot() -> BotConfig {
162 BotConfig {
163 name: None,
164 backend: None,
165 lark_app_id: "app-1".to_string(),
166 lark_app_secret: "secret".to_string(),
167 cli_id: "codex".to_string(),
168 cli_bin: None,
169 cgroup_slice: None,
170 cli_args: Vec::new(),
171 skip_working_dir_prompt: false,
172 model: None,
173 working_dir: None,
174 lark_encrypt_key: None,
175 lark_verification_token: None,
176 allowed_users: vec![],
177 private_card: false,
178 allowed_chat_groups: vec![],
179 chat_grants: HashMap::new(),
180 global_grants: vec![],
181 oncall_chats: vec![],
182 restrict_grant_commands: false,
183 message_quota: None,
184 quota_state: HashMap::new(),
185 custom_triggers: Vec::new(),
186 }
187 }
188
189 #[test]
190 fn evaluate_talk_open_when_no_restrictions() {
191 let bot = default_bot();
192 let eval = evaluate_talk(&bot, "chat-1", "ou_user", &[], &[]);
193 assert!(eval.allowed);
194 assert_eq!(eval.reason, TalkReason::Open);
195 }
196
197 #[test]
198 fn evaluate_talk_allowed_user() {
199 let bot = BotConfig {
200 allowed_users: vec!["ou_owner".to_string()],
201 ..default_bot()
202 };
203 let eval = evaluate_talk(&bot, "chat-1", "ou_owner", &["ou_owner".to_string()], &[]);
204 assert!(eval.allowed);
205 assert_eq!(eval.reason, TalkReason::AllowedUser);
206 }
207
208 #[test]
209 fn evaluate_talk_denies_non_allowed_user() {
210 let bot = BotConfig {
211 allowed_users: vec!["ou_owner".to_string()],
212 ..default_bot()
213 };
214 let eval = evaluate_talk(&bot, "chat-1", "ou_other", &["ou_owner".to_string()], &[]);
215 assert!(!eval.allowed);
216 assert_eq!(eval.reason, TalkReason::None);
217 }
218
219 #[test]
220 fn evaluate_talk_allowed_chat_group() {
221 let bot = BotConfig {
222 allowed_chat_groups: vec!["chat-open".to_string()],
223 ..default_bot()
224 };
225 let eval = evaluate_talk(&bot, "chat-open", "ou_anyone", &[], &[]);
226 assert!(eval.allowed);
227 assert_eq!(eval.reason, TalkReason::AllowedChatGroup);
228 }
229
230 #[test]
231 fn evaluate_talk_oncall_chat() {
232 let bot = BotConfig {
233 oncall_chats: vec![OncallChatBinding {
234 chat_id: "oncall-1".to_string(),
235 working_dir: None,
236 }],
237 ..default_bot()
238 };
239 let eval = evaluate_talk(&bot, "oncall-1", "ou_member", &[], &[]);
240 assert!(eval.allowed);
241 assert_eq!(eval.reason, TalkReason::Oncall);
242 }
243
244 #[test]
245 fn evaluate_talk_chat_grant_with_quota_key() {
246 let mut chat_grants = HashMap::new();
247 chat_grants.insert("chat-1".to_string(), vec!["ou_granted".to_string()]);
248 let bot = BotConfig {
249 allowed_users: vec!["ou_owner".to_string()],
250 chat_grants,
251 ..default_bot()
252 };
253 let eval = evaluate_talk(&bot, "chat-1", "ou_granted", &["ou_owner".to_string()], &[]);
254 assert!(eval.allowed);
255 assert_eq!(eval.reason, TalkReason::ChatGrant);
256 assert_eq!(eval.quota_key, Some("chat:chat-1:ou_granted".to_string()));
257 }
258
259 #[test]
260 fn evaluate_talk_global_grant() {
261 let bot = BotConfig {
262 allowed_users: vec!["ou_owner".to_string()],
263 global_grants: vec!["ou_global".to_string()],
264 ..default_bot()
265 };
266 let eval = evaluate_talk(
267 &bot,
268 "any-chat",
269 "ou_global",
270 &["ou_owner".to_string()],
271 &[],
272 );
273 assert!(eval.allowed);
274 assert_eq!(eval.reason, TalkReason::GlobalGrant);
275 assert_eq!(eval.quota_key, Some("global:ou_global".to_string()));
276 }
277
278 #[test]
279 fn can_operate_empty_allowlist() {
280 let bot = default_bot();
281 assert!(can_operate(&bot, "ou_any", &[], &[]));
282 }
283
284 #[test]
285 fn can_operate_respects_allowlist() {
286 let bot = BotConfig {
287 allowed_users: vec!["ou_owner".to_string()],
288 ..default_bot()
289 };
290 assert!(can_operate(
291 &bot,
292 "ou_owner",
293 &["ou_owner".to_string()],
294 &[]
295 ));
296 assert!(!can_operate(
297 &bot,
298 "ou_other",
299 &["ou_owner".to_string()],
300 &[]
301 ));
302 }
303
304 #[test]
305 fn can_operate_is_locked_by_talk_only_grants() {
306 let bot = BotConfig {
307 allowed_chat_groups: vec!["chat-open".to_string()],
308 ..default_bot()
309 };
310 assert!(!can_operate(&bot, "ou_any", &[], &[]));
311 }
312
313 #[test]
314 fn grant_restricted_blocks_chat_grant_when_enabled() {
315 let eval = TalkEvaluation {
316 allowed: true,
317 reason: TalkReason::ChatGrant,
318 quota_key: Some("chat:chat-1:ou_user".to_string()),
319 };
320 assert!(grant_restricted(&eval, true));
321 assert!(!grant_restricted(&eval, false));
322 }
323
324 #[test]
325 fn grant_restricted_allows_allowed_user_even_when_enabled() {
326 let eval = TalkEvaluation {
327 allowed: true,
328 reason: TalkReason::AllowedUser,
329 quota_key: None,
330 };
331 assert!(!grant_restricted(&eval, true));
332 }
333}