mirror of
https://github.com/tinyhumansai/openhuman.git
synced 2026-07-27 21:08:00 +00:00
channels: migrate providers to tinychannels + generic ChannelHost boundary (whatsapp_web, telegram) (#4569)
This commit is contained in:
Generated
+24
-5
@@ -4424,7 +4424,6 @@ dependencies = [
|
||||
"schemars",
|
||||
"sentry",
|
||||
"serde",
|
||||
"serde-big-array",
|
||||
"serde_json",
|
||||
"serde_repr",
|
||||
"serde_yaml",
|
||||
@@ -4466,13 +4465,9 @@ dependencies = [
|
||||
"url",
|
||||
"urlencoding",
|
||||
"uuid 1.23.1",
|
||||
"wacore",
|
||||
"wait-timeout",
|
||||
"walkdir",
|
||||
"webpki-roots 1.0.7",
|
||||
"whatsapp-rust",
|
||||
"whatsapp-rust-tokio-transport",
|
||||
"whatsapp-rust-ureq-http-client",
|
||||
"whisper-rs",
|
||||
"windows-sys 0.61.2",
|
||||
"wiremock",
|
||||
@@ -6830,20 +6825,44 @@ name = "tinychannels"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-imap",
|
||||
"async-trait",
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"chrono",
|
||||
"directories",
|
||||
"futures",
|
||||
"futures-util",
|
||||
"hex",
|
||||
"hmac 0.12.1",
|
||||
"lettre",
|
||||
"mail-parser",
|
||||
"parking_lot",
|
||||
"prost",
|
||||
"rand 0.10.1",
|
||||
"reqwest 0.12.28",
|
||||
"rusqlite",
|
||||
"rustls",
|
||||
"rustls-pki-types",
|
||||
"schemars",
|
||||
"serde",
|
||||
"serde-big-array",
|
||||
"serde_json",
|
||||
"sha1",
|
||||
"sha2 0.10.9",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-tungstenite 0.29.0",
|
||||
"tracing",
|
||||
"url",
|
||||
"urlencoding",
|
||||
"uuid 1.23.1",
|
||||
"wacore",
|
||||
"webpki-roots 1.0.7",
|
||||
"whatsapp-rust",
|
||||
"whatsapp-rust-tokio-transport",
|
||||
"whatsapp-rust-ureq-http-client",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
+5
-10
@@ -246,16 +246,10 @@ coins-bip39 = "0.8"
|
||||
curve25519-dalek = { version = "4", default-features = false, features = ["alloc"] }
|
||||
|
||||
fantoccini = { version = "0.22.0", optional = true, default-features = false, features = ["rustls-tls"] }
|
||||
serde-big-array = { version = "0.5", optional = true }
|
||||
pdf-extract = "0.10"
|
||||
# WhatsApp Web — upstream `whatsapp-rust` 0.5. Its Diesel-backed sqlite-storage
|
||||
# feature links sqlite3 separately from rusqlite 0.40, so the TinyAgents 1.5
|
||||
# baseline compiles this provider against wacore's in-memory Backend until a
|
||||
# rusqlite-backed durable store lands.
|
||||
whatsapp-rust = { version = "0.5", optional = true, default-features = false, features = ["tokio-runtime"] }
|
||||
whatsapp-rust-tokio-transport = { version = "0.5", optional = true, default-features = false }
|
||||
whatsapp-rust-ureq-http-client = { version = "0.5", optional = true }
|
||||
wacore = { version = "0.5", optional = true, default-features = false }
|
||||
# The WhatsApp Web provider (and its `whatsapp-rust` / `wacore` / `serde-big-array`
|
||||
# stack) now lives in the tinychannels crate; the `whatsapp-web` feature forwards
|
||||
# to `tinychannels/whatsapp-web`.
|
||||
ppt-rs = "0.2.14"
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
@@ -335,7 +329,8 @@ peripheral-rpi = ["dep:rppal"]
|
||||
browser-native = ["dep:fantoccini"]
|
||||
fantoccini = ["browser-native"]
|
||||
landlock = ["sandbox-landlock"]
|
||||
whatsapp-web = ["dep:whatsapp-rust", "dep:whatsapp-rust-tokio-transport", "dep:whatsapp-rust-ureq-http-client", "dep:wacore", "serde-big-array"]
|
||||
# The WhatsApp Web provider now lives in tinychannels; forward to its feature.
|
||||
whatsapp-web = ["tinychannels/whatsapp-web"]
|
||||
# Exposes the destructive `openhuman.test_reset` RPC. Off by default; the E2E
|
||||
# build (app/scripts/e2e-build.sh) flips it on. Shipped binaries never have
|
||||
# this feature so the wipe RPC isn't even registered, let alone reachable.
|
||||
|
||||
Generated
+19
@@ -9179,20 +9179,39 @@ name = "tinychannels"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-imap",
|
||||
"async-trait",
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"chrono",
|
||||
"directories 6.0.0",
|
||||
"futures",
|
||||
"futures-util",
|
||||
"hex",
|
||||
"hmac 0.12.1",
|
||||
"lettre",
|
||||
"mail-parser",
|
||||
"parking_lot",
|
||||
"prost",
|
||||
"rand 0.10.1",
|
||||
"reqwest 0.12.28",
|
||||
"rusqlite",
|
||||
"rustls",
|
||||
"rustls-pki-types",
|
||||
"schemars 1.2.1",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha1",
|
||||
"sha2 0.10.9",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-tungstenite 0.29.0",
|
||||
"tracing",
|
||||
"url",
|
||||
"urlencoding",
|
||||
"uuid 1.23.1",
|
||||
"webpki-roots 1.0.7",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -227,7 +227,7 @@ pub(super) async fn run_autonomous(
|
||||
if let Some(thread_id) = session_thread_id.as_deref() {
|
||||
match &result {
|
||||
Ok(response) => {
|
||||
crate::openhuman::channels::providers::presentation::deliver_response(
|
||||
crate::openhuman::channels::providers::web::presentation::deliver_response(
|
||||
"system",
|
||||
thread_id,
|
||||
run_id,
|
||||
|
||||
@@ -64,13 +64,14 @@ pub async fn doctor_channels(config: Config) -> Result<()> {
|
||||
if let Some(ref dc) = config.channels_config.discord {
|
||||
channels.push((
|
||||
"Discord",
|
||||
Arc::new(DiscordChannel::new(
|
||||
Arc::new(DiscordChannel::with_http_client(
|
||||
dc.bot_token.clone(),
|
||||
dc.guild_id.clone(),
|
||||
dc.channel_id.clone(),
|
||||
dc.allowed_users.clone(),
|
||||
dc.listen_to_bots,
|
||||
dc.mention_only,
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.discord"),
|
||||
)),
|
||||
));
|
||||
}
|
||||
@@ -209,10 +210,11 @@ pub async fn doctor_channels(config: Config) -> Result<()> {
|
||||
if let Some(ref dt) = config.channels_config.dingtalk {
|
||||
channels.push((
|
||||
"DingTalk",
|
||||
Arc::new(DingTalkChannel::new(
|
||||
Arc::new(DingTalkChannel::with_http_client(
|
||||
dt.client_id.clone(),
|
||||
dt.client_secret.clone(),
|
||||
dt.allowed_users.clone(),
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.dingtalk"),
|
||||
)),
|
||||
));
|
||||
}
|
||||
@@ -220,10 +222,11 @@ pub async fn doctor_channels(config: Config) -> Result<()> {
|
||||
if let Some(ref qq) = config.channels_config.qq {
|
||||
channels.push((
|
||||
"QQ",
|
||||
Arc::new(QQChannel::new(
|
||||
Arc::new(QQChannel::with_http_client(
|
||||
qq.app_id.clone(),
|
||||
qq.app_secret.clone(),
|
||||
qq.allowed_users.clone(),
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.qq"),
|
||||
)),
|
||||
));
|
||||
}
|
||||
|
||||
@@ -0,0 +1,377 @@
|
||||
//! OpenHuman-side implementations of the portable `tinychannels::host`
|
||||
//! capability traits.
|
||||
//!
|
||||
//! Each adapter wraps existing OpenHuman internals (voice factory, inference
|
||||
//! ops, approval gate, conversation store, shutdown registry, web event bus)
|
||||
//! and exposes them through the portable, `Config`-free trait surface a ported
|
||||
//! channel provider consumes. See [`super::build_channel_host`].
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use parking_lot::Mutex;
|
||||
use tinychannels::host::{
|
||||
AllowlistStore, ApprovalDecision, ApprovalGate, ConversationMessage, ConversationStore,
|
||||
EventSink, LifecycleRegistry, ReactionDecision, ReactionGate, ReactionQuery, ShutdownHook,
|
||||
SpeechRequest, SpeechResult, SpeechSynthesizer, Transcriber, TranscriptionRequest,
|
||||
TranscriptionResult,
|
||||
};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
|
||||
const LOG_PREFIX: &str = "[channels_host]";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// LifecycleRegistry → core::shutdown
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Bridges [`LifecycleRegistry`] onto the process-global async shutdown hook
|
||||
/// registry. Providers register teardown that runs once on shutdown.
|
||||
pub struct CoreShutdownRegistry;
|
||||
|
||||
impl LifecycleRegistry for CoreShutdownRegistry {
|
||||
fn register_shutdown(&self, name: &str, hook: ShutdownHook) {
|
||||
let name = name.to_string();
|
||||
// `core::shutdown::register` takes a re-callable `Fn`, but our hook is a
|
||||
// one-shot `FnOnce`; guard it behind a take-once slot so a (theoretical)
|
||||
// second invocation is a no-op.
|
||||
let slot = Arc::new(Mutex::new(Some(hook)));
|
||||
crate::core::shutdown::register(move || {
|
||||
let slot = Arc::clone(&slot);
|
||||
let name = name.clone();
|
||||
async move {
|
||||
let taken = slot.lock().take();
|
||||
if let Some(hook) = taken {
|
||||
tracing::debug!("{LOG_PREFIX} running shutdown hook: {name}");
|
||||
hook().await;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Transcriber → voice STT factory
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Speech-to-text backed by the OpenHuman voice provider factory.
|
||||
pub struct VoiceTranscriber {
|
||||
pub config: Arc<Config>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Transcriber for VoiceTranscriber {
|
||||
fn name(&self) -> &str {
|
||||
"openhuman-voice"
|
||||
}
|
||||
|
||||
async fn transcribe(
|
||||
&self,
|
||||
request: TranscriptionRequest,
|
||||
) -> anyhow::Result<TranscriptionResult> {
|
||||
let provider = crate::openhuman::voice::effective_stt_provider(&self.config);
|
||||
tracing::debug!(
|
||||
"{LOG_PREFIX} transcribe provider={provider} bytes_b64={}",
|
||||
request.audio_base64.len()
|
||||
);
|
||||
// Empty model → factory substitutes DEFAULT_WHISPER_MODEL.
|
||||
let stt = crate::openhuman::voice::create_stt_provider(&provider, "", &self.config)?;
|
||||
let outcome = stt
|
||||
.transcribe(
|
||||
&self.config,
|
||||
&request.audio_base64,
|
||||
request.mime_type.as_deref(),
|
||||
request.file_name.as_deref(),
|
||||
request.language.as_deref(),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!(e))?;
|
||||
Ok(TranscriptionResult {
|
||||
text: outcome.value.text,
|
||||
language: request.language,
|
||||
duration_secs: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SpeechSynthesizer → voice reply_speech
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Text-to-speech backed by the hosted reply-speech synthesizer.
|
||||
pub struct VoiceSynthesizer {
|
||||
pub config: Arc<Config>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SpeechSynthesizer for VoiceSynthesizer {
|
||||
async fn synthesize(&self, request: SpeechRequest) -> anyhow::Result<SpeechResult> {
|
||||
tracing::debug!(
|
||||
"{LOG_PREFIX} synthesize chars={} voice={:?}",
|
||||
request.text.len(),
|
||||
request.voice
|
||||
);
|
||||
let opts = crate::openhuman::voice::reply_speech::ReplySpeechOptions {
|
||||
voice_id: request.voice,
|
||||
model_id: None,
|
||||
output_format: request.format,
|
||||
voice_settings: None,
|
||||
};
|
||||
let outcome = crate::openhuman::voice::reply_speech::synthesize_reply(
|
||||
&self.config,
|
||||
&request.text,
|
||||
&opts,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!(e))?;
|
||||
let value = outcome.value;
|
||||
let visemes = serde_json::to_value(&value.visemes).ok();
|
||||
Ok(SpeechResult {
|
||||
audio_base64: value.audio_base64,
|
||||
mime_type: value.audio_mime,
|
||||
visemes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ReactionGate → inference should_react
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Inference-driven reaction gate backed by the local-AI should-react op.
|
||||
pub struct InferenceReactionGate {
|
||||
pub config: Arc<Config>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ReactionGate for InferenceReactionGate {
|
||||
async fn should_react(&self, query: ReactionQuery) -> anyhow::Result<ReactionDecision> {
|
||||
// Honour the runtime gate: when the local model runtime is disabled we
|
||||
// never react (matches presentation's prior inline guard).
|
||||
if !self.config.local_ai.runtime_enabled {
|
||||
tracing::debug!("{LOG_PREFIX} should_react skipped (local runtime disabled)");
|
||||
return Ok(ReactionDecision::default());
|
||||
}
|
||||
let outcome = crate::openhuman::inference::ops::inference_should_react(
|
||||
&self.config,
|
||||
&query.message,
|
||||
&query.channel_type,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!(e))?;
|
||||
Ok(ReactionDecision {
|
||||
should_react: outcome.value.should_react,
|
||||
emoji: outcome.value.emoji,
|
||||
reason: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ApprovalGate → approval reply parsing
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Parses inbound approval replies via the OpenHuman approval gate. Raising
|
||||
/// interactive approvals stays host-internal (the tool gate), so only
|
||||
/// [`ApprovalGate::parse_reply`] is implemented.
|
||||
pub struct CoreApprovalGate;
|
||||
|
||||
impl ApprovalGate for CoreApprovalGate {
|
||||
fn parse_reply(&self, message: &str) -> Option<ApprovalDecision> {
|
||||
crate::openhuman::approval::parse_approval_reply(message).map(|decision| {
|
||||
use crate::openhuman::approval::ApprovalDecision as Core;
|
||||
match decision {
|
||||
Core::ApproveOnce => ApprovalDecision::Approve,
|
||||
Core::ApproveAlwaysForTool => {
|
||||
ApprovalDecision::Choice("approve_always_for_tool".to_string())
|
||||
}
|
||||
Core::Deny => ApprovalDecision::Deny,
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ConversationStore → memory_conversations
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Durable conversation history backed by the OpenHuman conversation store.
|
||||
pub struct ConversationHistoryStore {
|
||||
pub workspace_dir: std::path::PathBuf,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ConversationStore for ConversationHistoryStore {
|
||||
async fn history(
|
||||
&self,
|
||||
session_key: &str,
|
||||
limit: usize,
|
||||
) -> anyhow::Result<Vec<ConversationMessage>> {
|
||||
let messages = crate::openhuman::memory_conversations::get_messages(
|
||||
self.workspace_dir.clone(),
|
||||
session_key,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!(e))?;
|
||||
let start = messages.len().saturating_sub(limit);
|
||||
Ok(messages[start..]
|
||||
.iter()
|
||||
.map(|m| ConversationMessage {
|
||||
role: m.message_type.clone(),
|
||||
content: m.content.clone(),
|
||||
timestamp: None,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn append(&self, session_key: &str, message: ConversationMessage) -> anyhow::Result<()> {
|
||||
let now = chrono::Utc::now().to_rfc3339();
|
||||
// `append_message` requires the thread to exist; create-or-noop first.
|
||||
crate::openhuman::memory_conversations::ensure_thread(
|
||||
self.workspace_dir.clone(),
|
||||
crate::openhuman::memory_conversations::CreateConversationThread {
|
||||
id: session_key.to_string(),
|
||||
title: session_key.to_string(),
|
||||
created_at: now.clone(),
|
||||
parent_thread_id: None,
|
||||
labels: None,
|
||||
personality_id: None,
|
||||
},
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!(e))?;
|
||||
let stored = crate::openhuman::memory_conversations::ConversationMessage {
|
||||
id: uuid::Uuid::new_v4().to_string(),
|
||||
content: message.content,
|
||||
message_type: message.role.clone(),
|
||||
extra_metadata: serde_json::Value::Null,
|
||||
sender: message.role,
|
||||
created_at: now,
|
||||
};
|
||||
crate::openhuman::memory_conversations::append_message(
|
||||
self.workspace_dir.clone(),
|
||||
session_key,
|
||||
stored,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!(e))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// AllowlistStore → config.toml channel allowlist
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Persists newly-authorized identities into the on-disk channel allowlist,
|
||||
/// replicating Telegram's former `persist_allowed_identity` (load
|
||||
/// `~/.openhuman/config.toml`, append to the channel's `allowed_users`, save).
|
||||
pub struct ConfigAllowlistStore;
|
||||
|
||||
#[async_trait]
|
||||
impl AllowlistStore for ConfigAllowlistStore {
|
||||
async fn persist_allowed_identity(&self, channel: &str, identity: &str) -> anyhow::Result<()> {
|
||||
use anyhow::Context;
|
||||
let normalized = identity.trim().trim_start_matches('@').to_string();
|
||||
if normalized.is_empty() {
|
||||
anyhow::bail!("cannot persist empty identity");
|
||||
}
|
||||
|
||||
let home = directories::UserDirs::new()
|
||||
.map(|u| u.home_dir().to_path_buf())
|
||||
.context("could not find home directory")?;
|
||||
let openhuman_dir = home.join(".openhuman");
|
||||
let config_path = openhuman_dir.join("config.toml");
|
||||
let contents = tokio::fs::read_to_string(&config_path)
|
||||
.await
|
||||
.with_context(|| format!("failed to read config file: {}", config_path.display()))?;
|
||||
let mut config: Config =
|
||||
toml::from_str(&contents).context("failed to parse config.toml for allowlist")?;
|
||||
config.config_path = config_path;
|
||||
config.workspace_dir = openhuman_dir.join("workspace");
|
||||
|
||||
match channel {
|
||||
"telegram" => {
|
||||
let Some(telegram) = config.channels_config.telegram.as_mut() else {
|
||||
anyhow::bail!("telegram channel config is missing in config.toml");
|
||||
};
|
||||
if !telegram.allowed_users.iter().any(|u| u == &normalized) {
|
||||
telegram.allowed_users.push(normalized);
|
||||
config
|
||||
.save()
|
||||
.await
|
||||
.context("failed to persist allowlist to config.toml")?;
|
||||
}
|
||||
}
|
||||
other => anyhow::bail!("allowlist persist unsupported for channel '{other}'"),
|
||||
}
|
||||
tracing::debug!("{LOG_PREFIX} persisted allowed identity for channel={channel}");
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// EventSink → routes provider events to the right OpenHuman bus
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Routes provider events by `domain`:
|
||||
/// - `"web"` → the web channel's `WebChannelEvent` broadcast bus (payload
|
||||
/// must deserialize into a `WebChannelEvent`; presentation builds that shape).
|
||||
/// - `"channel"` → the global `DomainEvent` bus (telegram reaction fan-out).
|
||||
///
|
||||
/// One capability, two backends — providers don't know which bus they hit.
|
||||
pub struct OpenHumanEventSink;
|
||||
|
||||
fn json_str(payload: &serde_json::Value, key: &str) -> String {
|
||||
payload
|
||||
.get(key)
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or_default()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl EventSink for OpenHumanEventSink {
|
||||
async fn publish(
|
||||
&self,
|
||||
domain: &str,
|
||||
kind: &str,
|
||||
payload: serde_json::Value,
|
||||
) -> anyhow::Result<()> {
|
||||
match domain {
|
||||
"web" => {
|
||||
let event: crate::core::socketio::WebChannelEvent = serde_json::from_value(payload)
|
||||
.map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"{LOG_PREFIX} web event payload not a WebChannelEvent ({kind}): {e}"
|
||||
)
|
||||
})?;
|
||||
crate::openhuman::channels::providers::web::publish_web_channel_event(event);
|
||||
}
|
||||
"channel" => {
|
||||
use crate::core::event_bus::{publish_global, DomainEvent};
|
||||
let event = match kind {
|
||||
"reaction_received" => DomainEvent::ChannelReactionReceived {
|
||||
channel: json_str(&payload, "channel"),
|
||||
sender: json_str(&payload, "sender"),
|
||||
target_message_id: json_str(&payload, "target_message_id"),
|
||||
emoji: json_str(&payload, "emoji"),
|
||||
},
|
||||
"reaction_sent" => DomainEvent::ChannelReactionSent {
|
||||
channel: json_str(&payload, "channel"),
|
||||
target_message_id: json_str(&payload, "target_message_id"),
|
||||
emoji: json_str(&payload, "emoji"),
|
||||
success: payload
|
||||
.get("success")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false),
|
||||
},
|
||||
other => {
|
||||
tracing::warn!("{LOG_PREFIX} unmapped channel event kind: {other}");
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
publish_global(event);
|
||||
}
|
||||
other => tracing::warn!("{LOG_PREFIX} unmapped event domain: {other}"),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
//! OpenHuman implementation of the `tinychannels::host` capability boundary.
|
||||
//!
|
||||
//! [`build_channel_host`] assembles the concrete [`tinychannels::ChannelHost`]
|
||||
//! from OpenHuman internals (voice, inference, approvals, conversation store,
|
||||
//! shutdown registry, web event bus). [`build_provider_context`] wraps it into
|
||||
//! the [`tinychannels::host::ProviderContext`] handed to channel providers.
|
||||
//!
|
||||
//! Ported providers reach host capabilities through this context instead of
|
||||
//! calling OpenHuman internals directly — the inversion that lets them live in
|
||||
//! the standalone `tinychannels` crate. Lean providers ignore the host.
|
||||
|
||||
mod adapters;
|
||||
|
||||
pub use adapters::{
|
||||
ConfigAllowlistStore, ConversationHistoryStore, CoreApprovalGate, CoreShutdownRegistry,
|
||||
InferenceReactionGate, OpenHumanEventSink, VoiceSynthesizer, VoiceTranscriber,
|
||||
};
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use tinychannels::host::{ChannelHostBuilder, ProviderContext};
|
||||
use tinychannels::ChannelHost;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
|
||||
/// Assemble the full OpenHuman [`ChannelHost`] from a config snapshot.
|
||||
///
|
||||
/// Wires every capability OpenHuman can back today: lifecycle (shutdown),
|
||||
/// STT, TTS, reaction gate, approval-reply parsing, conversation history, and
|
||||
/// the web-channel event sink. Capabilities OpenHuman cannot yet express
|
||||
/// portably (turn dispatch, run ledger, pairing) are simply left unset — a
|
||||
/// provider that needs one degrades gracefully.
|
||||
pub fn build_channel_host(config: Arc<Config>) -> Arc<dyn ChannelHost> {
|
||||
ChannelHostBuilder::new()
|
||||
.lifecycle(Arc::new(CoreShutdownRegistry))
|
||||
.transcriber(Arc::new(VoiceTranscriber {
|
||||
config: Arc::clone(&config),
|
||||
}))
|
||||
.synthesizer(Arc::new(VoiceSynthesizer {
|
||||
config: Arc::clone(&config),
|
||||
}))
|
||||
.reactions(Arc::new(InferenceReactionGate {
|
||||
config: Arc::clone(&config),
|
||||
}))
|
||||
.approvals(Arc::new(CoreApprovalGate))
|
||||
.conversations(Arc::new(ConversationHistoryStore {
|
||||
workspace_dir: config.workspace_dir.clone(),
|
||||
}))
|
||||
.events(Arc::new(OpenHumanEventSink))
|
||||
.allowlist(Arc::new(ConfigAllowlistStore))
|
||||
.build()
|
||||
}
|
||||
|
||||
/// Build the [`ProviderContext`] handed to a channel provider at construction:
|
||||
/// the assembled host + the channels config + a pre-built HTTP client.
|
||||
pub fn build_provider_context(config: &Config, http_client: reqwest::Client) -> ProviderContext {
|
||||
ProviderContext::new(
|
||||
build_channel_host(Arc::new(config.clone())),
|
||||
config.channels_config.clone(),
|
||||
http_client,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
@@ -0,0 +1,182 @@
|
||||
//! Unit tests for the OpenHuman ChannelHost capability adapters.
|
||||
|
||||
use super::*;
|
||||
use std::sync::Arc;
|
||||
use tinychannels::host::{
|
||||
ApprovalDecision, ApprovalGate, ConversationMessage, ConversationStore, EventSink,
|
||||
ReactionGate, ReactionQuery, Transcriber,
|
||||
};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
|
||||
// --- ApprovalGate::parse_reply -------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn approval_gate_maps_core_replies() {
|
||||
let gate = CoreApprovalGate;
|
||||
assert_eq!(gate.parse_reply("yes"), Some(ApprovalDecision::Approve));
|
||||
assert_eq!(gate.parse_reply("APPROVE"), Some(ApprovalDecision::Approve));
|
||||
assert_eq!(gate.parse_reply("no"), Some(ApprovalDecision::Deny));
|
||||
assert_eq!(gate.parse_reply("deny"), Some(ApprovalDecision::Deny));
|
||||
assert_eq!(gate.parse_reply("banana"), None);
|
||||
}
|
||||
|
||||
// --- ReactionGate (runtime-disabled short-circuit) ------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn reaction_gate_returns_default_when_runtime_disabled() {
|
||||
let mut config = Config::default();
|
||||
config.local_ai.runtime_enabled = false;
|
||||
let gate = InferenceReactionGate {
|
||||
config: Arc::new(config),
|
||||
};
|
||||
let decision = gate
|
||||
.should_react(ReactionQuery {
|
||||
message: "hello".into(),
|
||||
channel_type: "web".into(),
|
||||
})
|
||||
.await
|
||||
.expect("should_react ok");
|
||||
assert!(!decision.should_react);
|
||||
assert!(decision.emoji.is_none());
|
||||
}
|
||||
|
||||
// --- OpenHumanEventSink --------------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn event_sink_accepts_web_channel_event_shape() {
|
||||
let sink = OpenHumanEventSink;
|
||||
let ok = sink
|
||||
.publish(
|
||||
"web",
|
||||
"chat_done",
|
||||
serde_json::json!({
|
||||
"event": "chat_done",
|
||||
"client_id": "c1",
|
||||
"thread_id": "t1",
|
||||
"request_id": "r1",
|
||||
"full_response": "hi"
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert!(ok.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn event_sink_rejects_non_object_payload() {
|
||||
let sink = OpenHumanEventSink;
|
||||
let err = sink
|
||||
.publish("web", "bad", serde_json::json!([1, 2, 3]))
|
||||
.await;
|
||||
assert!(err.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn event_sink_routes_channel_reactions_to_domain_bus() {
|
||||
let sink = OpenHumanEventSink;
|
||||
assert!(sink
|
||||
.publish(
|
||||
"channel",
|
||||
"reaction_received",
|
||||
serde_json::json!({
|
||||
"channel": "telegram", "sender": "u1",
|
||||
"target_message_id": "telegram_1_2", "emoji": "👍"
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.is_ok());
|
||||
assert!(sink
|
||||
.publish(
|
||||
"channel",
|
||||
"reaction_sent",
|
||||
serde_json::json!({
|
||||
"channel": "telegram", "target_message_id": "telegram_1_2",
|
||||
"emoji": "👍", "success": true
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.is_ok());
|
||||
// Unknown domain/kind is a benign no-op.
|
||||
assert!(sink
|
||||
.publish("mystery", "x", serde_json::json!({}))
|
||||
.await
|
||||
.is_ok());
|
||||
assert!(sink
|
||||
.publish("channel", "unknown", serde_json::json!({}))
|
||||
.await
|
||||
.is_ok());
|
||||
}
|
||||
|
||||
// --- ConversationHistoryStore --------------------------------------------
|
||||
|
||||
#[tokio::test]
|
||||
async fn conversation_store_append_then_history_roundtrips() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let store = ConversationHistoryStore {
|
||||
workspace_dir: dir.path().to_path_buf(),
|
||||
};
|
||||
let thread = "thread-1";
|
||||
|
||||
store
|
||||
.append(
|
||||
thread,
|
||||
ConversationMessage {
|
||||
role: "user".into(),
|
||||
content: "first".into(),
|
||||
timestamp: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("append user");
|
||||
store
|
||||
.append(
|
||||
thread,
|
||||
ConversationMessage {
|
||||
role: "assistant".into(),
|
||||
content: "second".into(),
|
||||
timestamp: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("append assistant");
|
||||
|
||||
let history = store.history(thread, 10).await.expect("history");
|
||||
assert_eq!(history.len(), 2);
|
||||
assert_eq!(history[0].content, "first");
|
||||
assert_eq!(history[0].role, "user");
|
||||
assert_eq!(history[1].content, "second");
|
||||
assert_eq!(history[1].role, "assistant");
|
||||
|
||||
// `limit` keeps the most recent messages.
|
||||
let recent = store.history(thread, 1).await.expect("history limited");
|
||||
assert_eq!(recent.len(), 1);
|
||||
assert_eq!(recent[0].content, "second");
|
||||
}
|
||||
|
||||
// --- build_channel_host ---------------------------------------------------
|
||||
|
||||
#[test]
|
||||
fn build_channel_host_advertises_expected_capabilities() {
|
||||
let host = build_channel_host(Arc::new(Config::default()));
|
||||
let caps = host.capabilities();
|
||||
assert!(caps.lifecycle);
|
||||
assert!(caps.stt);
|
||||
assert!(caps.tts);
|
||||
assert!(caps.reaction_gate);
|
||||
assert!(caps.approvals);
|
||||
assert!(caps.conversation_store);
|
||||
assert!(caps.event_sink);
|
||||
assert!(caps.allowlist_store);
|
||||
// Not yet backed portably:
|
||||
assert!(!caps.turn_dispatch);
|
||||
assert!(!caps.run_ledger);
|
||||
assert!(!caps.memory_recall);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transcriber_reports_stable_name() {
|
||||
let stt = VoiceTranscriber {
|
||||
config: Arc::new(Config::default()),
|
||||
};
|
||||
assert_eq!(stt.name(), "openhuman-voice");
|
||||
}
|
||||
@@ -3,6 +3,7 @@
|
||||
pub mod bus;
|
||||
pub mod cli;
|
||||
pub mod controllers;
|
||||
pub mod host;
|
||||
pub mod proactive;
|
||||
pub mod providers;
|
||||
pub(crate) mod relay_runtime;
|
||||
|
||||
@@ -1,384 +1 @@
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
use async_trait::async_trait;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use uuid::Uuid;
|
||||
|
||||
const DINGTALK_BOT_CALLBACK_TOPIC: &str = "/v1.0/im/bot/messages/get";
|
||||
|
||||
/// DingTalk channel — connects via Stream Mode WebSocket for real-time messages.
|
||||
/// Replies are sent through per-message session webhook URLs.
|
||||
pub struct DingTalkChannel {
|
||||
client_id: String,
|
||||
client_secret: String,
|
||||
allowed_users: Vec<String>,
|
||||
/// Per-chat session webhooks for sending replies (chatID -> webhook URL).
|
||||
/// DingTalk provides a unique webhook URL with each incoming message.
|
||||
session_webhooks: Arc<RwLock<HashMap<String, String>>>,
|
||||
}
|
||||
|
||||
/// Response from DingTalk gateway connection registration.
|
||||
#[derive(serde::Deserialize)]
|
||||
struct GatewayResponse {
|
||||
endpoint: String,
|
||||
ticket: String,
|
||||
}
|
||||
|
||||
impl DingTalkChannel {
|
||||
pub fn new(client_id: String, client_secret: String, allowed_users: Vec<String>) -> Self {
|
||||
Self {
|
||||
client_id,
|
||||
client_secret,
|
||||
allowed_users,
|
||||
session_webhooks: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
fn http_client(&self) -> reqwest::Client {
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.dingtalk")
|
||||
}
|
||||
|
||||
fn is_user_allowed(&self, user_id: &str) -> bool {
|
||||
self.allowed_users.iter().any(|u| u == "*" || u == user_id)
|
||||
}
|
||||
|
||||
fn parse_stream_data(frame: &serde_json::Value) -> Option<serde_json::Value> {
|
||||
match frame.get("data") {
|
||||
Some(serde_json::Value::String(raw)) => serde_json::from_str(raw).ok(),
|
||||
Some(serde_json::Value::Object(_)) => frame.get("data").cloned(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_chat_id(data: &serde_json::Value, sender_id: &str) -> String {
|
||||
let is_private_chat = data
|
||||
.get("conversationType")
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_str()
|
||||
.map(|v| v == "1")
|
||||
.or_else(|| value.as_i64().map(|v| v == 1))
|
||||
})
|
||||
.unwrap_or(true);
|
||||
|
||||
if is_private_chat {
|
||||
sender_id.to_string()
|
||||
} else {
|
||||
data.get("conversationId")
|
||||
.and_then(|c| c.as_str())
|
||||
.unwrap_or(sender_id)
|
||||
.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a connection with DingTalk's gateway to get a WebSocket endpoint.
|
||||
async fn register_connection(&self) -> anyhow::Result<GatewayResponse> {
|
||||
let body = serde_json::json!({
|
||||
"clientId": self.client_id,
|
||||
"clientSecret": self.client_secret,
|
||||
"subscriptions": [
|
||||
{
|
||||
"type": "CALLBACK",
|
||||
"topic": DINGTALK_BOT_CALLBACK_TOPIC,
|
||||
}
|
||||
],
|
||||
});
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post("https://api.dingtalk.com/v1.0/gateway/connections/open")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let err = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!("DingTalk gateway registration failed ({status}): {err}");
|
||||
}
|
||||
|
||||
let gw: GatewayResponse = resp.json().await?;
|
||||
Ok(gw)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for DingTalkChannel {
|
||||
fn name(&self) -> &str {
|
||||
"dingtalk"
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> anyhow::Result<()> {
|
||||
let webhooks = self.session_webhooks.read().await;
|
||||
let webhook_url = webhooks.get(&message.recipient).ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"No session webhook found for chat {}. \
|
||||
The user must send a message first to establish a session.",
|
||||
message.recipient
|
||||
)
|
||||
})?;
|
||||
|
||||
let title = message.subject.as_deref().unwrap_or("OpenHuman");
|
||||
let body = serde_json::json!({
|
||||
"msgtype": "markdown",
|
||||
"markdown": {
|
||||
"title": title,
|
||||
"text": message.content,
|
||||
}
|
||||
});
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(webhook_url)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let err = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!("DingTalk webhook reply failed ({status}): {err}");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn listen(&self, tx: tokio::sync::mpsc::Sender<ChannelMessage>) -> anyhow::Result<()> {
|
||||
tracing::info!("DingTalk: registering gateway connection...");
|
||||
|
||||
let gw = self.register_connection().await?;
|
||||
let ws_url = format!("{}?ticket={}", gw.endpoint, gw.ticket);
|
||||
|
||||
tracing::info!("DingTalk: connecting to stream WebSocket...");
|
||||
let (ws_stream, _) = tokio_tungstenite::connect_async(&ws_url).await?;
|
||||
let (mut write, mut read) = ws_stream.split();
|
||||
|
||||
tracing::info!("DingTalk: connected and listening for messages...");
|
||||
|
||||
while let Some(msg) = read.next().await {
|
||||
let msg = match msg {
|
||||
Ok(Message::Text(t)) => t,
|
||||
Ok(Message::Close(_)) => break,
|
||||
Err(e) => {
|
||||
tracing::warn!("DingTalk WebSocket error: {e}");
|
||||
break;
|
||||
}
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let frame: serde_json::Value = match serde_json::from_str(msg.as_ref()) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let frame_type = frame.get("type").and_then(|t| t.as_str()).unwrap_or("");
|
||||
|
||||
match frame_type {
|
||||
"SYSTEM" => {
|
||||
// Respond to system pings to keep the connection alive
|
||||
let message_id = frame
|
||||
.get("headers")
|
||||
.and_then(|h| h.get("messageId"))
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
let pong = serde_json::json!({
|
||||
"code": 200,
|
||||
"headers": {
|
||||
"contentType": "application/json",
|
||||
"messageId": message_id,
|
||||
},
|
||||
"message": "OK",
|
||||
"data": "",
|
||||
});
|
||||
|
||||
if let Err(e) = write.send(Message::Text(pong.to_string())).await {
|
||||
tracing::warn!("DingTalk: failed to send pong: {e}");
|
||||
break;
|
||||
}
|
||||
}
|
||||
"EVENT" | "CALLBACK" => {
|
||||
// Parse the chatbot callback data from the frame.
|
||||
let data = match Self::parse_stream_data(&frame) {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
tracing::debug!("DingTalk: frame has no parseable data payload");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// Extract message content
|
||||
let content = data
|
||||
.get("text")
|
||||
.and_then(|t| t.get("content"))
|
||||
.and_then(|c| c.as_str())
|
||||
.unwrap_or("")
|
||||
.trim();
|
||||
|
||||
if content.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let sender_id = data
|
||||
.get("senderStaffId")
|
||||
.and_then(|s| s.as_str())
|
||||
.unwrap_or("unknown");
|
||||
|
||||
if !self.is_user_allowed(sender_id) {
|
||||
tracing::warn!(
|
||||
"DingTalk: ignoring message from unauthorized user: {sender_id}"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Private chat uses sender ID, group chat uses conversation ID.
|
||||
let chat_id = Self::resolve_chat_id(&data, sender_id);
|
||||
|
||||
// Store session webhook for later replies
|
||||
if let Some(webhook) = data.get("sessionWebhook").and_then(|w| w.as_str()) {
|
||||
let webhook = webhook.to_string();
|
||||
let mut webhooks = self.session_webhooks.write().await;
|
||||
// Use both keys so reply routing works for both group and private flows.
|
||||
webhooks.insert(chat_id.clone(), webhook.clone());
|
||||
webhooks.insert(sender_id.to_string(), webhook);
|
||||
}
|
||||
|
||||
// Acknowledge the event
|
||||
let message_id = frame
|
||||
.get("headers")
|
||||
.and_then(|h| h.get("messageId"))
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
let ack = serde_json::json!({
|
||||
"code": 200,
|
||||
"headers": {
|
||||
"contentType": "application/json",
|
||||
"messageId": message_id,
|
||||
},
|
||||
"message": "OK",
|
||||
"data": "",
|
||||
});
|
||||
let _ = write.send(Message::Text(ack.to_string())).await;
|
||||
|
||||
let channel_msg = ChannelMessage {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
sender: sender_id.to_string(),
|
||||
reply_target: chat_id,
|
||||
content: content.to_string(),
|
||||
channel: "dingtalk".to_string(),
|
||||
timestamp: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
thread_ts: None,
|
||||
};
|
||||
|
||||
if tx.send(channel_msg).await.is_err() {
|
||||
tracing::warn!("DingTalk: message channel closed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
anyhow::bail!("DingTalk WebSocket stream ended")
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
self.register_connection().await.is_ok()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_name() {
|
||||
let ch = DingTalkChannel::new("id".into(), "secret".into(), vec![]);
|
||||
assert_eq!(ch.name(), "dingtalk");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_allowed_wildcard() {
|
||||
let ch = DingTalkChannel::new("id".into(), "secret".into(), vec!["*".into()]);
|
||||
assert!(ch.is_user_allowed("anyone"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_allowed_specific() {
|
||||
let ch = DingTalkChannel::new("id".into(), "secret".into(), vec!["user123".into()]);
|
||||
assert!(ch.is_user_allowed("user123"));
|
||||
assert!(!ch.is_user_allowed("other"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_denied_empty() {
|
||||
let ch = DingTalkChannel::new("id".into(), "secret".into(), vec![]);
|
||||
assert!(!ch.is_user_allowed("anyone"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_serde() {
|
||||
let toml_str = r#"
|
||||
client_id = "app_id_123"
|
||||
client_secret = "secret_456"
|
||||
allowed_users = ["user1", "*"]
|
||||
"#;
|
||||
let config: crate::openhuman::config::schema::DingTalkConfig =
|
||||
toml::from_str(toml_str).unwrap();
|
||||
assert_eq!(config.client_id, "app_id_123");
|
||||
assert_eq!(config.client_secret, "secret_456");
|
||||
assert_eq!(config.allowed_users, vec!["user1", "*"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_serde_defaults() {
|
||||
let toml_str = r#"
|
||||
client_id = "id"
|
||||
client_secret = "secret"
|
||||
"#;
|
||||
let config: crate::openhuman::config::schema::DingTalkConfig =
|
||||
toml::from_str(toml_str).unwrap();
|
||||
assert!(config.allowed_users.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_stream_data_supports_string_payload() {
|
||||
let frame = serde_json::json!({
|
||||
"data": "{\"text\":{\"content\":\"hello\"}}"
|
||||
});
|
||||
let parsed = DingTalkChannel::parse_stream_data(&frame).unwrap();
|
||||
assert_eq!(
|
||||
parsed.get("text").and_then(|v| v.get("content")),
|
||||
Some(&serde_json::json!("hello"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_stream_data_supports_object_payload() {
|
||||
let frame = serde_json::json!({
|
||||
"data": {"text": {"content": "hello"}}
|
||||
});
|
||||
let parsed = DingTalkChannel::parse_stream_data(&frame).unwrap();
|
||||
assert_eq!(
|
||||
parsed.get("text").and_then(|v| v.get("content")),
|
||||
Some(&serde_json::json!("hello"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_chat_id_handles_numeric_group_conversation_type() {
|
||||
let data = serde_json::json!({
|
||||
"conversationType": 2,
|
||||
"conversationId": "cid-group",
|
||||
});
|
||||
let chat_id = DingTalkChannel::resolve_chat_id(&data, "staff-1");
|
||||
assert_eq!(chat_id, "cid-group");
|
||||
}
|
||||
}
|
||||
pub use tinychannels::providers::dingtalk::*;
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
pub use tinychannels::providers::discord::*;
|
||||
@@ -1,511 +0,0 @@
|
||||
//! Discord REST API helpers for guild/channel discovery and permission checks.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
const DISCORD_API_BASE: &str = "https://discord.com/api/v10";
|
||||
|
||||
/// Minimal guild (server) info returned by `GET /users/@me/guilds`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DiscordGuild {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub icon: Option<String>,
|
||||
}
|
||||
|
||||
/// Minimal channel info returned by `GET /guilds/{guild_id}/channels`.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DiscordTextChannel {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
/// Discord channel type — 0 = text, 2 = voice, 4 = category, etc.
|
||||
#[serde(rename = "type")]
|
||||
pub channel_type: u64,
|
||||
#[serde(default)]
|
||||
pub position: u64,
|
||||
/// Parent category ID (if nested under a category).
|
||||
pub parent_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Result of a bot permission check for a given channel.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct BotPermissionCheck {
|
||||
pub can_view_channel: bool,
|
||||
pub can_send_messages: bool,
|
||||
pub can_read_message_history: bool,
|
||||
pub missing_permissions: Vec<String>,
|
||||
}
|
||||
|
||||
// Discord permission flag bits
|
||||
const VIEW_CHANNEL: u64 = 1 << 10; // 0x400
|
||||
const SEND_MESSAGES: u64 = 1 << 11; // 0x800
|
||||
const READ_MESSAGE_HISTORY: u64 = 1 << 16; // 0x10000
|
||||
|
||||
fn build_client() -> reqwest::Client {
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.discord")
|
||||
}
|
||||
|
||||
fn auth_header(token: &str) -> String {
|
||||
format!("Bot {token}")
|
||||
}
|
||||
|
||||
/// Format a non-2xx response from the Discord REST API as a string
|
||||
/// suitable for a JSON-RPC error result.
|
||||
///
|
||||
/// **Load-bearing for issue #2285**: the global JSON-RPC dispatcher
|
||||
/// at `src/core/jsonrpc.rs::is_session_expired_error` matches any
|
||||
/// error string that contains `"401"` AND `"unauthorized"` as a
|
||||
/// signal that the OpenHuman backend session has expired, and
|
||||
/// publishes a `DomainEvent::SessionExpired` event that signs the
|
||||
/// user out. A raw upstream Discord 401 (revoked bot token) would
|
||||
/// previously trip that classifier — opening the connected-Discord
|
||||
/// card on the Channels page logged the user out of OpenHuman over
|
||||
/// a *Discord* credentials problem.
|
||||
///
|
||||
/// The fix is to convert auth failures here into a Discord-domain
|
||||
/// message that:
|
||||
///
|
||||
/// 1. Does NOT contain both `"401"` and `"unauthorized"` as a pair
|
||||
/// (so the global classifier ignores it).
|
||||
/// 2. Tells the user the actual remediation: rotate the bot token
|
||||
/// at `Settings → Channels → Discord`.
|
||||
/// 3. Preserves the originating endpoint identifier in the message
|
||||
/// so triage can still see WHICH Discord call failed without
|
||||
/// plumbing a separate error code.
|
||||
///
|
||||
/// Other non-2xx statuses (400 / 404 / 5xx) pass through with a
|
||||
/// `Discord API error` prefix — they don't match the
|
||||
/// `is_session_expired_error` predicate even when verbatim.
|
||||
fn format_discord_http_error(endpoint: &str, status: reqwest::StatusCode, body: &str) -> String {
|
||||
if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
|
||||
let kind = if status == reqwest::StatusCode::UNAUTHORIZED {
|
||||
"bot token was rejected"
|
||||
} else {
|
||||
"bot token lacks required Discord permissions"
|
||||
};
|
||||
// Spell out the HTTP code so `lower.contains("401")` does NOT
|
||||
// match — see #2285 rationale on this helper. Also avoid the
|
||||
// word `unauthorized` for the same reason; "rejected"/"forbidden"
|
||||
// are the user-visible equivalents.
|
||||
let code_word = if status == reqwest::StatusCode::UNAUTHORIZED {
|
||||
"four-oh-one"
|
||||
} else {
|
||||
"four-oh-three"
|
||||
};
|
||||
// Deliberately do NOT splice the upstream body into this
|
||||
// user-facing message — Discord's auth-error bodies often
|
||||
// include the literal words "401" and "Unauthorized", which
|
||||
// would smuggle the cascade trigger back in. The body is
|
||||
// still in `tracing::debug!` logs above the call site for
|
||||
// triage; the user-facing message only needs the remediation.
|
||||
let _ = body;
|
||||
format!(
|
||||
"Discord {endpoint}: {kind} (upstream HTTP {code_word}). \
|
||||
Open Settings → Channels → Discord and rotate / reconnect the bot \
|
||||
token."
|
||||
)
|
||||
} else {
|
||||
format!("Discord {endpoint} failed ({status}): {body}")
|
||||
}
|
||||
}
|
||||
|
||||
/// List all guilds (servers) the bot is a member of.
|
||||
pub async fn list_bot_guilds(token: &str) -> anyhow::Result<Vec<DiscordGuild>> {
|
||||
list_bot_guilds_at_base(DISCORD_API_BASE, token).await
|
||||
}
|
||||
|
||||
/// Test seam: list guilds against an arbitrary API base. Used by
|
||||
/// `list_bot_guilds` in production and by unit tests that drive a
|
||||
/// local mock Discord API.
|
||||
async fn list_bot_guilds_at_base(base: &str, token: &str) -> anyhow::Result<Vec<DiscordGuild>> {
|
||||
let url = format!("{base}/users/@me/guilds");
|
||||
tracing::debug!("[discord-api] listing guilds for bot");
|
||||
|
||||
let resp = build_client()
|
||||
.get(&url)
|
||||
.header("Authorization", auth_header(token))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
tracing::debug!(
|
||||
target: "discord-api",
|
||||
endpoint = "list_guilds",
|
||||
%url,
|
||||
%status,
|
||||
body = %body,
|
||||
"[discord-api] non-success response"
|
||||
);
|
||||
anyhow::bail!(
|
||||
"{}",
|
||||
format_discord_http_error("list_guilds", status, &body)
|
||||
);
|
||||
}
|
||||
|
||||
let guilds: Vec<DiscordGuild> = resp.json().await?;
|
||||
tracing::debug!("[discord-api] found {} guilds", guilds.len());
|
||||
Ok(guilds)
|
||||
}
|
||||
|
||||
/// List text channels in a guild. Filters to type=0 (text channels) only.
|
||||
pub async fn list_guild_channels(
|
||||
token: &str,
|
||||
guild_id: &str,
|
||||
) -> anyhow::Result<Vec<DiscordTextChannel>> {
|
||||
list_guild_channels_at_base(DISCORD_API_BASE, token, guild_id).await
|
||||
}
|
||||
|
||||
/// Test seam: list guild channels against an arbitrary API base.
|
||||
async fn list_guild_channels_at_base(
|
||||
base: &str,
|
||||
token: &str,
|
||||
guild_id: &str,
|
||||
) -> anyhow::Result<Vec<DiscordTextChannel>> {
|
||||
let url = format!("{base}/guilds/{guild_id}/channels");
|
||||
tracing::debug!("[discord-api] listing channels for guild {guild_id}");
|
||||
|
||||
let resp = build_client()
|
||||
.get(&url)
|
||||
.header("Authorization", auth_header(token))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
tracing::debug!(
|
||||
target: "discord-api",
|
||||
endpoint = "list_guild_channels",
|
||||
%guild_id,
|
||||
%url,
|
||||
%status,
|
||||
body = %body,
|
||||
"[discord-api] non-success response"
|
||||
);
|
||||
anyhow::bail!(
|
||||
"{}",
|
||||
format_discord_http_error("list_channels", status, &body)
|
||||
);
|
||||
}
|
||||
|
||||
let all_channels: Vec<DiscordTextChannel> = resp.json().await?;
|
||||
|
||||
// Filter to text channels (type 0) and sort by position
|
||||
let mut text_channels: Vec<DiscordTextChannel> = all_channels
|
||||
.into_iter()
|
||||
.filter(|c| c.channel_type == 0)
|
||||
.collect();
|
||||
text_channels.sort_by_key(|c| c.position);
|
||||
|
||||
tracing::debug!(
|
||||
"[discord-api] found {} text channels in guild {guild_id}",
|
||||
text_channels.len()
|
||||
);
|
||||
Ok(text_channels)
|
||||
}
|
||||
|
||||
/// Check bot permissions in a specific channel.
|
||||
///
|
||||
/// Uses `GET /channels/{channel_id}` combined with the bot's guild member
|
||||
/// permissions to determine if the bot can view, send, and read history.
|
||||
pub async fn check_channel_permissions(
|
||||
token: &str,
|
||||
guild_id: &str,
|
||||
channel_id: &str,
|
||||
) -> anyhow::Result<BotPermissionCheck> {
|
||||
check_channel_permissions_at_base(DISCORD_API_BASE, token, guild_id, channel_id).await
|
||||
}
|
||||
|
||||
/// Test seam: see [`check_channel_permissions`].
|
||||
async fn check_channel_permissions_at_base(
|
||||
base: &str,
|
||||
token: &str,
|
||||
guild_id: &str,
|
||||
channel_id: &str,
|
||||
) -> anyhow::Result<BotPermissionCheck> {
|
||||
tracing::debug!(
|
||||
"[discord-api] checking permissions in channel {channel_id} (guild {guild_id})"
|
||||
);
|
||||
|
||||
// Resolve bot user id first (`members/@me` is not a valid Discord route).
|
||||
let me_url = format!("{base}/users/@me");
|
||||
let me_resp = build_client()
|
||||
.get(&me_url)
|
||||
.header("Authorization", auth_header(token))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !me_resp.status().is_success() {
|
||||
let status = me_resp.status();
|
||||
let body = me_resp.text().await.unwrap_or_default();
|
||||
tracing::debug!(
|
||||
target: "discord-api",
|
||||
endpoint = "check_bot_permissions.me",
|
||||
%guild_id,
|
||||
%channel_id,
|
||||
url = %me_url,
|
||||
%status,
|
||||
body = %body,
|
||||
"[discord-api] non-success response"
|
||||
);
|
||||
anyhow::bail!(
|
||||
"{}",
|
||||
format_discord_http_error("get_bot_user", status, &body)
|
||||
);
|
||||
}
|
||||
let me: serde_json::Value = me_resp.json().await?;
|
||||
let bot_user_id = me.get("id").and_then(|i| i.as_str()).unwrap_or("").trim();
|
||||
if bot_user_id.is_empty() {
|
||||
anyhow::bail!("Discord get bot user returned empty id");
|
||||
}
|
||||
|
||||
// Fetch the bot's guild member info which includes role ids.
|
||||
let member_url = format!("{base}/guilds/{guild_id}/members/{bot_user_id}");
|
||||
let member_resp = build_client()
|
||||
.get(&member_url)
|
||||
.header("Authorization", auth_header(token))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !member_resp.status().is_success() {
|
||||
let status = member_resp.status();
|
||||
let body = member_resp.text().await.unwrap_or_default();
|
||||
tracing::debug!(
|
||||
target: "discord-api",
|
||||
endpoint = "check_bot_permissions.member",
|
||||
%guild_id,
|
||||
%channel_id,
|
||||
url = %member_url,
|
||||
%status,
|
||||
body = %body,
|
||||
"[discord-api] non-success response"
|
||||
);
|
||||
anyhow::bail!(
|
||||
"{}",
|
||||
format_discord_http_error("get_member_info", status, &body)
|
||||
);
|
||||
}
|
||||
|
||||
let member: serde_json::Value = member_resp.json().await?;
|
||||
|
||||
// Fetch guild roles to compute permissions
|
||||
let roles_url = format!("{base}/guilds/{guild_id}/roles");
|
||||
let roles_resp = build_client()
|
||||
.get(&roles_url)
|
||||
.header("Authorization", auth_header(token))
|
||||
.send()
|
||||
.await?;
|
||||
if !roles_resp.status().is_success() {
|
||||
let status = roles_resp.status();
|
||||
let body = roles_resp.text().await.unwrap_or_default();
|
||||
tracing::debug!(
|
||||
target: "discord-api",
|
||||
endpoint = "check_bot_permissions.roles",
|
||||
%guild_id,
|
||||
%channel_id,
|
||||
url = %roles_url,
|
||||
%status,
|
||||
body = %body,
|
||||
"[discord-api] non-success response"
|
||||
);
|
||||
anyhow::bail!(
|
||||
"{}",
|
||||
format_discord_http_error("get_guild_roles", status, &body)
|
||||
);
|
||||
}
|
||||
let guild_roles: Vec<serde_json::Value> = roles_resp.json().await?;
|
||||
|
||||
// Get the member's role IDs
|
||||
let member_role_ids: Vec<&str> = member
|
||||
.get("roles")
|
||||
.and_then(|r| r.as_array())
|
||||
.map(|arr| arr.iter().filter_map(|v| v.as_str()).collect::<Vec<&str>>())
|
||||
.unwrap_or_default();
|
||||
|
||||
// Compute base permissions from @everyone role + member roles
|
||||
let mut permissions: u64 = 0;
|
||||
for role in &guild_roles {
|
||||
let role_id = role.get("id").and_then(|i| i.as_str()).unwrap_or("");
|
||||
let is_everyone = role_id == guild_id; // @everyone role ID == guild ID
|
||||
let is_member_role = member_role_ids.contains(&role_id);
|
||||
|
||||
if is_everyone || is_member_role {
|
||||
if let Some(perms_str) = role.get("permissions").and_then(|p| p.as_str()) {
|
||||
if let Ok(perms) = perms_str.parse::<u64>() {
|
||||
permissions |= perms;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Administrator bypasses all permission checks
|
||||
const ADMINISTRATOR: u64 = 1 << 3;
|
||||
if permissions & ADMINISTRATOR != 0 {
|
||||
return Ok(BotPermissionCheck {
|
||||
can_view_channel: true,
|
||||
can_send_messages: true,
|
||||
can_read_message_history: true,
|
||||
missing_permissions: vec![],
|
||||
});
|
||||
}
|
||||
|
||||
// Now check channel-level permission overwrites
|
||||
let channel_url = format!("{base}/channels/{channel_id}");
|
||||
let ch_resp = build_client()
|
||||
.get(&channel_url)
|
||||
.header("Authorization", auth_header(token))
|
||||
.send()
|
||||
.await?;
|
||||
if !ch_resp.status().is_success() {
|
||||
let status = ch_resp.status();
|
||||
let body = ch_resp.text().await.unwrap_or_default();
|
||||
tracing::debug!(
|
||||
target: "discord-api",
|
||||
endpoint = "check_bot_permissions.channel",
|
||||
%guild_id,
|
||||
%channel_id,
|
||||
url = %channel_url,
|
||||
%status,
|
||||
body = %body,
|
||||
"[discord-api] non-success response"
|
||||
);
|
||||
anyhow::bail!(
|
||||
"{}",
|
||||
format_discord_http_error("get_channel", status, &body)
|
||||
);
|
||||
}
|
||||
let channel_data: serde_json::Value = ch_resp.json().await?;
|
||||
if let Some(overwrites) = channel_data
|
||||
.get("permission_overwrites")
|
||||
.and_then(|o| o.as_array())
|
||||
{
|
||||
// Intentional shadowing: prefer the ID returned inside the member
|
||||
// object over the one fetched from /users/@me, because the guild
|
||||
// member record is more authoritative for permission overwrite lookups.
|
||||
let bot_user_id = member
|
||||
.get("user")
|
||||
.and_then(|u| u.get("id"))
|
||||
.and_then(|i| i.as_str())
|
||||
.unwrap_or(bot_user_id);
|
||||
|
||||
let mut everyone_allow = 0_u64;
|
||||
let mut everyone_deny = 0_u64;
|
||||
let mut role_allow = 0_u64;
|
||||
let mut role_deny = 0_u64;
|
||||
let mut member_allow = 0_u64;
|
||||
let mut member_deny = 0_u64;
|
||||
|
||||
for overwrite in overwrites {
|
||||
let ow_id = overwrite.get("id").and_then(|i| i.as_str()).unwrap_or("");
|
||||
let ow_type = overwrite.get("type").and_then(|t| t.as_u64()).unwrap_or(0);
|
||||
let allow = overwrite
|
||||
.get("allow")
|
||||
.and_then(|a| a.as_str())
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(0);
|
||||
let deny = overwrite
|
||||
.get("deny")
|
||||
.and_then(|d| d.as_str())
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(0);
|
||||
|
||||
match ow_type {
|
||||
// @everyone overwrite (role id == guild id)
|
||||
0 if ow_id == guild_id => {
|
||||
everyone_allow = allow;
|
||||
everyone_deny = deny;
|
||||
}
|
||||
// Aggregate all role overwrites
|
||||
0 if member_role_ids.contains(&ow_id) => {
|
||||
role_allow |= allow;
|
||||
role_deny |= deny;
|
||||
}
|
||||
// Member-specific overwrite
|
||||
1 if ow_id == bot_user_id => {
|
||||
member_allow = allow;
|
||||
member_deny = deny;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// Apply Discord overwrite precedence: everyone -> roles -> member.
|
||||
permissions &= !everyone_deny;
|
||||
permissions |= everyone_allow;
|
||||
permissions &= !role_deny;
|
||||
permissions |= role_allow;
|
||||
permissions &= !member_deny;
|
||||
permissions |= member_allow;
|
||||
}
|
||||
|
||||
let can_view = permissions & VIEW_CHANNEL != 0;
|
||||
let can_send = permissions & SEND_MESSAGES != 0;
|
||||
let can_read_history = permissions & READ_MESSAGE_HISTORY != 0;
|
||||
|
||||
let mut missing = Vec::new();
|
||||
if !can_view {
|
||||
missing.push("VIEW_CHANNEL".to_string());
|
||||
}
|
||||
if !can_send {
|
||||
missing.push("SEND_MESSAGES".to_string());
|
||||
}
|
||||
if !can_read_history {
|
||||
missing.push("READ_MESSAGE_HISTORY".to_string());
|
||||
}
|
||||
|
||||
tracing::debug!(
|
||||
"[discord-api] permissions for channel {channel_id}: view={can_view}, send={can_send}, history={can_read_history}"
|
||||
);
|
||||
|
||||
Ok(BotPermissionCheck {
|
||||
can_view_channel: can_view,
|
||||
can_send_messages: can_send,
|
||||
can_read_message_history: can_read_history,
|
||||
missing_permissions: missing,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "api_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
#[cfg(any(test, debug_assertions))]
|
||||
pub mod test_support {
|
||||
//! Debug-build wrappers for raw integration tests that drive the Discord
|
||||
//! REST helpers against loopback servers.
|
||||
|
||||
use super::*;
|
||||
|
||||
pub fn format_discord_http_error_for_test(
|
||||
endpoint: &str,
|
||||
status: reqwest::StatusCode,
|
||||
body: &str,
|
||||
) -> String {
|
||||
format_discord_http_error(endpoint, status, body)
|
||||
}
|
||||
|
||||
pub async fn list_bot_guilds_at_base_for_test(
|
||||
base: &str,
|
||||
token: &str,
|
||||
) -> anyhow::Result<Vec<DiscordGuild>> {
|
||||
list_bot_guilds_at_base(base, token).await
|
||||
}
|
||||
|
||||
pub async fn list_guild_channels_at_base_for_test(
|
||||
base: &str,
|
||||
token: &str,
|
||||
guild_id: &str,
|
||||
) -> anyhow::Result<Vec<DiscordTextChannel>> {
|
||||
list_guild_channels_at_base(base, token, guild_id).await
|
||||
}
|
||||
|
||||
pub async fn check_channel_permissions_at_base_for_test(
|
||||
base: &str,
|
||||
token: &str,
|
||||
guild_id: &str,
|
||||
channel_id: &str,
|
||||
) -> anyhow::Result<BotPermissionCheck> {
|
||||
check_channel_permissions_at_base(base, token, guild_id, channel_id).await
|
||||
}
|
||||
}
|
||||
@@ -1,451 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn guild_deserializes() {
|
||||
let json = r#"{"id":"123","name":"Test Server","icon":"abc123"}"#;
|
||||
let guild: DiscordGuild = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(guild.id, "123");
|
||||
assert_eq!(guild.name, "Test Server");
|
||||
assert_eq!(guild.icon, Some("abc123".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn guild_deserializes_without_icon() {
|
||||
let json = r#"{"id":"456","name":"No Icon","icon":null}"#;
|
||||
let guild: DiscordGuild = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(guild.id, "456");
|
||||
assert!(guild.icon.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_channel_deserializes() {
|
||||
let json = r#"{"id":"789","name":"general","type":0,"position":1,"parent_id":"100"}"#;
|
||||
let ch: DiscordTextChannel = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(ch.id, "789");
|
||||
assert_eq!(ch.name, "general");
|
||||
assert_eq!(ch.channel_type, 0);
|
||||
assert_eq!(ch.position, 1);
|
||||
assert_eq!(ch.parent_id, Some("100".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_channel_without_parent() {
|
||||
let json = r#"{"id":"789","name":"general","type":0,"position":0,"parent_id":null}"#;
|
||||
let ch: DiscordTextChannel = serde_json::from_str(json).unwrap();
|
||||
assert!(ch.parent_id.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn permission_check_serializes() {
|
||||
let check = BotPermissionCheck {
|
||||
can_view_channel: true,
|
||||
can_send_messages: true,
|
||||
can_read_message_history: false,
|
||||
missing_permissions: vec!["READ_MESSAGE_HISTORY".to_string()],
|
||||
};
|
||||
let json = serde_json::to_string(&check).unwrap();
|
||||
assert!(json.contains("READ_MESSAGE_HISTORY"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn permission_bits_are_correct() {
|
||||
assert_eq!(VIEW_CHANNEL, 1024);
|
||||
assert_eq!(SEND_MESSAGES, 2048);
|
||||
assert_eq!(READ_MESSAGE_HISTORY, 65536);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_header_has_bot_prefix() {
|
||||
assert_eq!(auth_header("abc"), "Bot abc");
|
||||
assert_eq!(auth_header(""), "Bot ");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn permission_check_lists_all_missing_permissions_when_bot_lacks_any() {
|
||||
let check = BotPermissionCheck {
|
||||
can_view_channel: false,
|
||||
can_send_messages: false,
|
||||
can_read_message_history: false,
|
||||
missing_permissions: vec![
|
||||
"VIEW_CHANNEL".into(),
|
||||
"SEND_MESSAGES".into(),
|
||||
"READ_MESSAGE_HISTORY".into(),
|
||||
],
|
||||
};
|
||||
let json = serde_json::to_string(&check).unwrap();
|
||||
assert!(json.contains("VIEW_CHANNEL"));
|
||||
assert!(json.contains("SEND_MESSAGES"));
|
||||
assert!(json.contains("READ_MESSAGE_HISTORY"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn permission_check_with_all_granted_has_empty_missing_list() {
|
||||
let check = BotPermissionCheck {
|
||||
can_view_channel: true,
|
||||
can_send_messages: true,
|
||||
can_read_message_history: true,
|
||||
missing_permissions: vec![],
|
||||
};
|
||||
let json = serde_json::to_string(&check).unwrap();
|
||||
assert!(json.contains("\"missing_permissions\":[]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_channel_type_zero_is_standard_text() {
|
||||
let json = r#"{"id":"1","name":"general","type":0,"position":0,"parent_id":null}"#;
|
||||
let ch: DiscordTextChannel = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(ch.channel_type, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn guild_deserializes_with_full_payload() {
|
||||
let json = r#"{
|
||||
"id": "999",
|
||||
"name": "Full Guild",
|
||||
"icon": "hash"
|
||||
}"#;
|
||||
let g: DiscordGuild = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(g.id, "999");
|
||||
assert_eq!(g.name, "Full Guild");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn permission_bit_flags_are_disjoint() {
|
||||
// Sanity: each permission is a single bit and distinct.
|
||||
assert_eq!(VIEW_CHANNEL.count_ones(), 1);
|
||||
assert_eq!(SEND_MESSAGES.count_ones(), 1);
|
||||
assert_eq!(READ_MESSAGE_HISTORY.count_ones(), 1);
|
||||
assert_ne!(VIEW_CHANNEL, SEND_MESSAGES);
|
||||
assert_ne!(SEND_MESSAGES, READ_MESSAGE_HISTORY);
|
||||
}
|
||||
|
||||
// ── Mock Discord server integration tests ──────────────────────
|
||||
|
||||
use axum::{extract::Path, http::StatusCode, routing::get, Json, Router};
|
||||
use serde_json::json;
|
||||
|
||||
async fn spawn_mock(app: Router) -> String {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
format!("http://127.0.0.1:{}", addr.port())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_bot_guilds_parses_discord_response() {
|
||||
let app = Router::new().route(
|
||||
"/users/@me/guilds",
|
||||
get(|| async {
|
||||
Json(json!([
|
||||
{"id": "g1", "name": "Guild One", "icon": "hash1"},
|
||||
{"id": "g2", "name": "Guild Two", "icon": null}
|
||||
]))
|
||||
}),
|
||||
);
|
||||
let base = spawn_mock(app).await;
|
||||
let guilds = list_bot_guilds_at_base(&base, "test-token").await.unwrap();
|
||||
assert_eq!(guilds.len(), 2);
|
||||
assert_eq!(guilds[0].id, "g1");
|
||||
assert_eq!(guilds[0].name, "Guild One");
|
||||
assert_eq!(guilds[1].icon, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_bot_guilds_rewraps_401_so_global_session_cascade_does_not_fire() {
|
||||
// Upstream returns 401 with the canonical Discord auth-error body.
|
||||
// BEFORE #2285 the error string flowed up to JSON-RPC as
|
||||
// "Discord list guilds failed (401 Unauthorized): {\"message\":
|
||||
// \"401: Unauthorized\",\"code\":0}" — that pair tripped
|
||||
// `jsonrpc::is_session_expired_error` ("401" + "unauthorized")
|
||||
// and logged the user out of OpenHuman over a *Discord*
|
||||
// credentials problem.
|
||||
//
|
||||
// After the fix the user-facing message:
|
||||
// - does NOT contain "401" or "unauthorized" as substrings
|
||||
// (so `is_session_expired_error` returns false), AND
|
||||
// - names the endpoint + the actionable Settings → Channels →
|
||||
// Discord remediation path.
|
||||
let app = Router::new().route(
|
||||
"/users/@me/guilds",
|
||||
get(|| async {
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
r#"{"message":"401: Unauthorized","code":0}"#,
|
||||
)
|
||||
}),
|
||||
);
|
||||
let base = spawn_mock(app).await;
|
||||
let err = list_bot_guilds_at_base(&base, "t")
|
||||
.await
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
let lower = err.to_ascii_lowercase();
|
||||
assert!(
|
||||
!lower.contains("401"),
|
||||
"must NOT contain '401' substring: {err}"
|
||||
);
|
||||
assert!(
|
||||
!lower.contains("unauthorized"),
|
||||
"must NOT contain 'unauthorized' substring: {err}"
|
||||
);
|
||||
assert!(
|
||||
err.contains("list_guilds"),
|
||||
"endpoint identifier preserved for triage: {err}"
|
||||
);
|
||||
assert!(
|
||||
err.contains("Settings → Channels → Discord"),
|
||||
"remediation path present: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_bot_guilds_5xx_still_carries_raw_status() {
|
||||
// Non-auth errors fall through to the legacy verbose format —
|
||||
// those don't match `is_session_expired_error` even verbatim, so
|
||||
// surfacing the raw status code helps the user / triage.
|
||||
let app = Router::new().route(
|
||||
"/users/@me/guilds",
|
||||
get(|| async { (StatusCode::INTERNAL_SERVER_ERROR, "discord melting") }),
|
||||
);
|
||||
let base = spawn_mock(app).await;
|
||||
let err = list_bot_guilds_at_base(&base, "t")
|
||||
.await
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
assert!(
|
||||
err.contains("500"),
|
||||
"5xx must surface verbatim status: {err}"
|
||||
);
|
||||
assert!(err.contains("list_guilds"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_guild_channels_filters_text_channels_and_sorts_by_position() {
|
||||
let app = Router::new().route(
|
||||
"/guilds/{guild_id}/channels",
|
||||
get(|Path(guild_id): Path<String>| async move {
|
||||
assert_eq!(guild_id, "g1");
|
||||
Json(json!([
|
||||
{"id": "c3", "name": "category", "type": 4, "position": 0, "parent_id": null},
|
||||
{"id": "c1", "name": "general", "type": 0, "position": 2, "parent_id": null},
|
||||
{"id": "c2", "name": "random", "type": 0, "position": 1, "parent_id": null},
|
||||
{"id": "c4", "name": "voice", "type": 2, "position": 3, "parent_id": null}
|
||||
]))
|
||||
}),
|
||||
);
|
||||
let base = spawn_mock(app).await;
|
||||
let channels = list_guild_channels_at_base(&base, "t", "g1").await.unwrap();
|
||||
// Only text channels (type=0) remain, sorted by position ascending.
|
||||
assert_eq!(channels.len(), 2);
|
||||
assert_eq!(channels[0].id, "c2");
|
||||
assert_eq!(channels[1].id, "c1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_guild_channels_rewraps_403_with_remediation_and_no_session_keywords() {
|
||||
// 403 follows the same rewrap path as 401 (#2285) — both can
|
||||
// happen on a stale/disabled bot token AND both share enough
|
||||
// substrings with `is_session_expired_error` to be a problem if
|
||||
// raw upstream text reaches the JSON-RPC layer. The user-facing
|
||||
// message must use the safer wording.
|
||||
let app = Router::new().route(
|
||||
"/guilds/{guild_id}/channels",
|
||||
get(|| async {
|
||||
(
|
||||
StatusCode::FORBIDDEN,
|
||||
r#"{"message":"Missing Access","code":50001}"#,
|
||||
)
|
||||
}),
|
||||
);
|
||||
let base = spawn_mock(app).await;
|
||||
let err = list_guild_channels_at_base(&base, "t", "g1")
|
||||
.await
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
let lower = err.to_ascii_lowercase();
|
||||
assert!(!lower.contains("403"), "raw 403 must not leak: {err}");
|
||||
assert!(
|
||||
err.contains("list_channels"),
|
||||
"endpoint identifier preserved: {err}"
|
||||
);
|
||||
assert!(
|
||||
err.contains("Settings → Channels → Discord"),
|
||||
"remediation path present: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_guild_channels_empty_returns_empty_vec() {
|
||||
let app = Router::new().route(
|
||||
"/guilds/{guild_id}/channels",
|
||||
get(|| async { Json(json!([])) }),
|
||||
);
|
||||
let base = spawn_mock(app).await;
|
||||
let channels = list_guild_channels_at_base(&base, "t", "g").await.unwrap();
|
||||
assert!(channels.is_empty());
|
||||
}
|
||||
|
||||
// ── check_channel_permissions ─────────────────────────────────
|
||||
|
||||
/// Build a mock Discord that answers all endpoints the permissions check
|
||||
/// touches: `/users/@me`, `/guilds/<id>/members/<bot_id>`,
|
||||
/// `/guilds/<id>/roles`, and `/channels/<id>`.
|
||||
fn permissions_mock(
|
||||
member: serde_json::Value,
|
||||
roles: serde_json::Value,
|
||||
channel: serde_json::Value,
|
||||
) -> Router {
|
||||
use axum::extract::Path;
|
||||
Router::new()
|
||||
.route(
|
||||
"/users/@me",
|
||||
get(|| async { Json(json!({ "id": "bot-1" })) }),
|
||||
)
|
||||
.route(
|
||||
"/guilds/{guild_id}/members/{member_id}",
|
||||
get(move |Path((_g, member_id)): Path<(String, String)>| {
|
||||
assert_eq!(member_id, "bot-1");
|
||||
let m = member.clone();
|
||||
async move { Json(m) }
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/guilds/{guild_id}/roles",
|
||||
get(move |Path(_g): Path<String>| {
|
||||
let r = roles.clone();
|
||||
async move { Json(r) }
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/channels/{channel_id}",
|
||||
get(move |Path(_c): Path<String>| {
|
||||
let c = channel.clone();
|
||||
async move { Json(c) }
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn check_channel_permissions_administrator_bypasses_everything() {
|
||||
let member = json!({ "roles": ["role-admin"], "user": { "id": "bot-1" } });
|
||||
// Role with Administrator bit (1<<3 = 8) — overrides all other checks.
|
||||
let roles = json!([
|
||||
{ "id": "role-admin", "permissions": "8" }
|
||||
]);
|
||||
let channel = json!({ "permission_overwrites": [] });
|
||||
let base = spawn_mock(permissions_mock(member, roles, channel)).await;
|
||||
let out = check_channel_permissions_at_base(&base, "token", "guild-1", "channel-1")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(out.can_view_channel);
|
||||
assert!(out.can_send_messages);
|
||||
assert!(out.can_read_message_history);
|
||||
assert!(out.missing_permissions.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn check_channel_permissions_flags_missing_bits_when_role_lacks_them() {
|
||||
// No roles grant any of the 3 permissions → all missing.
|
||||
let member = json!({ "roles": ["role-nobody"], "user": { "id": "bot-1" } });
|
||||
let roles = json!([
|
||||
{ "id": "role-nobody", "permissions": "0" }
|
||||
]);
|
||||
let channel = json!({ "permission_overwrites": [] });
|
||||
let base = spawn_mock(permissions_mock(member, roles, channel)).await;
|
||||
let out = check_channel_permissions_at_base(&base, "t", "guild-1", "channel-1")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!out.can_view_channel);
|
||||
assert!(!out.can_send_messages);
|
||||
assert!(!out.can_read_message_history);
|
||||
assert!(out
|
||||
.missing_permissions
|
||||
.contains(&"VIEW_CHANNEL".to_string()));
|
||||
assert!(out
|
||||
.missing_permissions
|
||||
.contains(&"SEND_MESSAGES".to_string()));
|
||||
assert!(out
|
||||
.missing_permissions
|
||||
.contains(&"READ_MESSAGE_HISTORY".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn check_channel_permissions_grants_everything_when_everyone_role_allows() {
|
||||
// @everyone role (id == guild_id) grants VIEW|SEND|HISTORY
|
||||
// = 1024 | 2048 | 65536 = 68608
|
||||
let member = json!({ "roles": [], "user": { "id": "bot-1" } });
|
||||
let roles = json!([
|
||||
{ "id": "guild-1", "permissions": "68608" }
|
||||
]);
|
||||
let channel = json!({ "permission_overwrites": [] });
|
||||
let base = spawn_mock(permissions_mock(member, roles, channel)).await;
|
||||
let out = check_channel_permissions_at_base(&base, "t", "guild-1", "channel-1")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(out.can_view_channel);
|
||||
assert!(out.can_send_messages);
|
||||
assert!(out.can_read_message_history);
|
||||
assert!(out.missing_permissions.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn check_channel_permissions_channel_overwrite_can_deny_permission() {
|
||||
// @everyone role grants everything, but the channel's @everyone
|
||||
// overwrite denies VIEW_CHANNEL — expect VIEW missing.
|
||||
let member = json!({ "roles": [], "user": { "id": "bot-1" } });
|
||||
let roles = json!([
|
||||
{ "id": "guild-1", "permissions": "68608" }
|
||||
]);
|
||||
let channel = json!({
|
||||
"permission_overwrites": [
|
||||
{
|
||||
"id": "guild-1",
|
||||
"type": 0,
|
||||
"allow": "0",
|
||||
"deny": "1024" // VIEW_CHANNEL
|
||||
}
|
||||
]
|
||||
});
|
||||
let base = spawn_mock(permissions_mock(member, roles, channel)).await;
|
||||
let out = check_channel_permissions_at_base(&base, "t", "guild-1", "channel-1")
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!out.can_view_channel);
|
||||
assert!(out
|
||||
.missing_permissions
|
||||
.contains(&"VIEW_CHANNEL".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn check_channel_permissions_errors_on_member_lookup_failure() {
|
||||
use axum::http::StatusCode;
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/users/@me",
|
||||
get(|| async { Json(json!({ "id": "bot-1" })) }),
|
||||
)
|
||||
.route(
|
||||
"/guilds/{guild_id}/members/{member_id}",
|
||||
get(|Path((_g, _member_id)): Path<(String, String)>| async {
|
||||
(StatusCode::UNAUTHORIZED, "bad token")
|
||||
}),
|
||||
);
|
||||
let base = spawn_mock(app).await;
|
||||
let err = check_channel_permissions_at_base(&base, "t", "g", "c")
|
||||
.await
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
// Endpoint identifier preserved in the rewrap (#2285), and the
|
||||
// 401 path keeps the substrings "401"/"unauthorized" out of the
|
||||
// user-facing message so the JSON-RPC session-expired classifier
|
||||
// ignores it.
|
||||
assert!(err.contains("get_member_info"));
|
||||
assert!(
|
||||
!err.to_ascii_lowercase().contains("401"),
|
||||
"rewrapped message must not contain '401': {err}"
|
||||
);
|
||||
assert!(
|
||||
!err.to_ascii_lowercase().contains("unauthorized"),
|
||||
"rewrapped message must not contain 'unauthorized': {err}"
|
||||
);
|
||||
}
|
||||
@@ -1,525 +0,0 @@
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
use async_trait::async_trait;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use parking_lot::Mutex;
|
||||
use serde_json::json;
|
||||
use tinychannels::channel::LengthUnit;
|
||||
use tinychannels::text::{chunk_text_with_options, ChunkMode, TextChunkOptions};
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Discord channel — connects via Gateway WebSocket for real-time messages
|
||||
pub struct DiscordChannel {
|
||||
bot_token: String,
|
||||
guild_id: Option<String>,
|
||||
channel_id: Option<String>,
|
||||
allowed_users: Vec<String>,
|
||||
listen_to_bots: bool,
|
||||
mention_only: bool,
|
||||
typing_handle: Mutex<Option<tokio::task::JoinHandle<()>>>,
|
||||
}
|
||||
|
||||
impl DiscordChannel {
|
||||
pub fn new(
|
||||
bot_token: String,
|
||||
guild_id: Option<String>,
|
||||
channel_id: Option<String>,
|
||||
allowed_users: Vec<String>,
|
||||
listen_to_bots: bool,
|
||||
mention_only: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
bot_token,
|
||||
guild_id,
|
||||
channel_id,
|
||||
allowed_users,
|
||||
listen_to_bots,
|
||||
mention_only,
|
||||
typing_handle: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn http_client(&self) -> reqwest::Client {
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.discord")
|
||||
}
|
||||
|
||||
/// Check if a Discord user ID is in the allowlist.
|
||||
///
|
||||
/// Empty list ⇒ allow-all: an unconfigured allowlist applies no per-user
|
||||
/// restriction (the bot is still scoped to its configured guild/channel).
|
||||
/// Previously an empty list denied *everyone*, so a bot connected via the UI
|
||||
/// with the default-empty allowlist silently ignored every message and never
|
||||
/// replied (issue #3712). This now matches the WhatsApp provider's
|
||||
/// empty-⇒-allow-all convention. `"*"` also allows everyone; populate the
|
||||
/// list with specific user IDs to restrict.
|
||||
fn is_user_allowed(&self, user_id: &str) -> bool {
|
||||
self.allowed_users.is_empty() || self.allowed_users.iter().any(|u| u == "*" || u == user_id)
|
||||
}
|
||||
|
||||
/// Decide whether a message passes the bot's guild scoping.
|
||||
///
|
||||
/// - No configured guild → all messages pass (guild filter inactive).
|
||||
/// - Configured guild, message in a *different* guild → blocked.
|
||||
/// - Configured guild, message in that guild → allowed.
|
||||
/// - Configured guild, DM (no `guild_id`) → only when the allowlist is
|
||||
/// non-empty (explicit). `is_user_allowed` treats an empty list as
|
||||
/// allow-all (the intended *within-guild* default), but DMs bypass the
|
||||
/// guild filter — so a blank allowlist must not open a guild-scoped bot to
|
||||
/// arbitrary DMs (#3794 review — Codex P1). `"*"` (non-empty) still allows.
|
||||
fn passes_guild_scope(
|
||||
configured_guild: Option<&str>,
|
||||
msg_guild: Option<&str>,
|
||||
allowlist_empty: bool,
|
||||
) -> bool {
|
||||
let Some(gid) = configured_guild else {
|
||||
return true;
|
||||
};
|
||||
match msg_guild {
|
||||
Some(g) => g == gid,
|
||||
None => !allowlist_empty,
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the outbound recipient channel id. Prefer the message's explicit
|
||||
/// recipient (e.g. the channel a reply targets); fall back to the bot's
|
||||
/// configured `channel_id` for recipient-less sends such as proactive
|
||||
/// cron/heartbeat delivery (#3794 review — Codex P2). `None` when neither is
|
||||
/// available, so the caller surfaces an error instead of POSTing to an empty
|
||||
/// channel id.
|
||||
fn resolve_recipient<'a>(
|
||||
msg_recipient: &'a str,
|
||||
configured: Option<&'a str>,
|
||||
) -> Option<&'a str> {
|
||||
let recipient = if msg_recipient.is_empty() {
|
||||
configured.unwrap_or("")
|
||||
} else {
|
||||
msg_recipient
|
||||
};
|
||||
(!recipient.is_empty()).then_some(recipient)
|
||||
}
|
||||
|
||||
fn bot_user_id_from_token(token: &str) -> Option<String> {
|
||||
// Discord bot tokens are base64(bot_user_id).timestamp.hmac
|
||||
let part = token.split('.').next()?;
|
||||
base64_decode(part)
|
||||
}
|
||||
}
|
||||
|
||||
const BASE64_ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
|
||||
|
||||
/// Discord's maximum message length for regular messages.
|
||||
///
|
||||
/// Discord rejects longer payloads with `50035 Invalid Form Body`.
|
||||
const DISCORD_MAX_MESSAGE_LENGTH: usize = 2000;
|
||||
|
||||
/// Split a message into chunks that respect Discord's 2000-character limit.
|
||||
/// Tries to split at word boundaries when possible.
|
||||
fn split_message_for_discord(message: &str) -> Vec<String> {
|
||||
if message.is_empty() {
|
||||
return vec![String::new()];
|
||||
}
|
||||
chunk_text_with_options(
|
||||
message,
|
||||
TextChunkOptions {
|
||||
limit: DISCORD_MAX_MESSAGE_LENGTH,
|
||||
length_unit: LengthUnit::Utf16Units,
|
||||
mode: ChunkMode::Length,
|
||||
markdown: true,
|
||||
indicators: false,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
fn mention_tags(bot_user_id: &str) -> [String; 2] {
|
||||
[format!("<@{bot_user_id}>"), format!("<@!{bot_user_id}>")]
|
||||
}
|
||||
|
||||
fn contains_bot_mention(content: &str, bot_user_id: &str) -> bool {
|
||||
let tags = mention_tags(bot_user_id);
|
||||
content.contains(&tags[0]) || content.contains(&tags[1])
|
||||
}
|
||||
|
||||
fn normalize_incoming_content(
|
||||
content: &str,
|
||||
mention_only: bool,
|
||||
bot_user_id: &str,
|
||||
) -> Option<String> {
|
||||
if content.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
if mention_only && !contains_bot_mention(content, bot_user_id) {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut normalized = content.to_string();
|
||||
if mention_only {
|
||||
for tag in mention_tags(bot_user_id) {
|
||||
normalized = normalized.replace(&tag, " ");
|
||||
}
|
||||
}
|
||||
|
||||
let normalized = normalized.trim().to_string();
|
||||
if normalized.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(normalized)
|
||||
}
|
||||
|
||||
/// Minimal base64 decode (no extra dep) — only needs to decode the user ID portion
|
||||
#[allow(clippy::cast_possible_truncation)]
|
||||
fn base64_decode(input: &str) -> Option<String> {
|
||||
let padded = match input.len() % 4 {
|
||||
2 => format!("{input}=="),
|
||||
3 => format!("{input}="),
|
||||
_ => input.to_string(),
|
||||
};
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
let chars: Vec<u8> = padded.bytes().collect();
|
||||
|
||||
for chunk in chars.chunks(4) {
|
||||
if chunk.len() < 4 {
|
||||
break;
|
||||
}
|
||||
|
||||
let mut v = [0usize; 4];
|
||||
for (i, &b) in chunk.iter().enumerate() {
|
||||
if b == b'=' {
|
||||
v[i] = 0;
|
||||
} else {
|
||||
v[i] = BASE64_ALPHABET.iter().position(|&a| a == b)?;
|
||||
}
|
||||
}
|
||||
|
||||
bytes.push(((v[0] << 2) | (v[1] >> 4)) as u8);
|
||||
if chunk[2] != b'=' {
|
||||
bytes.push((((v[1] & 0xF) << 4) | (v[2] >> 2)) as u8);
|
||||
}
|
||||
if chunk[3] != b'=' {
|
||||
bytes.push((((v[2] & 0x3) << 6) | v[3]) as u8);
|
||||
}
|
||||
}
|
||||
|
||||
String::from_utf8(bytes).ok()
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for DiscordChannel {
|
||||
fn name(&self) -> &str {
|
||||
"discord"
|
||||
}
|
||||
|
||||
/// Recipient-less proactive sends (cron/heartbeat) deliver to the bot's
|
||||
/// configured default `channel_id`. `None` when unconfigured, so proactive
|
||||
/// routing skips Discord rather than letting `send` bail on an empty target
|
||||
/// (#3794 review — Codex P2).
|
||||
fn proactive_target(&self) -> Option<String> {
|
||||
self.channel_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(str::to_string)
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> anyhow::Result<()> {
|
||||
// Resolve the target channel: explicit recipient (replies) or the
|
||||
// configured default channel (recipient-less proactive sends). Bail with
|
||||
// a clear error rather than POSTing to an empty channel id (#3794 review).
|
||||
let Some(recipient) =
|
||||
Self::resolve_recipient(&message.recipient, self.channel_id.as_deref())
|
||||
else {
|
||||
anyhow::bail!(
|
||||
"Discord send: no target channel — message had no recipient and no channel_id is configured"
|
||||
);
|
||||
};
|
||||
|
||||
let chunks = split_message_for_discord(&message.content);
|
||||
|
||||
for (i, chunk) in chunks.iter().enumerate() {
|
||||
let url = format!("https://discord.com/api/v10/channels/{recipient}/messages");
|
||||
|
||||
let body = json!({ "content": chunk });
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bot {}", self.bot_token))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let err = resp
|
||||
.text()
|
||||
.await
|
||||
.unwrap_or_else(|e| format!("<failed to read response body: {e}>"));
|
||||
anyhow::bail!("Discord send message failed ({status}): {err}");
|
||||
}
|
||||
|
||||
// Add a small delay between chunks to avoid rate limiting
|
||||
if i < chunks.len() - 1 {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(500)).await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_lines)]
|
||||
async fn listen(&self, tx: tokio::sync::mpsc::Sender<ChannelMessage>) -> anyhow::Result<()> {
|
||||
let bot_user_id = Self::bot_user_id_from_token(&self.bot_token).unwrap_or_default();
|
||||
|
||||
// Get Gateway URL
|
||||
let gw_resp: serde_json::Value = self
|
||||
.http_client()
|
||||
.get("https://discord.com/api/v10/gateway/bot")
|
||||
.header("Authorization", format!("Bot {}", self.bot_token))
|
||||
.send()
|
||||
.await?
|
||||
.json()
|
||||
.await?;
|
||||
|
||||
let gw_url = gw_resp
|
||||
.get("url")
|
||||
.and_then(|u| u.as_str())
|
||||
.unwrap_or("wss://gateway.discord.gg");
|
||||
|
||||
let ws_url = format!("{gw_url}/?v=10&encoding=json");
|
||||
tracing::info!("Discord: connecting to gateway...");
|
||||
|
||||
let (ws_stream, _) = tokio_tungstenite::connect_async(&ws_url).await?;
|
||||
let (mut write, mut read) = ws_stream.split();
|
||||
|
||||
// Read Hello (opcode 10)
|
||||
let hello = read.next().await.ok_or(anyhow::anyhow!("No hello"))??;
|
||||
let hello_data: serde_json::Value = serde_json::from_str(&hello.to_string())?;
|
||||
let heartbeat_interval = hello_data
|
||||
.get("d")
|
||||
.and_then(|d| d.get("heartbeat_interval"))
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.unwrap_or(41250);
|
||||
|
||||
// Send Identify (opcode 2)
|
||||
let identify = json!({
|
||||
"op": 2,
|
||||
"d": {
|
||||
"token": self.bot_token,
|
||||
"intents": 37377, // GUILDS | GUILD_MESSAGES | MESSAGE_CONTENT | DIRECT_MESSAGES
|
||||
"properties": {
|
||||
"os": "linux",
|
||||
"browser": "openhuman",
|
||||
"device": "openhuman"
|
||||
}
|
||||
}
|
||||
});
|
||||
write.send(Message::Text(identify.to_string())).await?;
|
||||
|
||||
tracing::info!("Discord: connected and identified");
|
||||
|
||||
// Track the last sequence number for heartbeats and resume.
|
||||
// Only accessed in the select! loop below, so a plain i64 suffices.
|
||||
let mut sequence: i64 = -1;
|
||||
|
||||
// Spawn heartbeat timer — sends a tick signal, actual heartbeat
|
||||
// is assembled in the select! loop where `sequence` lives.
|
||||
let (hb_tx, mut hb_rx) = tokio::sync::mpsc::channel::<()>(1);
|
||||
let hb_interval = heartbeat_interval;
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(std::time::Duration::from_millis(hb_interval));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
if hb_tx.send(()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let guild_filter = self.guild_id.clone();
|
||||
let channel_filter = self.channel_id.clone();
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = hb_rx.recv() => {
|
||||
let d = if sequence >= 0 { json!(sequence) } else { json!(null) };
|
||||
let hb = json!({"op": 1, "d": d});
|
||||
if write.send(Message::Text(hb.to_string())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
msg = read.next() => {
|
||||
let msg = match msg {
|
||||
Some(Ok(Message::Text(t))) => t,
|
||||
Some(Ok(Message::Close(_))) | None => break,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let event: serde_json::Value = match serde_json::from_str(msg.as_ref()) {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
// Track sequence number from all dispatch events
|
||||
if let Some(s) = event.get("s").and_then(serde_json::Value::as_i64) {
|
||||
sequence = s;
|
||||
}
|
||||
|
||||
let op = event.get("op").and_then(serde_json::Value::as_u64).unwrap_or(0);
|
||||
|
||||
match op {
|
||||
// Op 1: Server requests an immediate heartbeat
|
||||
1 => {
|
||||
let d = if sequence >= 0 { json!(sequence) } else { json!(null) };
|
||||
let hb = json!({"op": 1, "d": d});
|
||||
if write.send(Message::Text(hb.to_string())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
// Op 7: Reconnect
|
||||
7 => {
|
||||
tracing::warn!("Discord: received Reconnect (op 7), closing for restart");
|
||||
break;
|
||||
}
|
||||
// Op 9: Invalid Session
|
||||
9 => {
|
||||
tracing::warn!("Discord: received Invalid Session (op 9), closing for restart");
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
// Only handle MESSAGE_CREATE (opcode 0, type "MESSAGE_CREATE")
|
||||
let event_type = event.get("t").and_then(|t| t.as_str()).unwrap_or("");
|
||||
if event_type != "MESSAGE_CREATE" {
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(d) = event.get("d") else {
|
||||
continue;
|
||||
};
|
||||
|
||||
// Skip messages from the bot itself
|
||||
let author_id = d.get("author").and_then(|a| a.get("id")).and_then(|i| i.as_str()).unwrap_or("");
|
||||
if author_id == bot_user_id {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Skip bot messages (unless listen_to_bots is enabled)
|
||||
if !self.listen_to_bots && d.get("author").and_then(|a| a.get("bot")).and_then(serde_json::Value::as_bool).unwrap_or(false) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Sender validation
|
||||
if !self.is_user_allowed(author_id) {
|
||||
tracing::warn!("Discord: ignoring message from unauthorized user: {author_id}");
|
||||
continue;
|
||||
}
|
||||
|
||||
// Guild filter + DM scoping (#3794 review — Codex P1)
|
||||
if !Self::passes_guild_scope(
|
||||
guild_filter.as_deref(),
|
||||
d.get("guild_id").and_then(serde_json::Value::as_str),
|
||||
self.allowed_users.is_empty(),
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Channel filter — only process messages from the configured channel
|
||||
if let Some(ref cid) = channel_filter {
|
||||
let msg_channel = d.get("channel_id").and_then(serde_json::Value::as_str).unwrap_or("");
|
||||
if msg_channel != cid {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
let content = d.get("content").and_then(|c| c.as_str()).unwrap_or("");
|
||||
let Some(clean_content) =
|
||||
normalize_incoming_content(content, self.mention_only, &bot_user_id)
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
let message_id = d.get("id").and_then(|i| i.as_str()).unwrap_or("");
|
||||
let channel_id = d.get("channel_id").and_then(|c| c.as_str()).unwrap_or("").to_string();
|
||||
|
||||
let channel_msg = ChannelMessage {
|
||||
id: if message_id.is_empty() {
|
||||
format!("discord_{}", Uuid::new_v4())
|
||||
} else {
|
||||
format!("discord_{message_id}")
|
||||
},
|
||||
sender: author_id.to_string(),
|
||||
reply_target: if channel_id.is_empty() {
|
||||
author_id.to_string()
|
||||
} else {
|
||||
channel_id.clone()
|
||||
},
|
||||
content: clean_content,
|
||||
channel: "discord".to_string(),
|
||||
timestamp: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
thread_ts: None,
|
||||
};
|
||||
|
||||
if tx.send(channel_msg).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
self.http_client()
|
||||
.get("https://discord.com/api/v10/users/@me")
|
||||
.header("Authorization", format!("Bot {}", self.bot_token))
|
||||
.send()
|
||||
.await
|
||||
.map(|r| r.status().is_success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
async fn start_typing(&self, recipient: &str) -> anyhow::Result<()> {
|
||||
self.stop_typing(recipient).await?;
|
||||
|
||||
let client = self.http_client();
|
||||
let token = self.bot_token.clone();
|
||||
let channel_id = recipient.to_string();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
let url = format!("https://discord.com/api/v10/channels/{channel_id}/typing");
|
||||
loop {
|
||||
let _ = client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bot {token}"))
|
||||
.send()
|
||||
.await;
|
||||
tokio::time::sleep(std::time::Duration::from_secs(8)).await;
|
||||
}
|
||||
});
|
||||
|
||||
let mut guard = self.typing_handle.lock();
|
||||
*guard = Some(handle);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stop_typing(&self, _recipient: &str) -> anyhow::Result<()> {
|
||||
let mut guard = self.typing_handle.lock();
|
||||
if let Some(handle) = guard.take() {
|
||||
handle.abort();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "channel_tests.rs"]
|
||||
mod tests;
|
||||
@@ -1,581 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
fn discord_units(text: &str) -> usize {
|
||||
text.encode_utf16().count()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discord_channel_name() {
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec![], false, false);
|
||||
assert_eq!(ch.name(), "discord");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn base64_decode_bot_id() {
|
||||
// "MTIzNDU2" decodes to "123456"
|
||||
let decoded = base64_decode("MTIzNDU2");
|
||||
assert_eq!(decoded, Some("123456".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bot_user_id_extraction() {
|
||||
// Token format: base64(user_id).timestamp.hmac
|
||||
let token = "MTIzNDU2.fake.hmac";
|
||||
let id = DiscordChannel::bot_user_id_from_token(token);
|
||||
assert_eq!(id, Some("123456".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_allowlist_allows_everyone() {
|
||||
// Issue #3712: an unconfigured (empty) allowlist must apply no per-user
|
||||
// restriction — otherwise a UI-connected bot silently ignores every message
|
||||
// and never replies. Scope is still enforced by guild/channel filters.
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec![], false, false);
|
||||
assert!(ch.is_user_allowed("12345"));
|
||||
assert!(ch.is_user_allowed("anyone"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wildcard_allows_everyone() {
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec!["*".into()], false, false);
|
||||
assert!(ch.is_user_allowed("12345"));
|
||||
assert!(ch.is_user_allowed("anyone"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn specific_allowlist_filters() {
|
||||
let ch = DiscordChannel::new(
|
||||
"fake".into(),
|
||||
None,
|
||||
None,
|
||||
vec!["111".into(), "222".into()],
|
||||
false,
|
||||
false,
|
||||
);
|
||||
assert!(ch.is_user_allowed("111"));
|
||||
assert!(ch.is_user_allowed("222"));
|
||||
assert!(!ch.is_user_allowed("333"));
|
||||
assert!(!ch.is_user_allowed("unknown"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allowlist_is_exact_match_not_substring() {
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec!["111".into()], false, false);
|
||||
assert!(!ch.is_user_allowed("1111"));
|
||||
assert!(!ch.is_user_allowed("11"));
|
||||
assert!(!ch.is_user_allowed("0111"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allowlist_empty_string_user_id() {
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec!["111".into()], false, false);
|
||||
assert!(!ch.is_user_allowed(""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allowlist_with_wildcard_and_specific() {
|
||||
let ch = DiscordChannel::new(
|
||||
"fake".into(),
|
||||
None,
|
||||
None,
|
||||
vec!["111".into(), "*".into()],
|
||||
false,
|
||||
false,
|
||||
);
|
||||
assert!(ch.is_user_allowed("111"));
|
||||
assert!(ch.is_user_allowed("anyone_else"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allowlist_case_sensitive() {
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec!["ABC".into()], false, false);
|
||||
assert!(ch.is_user_allowed("ABC"));
|
||||
assert!(!ch.is_user_allowed("abc"));
|
||||
assert!(!ch.is_user_allowed("Abc"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn base64_decode_empty_string() {
|
||||
let decoded = base64_decode("");
|
||||
assert_eq!(decoded, Some(String::new()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn base64_decode_invalid_chars() {
|
||||
let decoded = base64_decode("!!!!");
|
||||
assert!(decoded.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bot_user_id_from_empty_token() {
|
||||
let id = DiscordChannel::bot_user_id_from_token("");
|
||||
assert_eq!(id, Some(String::new()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn contains_bot_mention_supports_plain_and_nick_forms() {
|
||||
assert!(contains_bot_mention("hi <@12345>", "12345"));
|
||||
assert!(contains_bot_mention("hi <@!12345>", "12345"));
|
||||
assert!(!contains_bot_mention("hi <@99999>", "12345"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_incoming_content_requires_mention_when_enabled() {
|
||||
let cleaned = normalize_incoming_content("hello there", true, "12345");
|
||||
assert!(cleaned.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_incoming_content_strips_mentions_and_trims() {
|
||||
let cleaned = normalize_incoming_content(" <@!12345> run status ", true, "12345");
|
||||
assert_eq!(cleaned.as_deref(), Some("run status"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_incoming_content_rejects_empty_after_strip() {
|
||||
let cleaned = normalize_incoming_content("<@12345>", true, "12345");
|
||||
assert!(cleaned.is_none());
|
||||
}
|
||||
|
||||
// Message splitting tests
|
||||
|
||||
#[test]
|
||||
fn split_empty_message() {
|
||||
let chunks = split_message_for_discord("");
|
||||
assert_eq!(chunks, vec![""]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_short_message_under_limit() {
|
||||
let msg = "Hello, world!";
|
||||
let chunks = split_message_for_discord(msg);
|
||||
assert_eq!(chunks, vec![msg]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_message_exactly_2000_chars() {
|
||||
let msg = "a".repeat(DISCORD_MAX_MESSAGE_LENGTH);
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0].chars().count(), DISCORD_MAX_MESSAGE_LENGTH);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_message_just_over_limit() {
|
||||
let msg = "a".repeat(DISCORD_MAX_MESSAGE_LENGTH + 1);
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
assert_eq!(chunks.len(), 2);
|
||||
assert_eq!(chunks[0].chars().count(), DISCORD_MAX_MESSAGE_LENGTH);
|
||||
assert_eq!(chunks[1].chars().count(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_very_long_message() {
|
||||
let msg = "word ".repeat(2000); // 10000 characters (5 chars per "word ")
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
// The shared chunker prefers whitespace boundaries, so it may produce more
|
||||
// than the mathematical minimum while preserving Discord's UTF-16 limit.
|
||||
assert!(chunks.len() > 1);
|
||||
assert!(chunks
|
||||
.iter()
|
||||
.all(|chunk| discord_units(chunk) <= DISCORD_MAX_MESSAGE_LENGTH));
|
||||
// Verify total content is preserved
|
||||
let reconstructed = chunks.concat();
|
||||
assert_eq!(reconstructed, msg);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_prefer_newline_break() {
|
||||
let msg = format!("{}\n{}", "a".repeat(1500), "b".repeat(500));
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
// Should split at the newline
|
||||
assert_eq!(chunks.len(), 2);
|
||||
assert!(chunks[1].starts_with('\n'));
|
||||
assert!(chunks[1].trim_start_matches('\n').starts_with('b'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_prefer_space_break() {
|
||||
let msg = format!("{} {}", "a".repeat(1500), "b".repeat(600));
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
assert_eq!(chunks.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_without_good_break_points_hard_split() {
|
||||
// No spaces or newlines - should hard split at 2000
|
||||
let msg = "a".repeat(5000);
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
assert_eq!(chunks.len(), 3);
|
||||
assert_eq!(chunks[0].chars().count(), DISCORD_MAX_MESSAGE_LENGTH);
|
||||
assert_eq!(chunks[1].chars().count(), DISCORD_MAX_MESSAGE_LENGTH);
|
||||
assert_eq!(chunks[2].chars().count(), 1000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_multiple_breaks() {
|
||||
// Create a message with multiple newlines
|
||||
let part1 = "a".repeat(900);
|
||||
let part2 = "b".repeat(900);
|
||||
let part3 = "c".repeat(900);
|
||||
let msg = format!("{part1}\n{part2}\n{part3}");
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
// Should split into 2 chunks (first two parts + third part)
|
||||
assert_eq!(chunks.len(), 2);
|
||||
assert!(chunks[0].chars().count() <= DISCORD_MAX_MESSAGE_LENGTH);
|
||||
assert!(chunks[1].chars().count() <= DISCORD_MAX_MESSAGE_LENGTH);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_preserves_content() {
|
||||
let original = "Hello world! This is a test message with some content. ".repeat(200);
|
||||
let chunks = split_message_for_discord(&original);
|
||||
let reconstructed = chunks.concat();
|
||||
assert_eq!(reconstructed, original);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_unicode_content() {
|
||||
// Test with emoji and multi-byte characters
|
||||
let msg = "🦀 Rust is awesome! ".repeat(500);
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
// All chunks should be valid UTF-8
|
||||
for chunk in &chunks {
|
||||
assert!(std::str::from_utf8(chunk.as_bytes()).is_ok());
|
||||
assert!(chunk.chars().count() <= DISCORD_MAX_MESSAGE_LENGTH);
|
||||
}
|
||||
// Reconstruct and verify
|
||||
let reconstructed = chunks.concat();
|
||||
assert_eq!(reconstructed, msg);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_newline_too_close_to_end() {
|
||||
// If newline is in the first half, don't use it - use space instead or hard split
|
||||
let msg = format!("{}\n{}", "a".repeat(1900), "b".repeat(500));
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
// Should split at newline since it's in the second half of the window
|
||||
assert_eq!(chunks.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_multibyte_only_content_without_panics() {
|
||||
let msg = "🦀".repeat(2500);
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
assert_eq!(chunks.len(), 3);
|
||||
assert!(chunks
|
||||
.iter()
|
||||
.all(|chunk| discord_units(chunk) <= DISCORD_MAX_MESSAGE_LENGTH));
|
||||
let reconstructed = chunks.concat();
|
||||
assert_eq!(reconstructed, msg);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_chunks_always_within_discord_limit() {
|
||||
let msg = "x".repeat(12_345);
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
assert!(chunks
|
||||
.iter()
|
||||
.all(|chunk| chunk.chars().count() <= DISCORD_MAX_MESSAGE_LENGTH));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_message_with_multiple_newlines() {
|
||||
let msg = "Line 1\nLine 2\nLine 3\n".repeat(1000);
|
||||
let chunks = split_message_for_discord(&msg);
|
||||
assert!(chunks.len() > 1);
|
||||
let reconstructed = chunks.concat();
|
||||
assert_eq!(reconstructed, msg);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn typing_handle_starts_as_none() {
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec![], false, false);
|
||||
let guard = ch.typing_handle.lock();
|
||||
assert!(guard.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn start_typing_sets_handle() {
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec![], false, false);
|
||||
let _ = ch.start_typing("123456").await;
|
||||
let guard = ch.typing_handle.lock();
|
||||
assert!(guard.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stop_typing_clears_handle() {
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec![], false, false);
|
||||
let _ = ch.start_typing("123456").await;
|
||||
let _ = ch.stop_typing("123456").await;
|
||||
let guard = ch.typing_handle.lock();
|
||||
assert!(guard.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stop_typing_is_idempotent() {
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec![], false, false);
|
||||
assert!(ch.stop_typing("123456").await.is_ok());
|
||||
assert!(ch.stop_typing("123456").await.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn start_typing_replaces_existing_task() {
|
||||
let ch = DiscordChannel::new("fake".into(), None, None, vec![], false, false);
|
||||
let _ = ch.start_typing("111").await;
|
||||
let _ = ch.start_typing("222").await;
|
||||
let guard = ch.typing_handle.lock();
|
||||
assert!(guard.is_some());
|
||||
}
|
||||
|
||||
// ── Message ID edge cases ─────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn discord_message_id_format_includes_discord_prefix() {
|
||||
// Verify that message IDs follow the format: discord_{message_id}
|
||||
let message_id = "123456789012345678";
|
||||
let expected_id = format!("discord_{message_id}");
|
||||
assert_eq!(expected_id, "discord_123456789012345678");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discord_message_id_is_deterministic() {
|
||||
// Same message_id = same ID (prevents duplicates after restart)
|
||||
let message_id = "123456789012345678";
|
||||
let id1 = format!("discord_{message_id}");
|
||||
let id2 = format!("discord_{message_id}");
|
||||
assert_eq!(id1, id2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discord_message_id_different_message_different_id() {
|
||||
// Different message IDs produce different IDs
|
||||
let id1 = "discord_123456789012345678".to_string();
|
||||
let id2 = "discord_987654321098765432".to_string();
|
||||
assert_ne!(id1, id2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discord_message_id_uses_snowflake_id() {
|
||||
// Discord snowflake IDs are numeric strings
|
||||
let message_id = "123456789012345678"; // Typical snowflake format
|
||||
let id = format!("discord_{message_id}");
|
||||
assert!(id.starts_with("discord_"));
|
||||
// Snowflake IDs are numeric
|
||||
assert!(message_id.chars().all(|c| c.is_ascii_digit()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discord_message_id_fallback_to_uuid_on_empty() {
|
||||
// Edge case: empty message_id falls back to UUID
|
||||
let message_id = "";
|
||||
let id = if message_id.is_empty() {
|
||||
format!("discord_{}", uuid::Uuid::new_v4())
|
||||
} else {
|
||||
format!("discord_{message_id}")
|
||||
};
|
||||
assert!(id.starts_with("discord_"));
|
||||
// Should have UUID dashes
|
||||
assert!(id.contains('-'));
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// TG6: Channel platform limit edge cases for Discord (2000 char limit)
|
||||
// Prevents: Pattern 6 — issues #574, #499
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn split_message_code_block_at_boundary() {
|
||||
// Code block that spans the split boundary
|
||||
let mut msg = String::new();
|
||||
msg.push_str("```rust\n");
|
||||
msg.push_str(&"x".repeat(1990));
|
||||
msg.push_str("\n```\nMore text after code block");
|
||||
let parts = split_message_for_discord(&msg);
|
||||
assert!(
|
||||
parts.len() >= 2,
|
||||
"code block spanning boundary should split"
|
||||
);
|
||||
for part in &parts {
|
||||
assert!(
|
||||
part.len() <= DISCORD_MAX_MESSAGE_LENGTH,
|
||||
"each part must be <= {DISCORD_MAX_MESSAGE_LENGTH}, got {}",
|
||||
part.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_message_single_long_word_exceeds_limit() {
|
||||
// A single word longer than 2000 chars must be hard-split
|
||||
let long_word = "a".repeat(2500);
|
||||
let parts = split_message_for_discord(&long_word);
|
||||
assert!(parts.len() >= 2, "word exceeding limit must be split");
|
||||
for part in &parts {
|
||||
assert!(
|
||||
part.len() <= DISCORD_MAX_MESSAGE_LENGTH,
|
||||
"hard-split part must be <= {DISCORD_MAX_MESSAGE_LENGTH}, got {}",
|
||||
part.len()
|
||||
);
|
||||
}
|
||||
// Reassembled content should match original
|
||||
let reassembled: String = parts.join("");
|
||||
assert_eq!(reassembled, long_word);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_message_exactly_at_limit_no_split() {
|
||||
let msg = "a".repeat(DISCORD_MAX_MESSAGE_LENGTH);
|
||||
let parts = split_message_for_discord(&msg);
|
||||
assert_eq!(parts.len(), 1, "message exactly at limit should not split");
|
||||
assert_eq!(parts[0].len(), DISCORD_MAX_MESSAGE_LENGTH);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_message_one_over_limit_splits() {
|
||||
let msg = "a".repeat(DISCORD_MAX_MESSAGE_LENGTH + 1);
|
||||
let parts = split_message_for_discord(&msg);
|
||||
assert!(parts.len() >= 2, "message 1 char over limit must split");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_message_many_short_lines() {
|
||||
// Many short lines should be batched into chunks under the limit
|
||||
let msg: String = (0..500).map(|i| format!("line {i}\n")).collect();
|
||||
let parts = split_message_for_discord(&msg);
|
||||
for part in &parts {
|
||||
assert!(
|
||||
part.len() <= DISCORD_MAX_MESSAGE_LENGTH,
|
||||
"short-line batch must be <= limit"
|
||||
);
|
||||
}
|
||||
// All content should be preserved
|
||||
let reassembled: String = parts.join("");
|
||||
assert_eq!(reassembled.trim(), msg.trim());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_message_only_whitespace() {
|
||||
let msg = " \n\n\t ";
|
||||
let parts = split_message_for_discord(msg);
|
||||
// Should handle gracefully without panic
|
||||
assert!(parts.len() <= 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_message_emoji_at_boundary() {
|
||||
// Emoji are multi-byte; ensure we don't split mid-emoji
|
||||
let mut msg = "a".repeat(1998);
|
||||
msg.push_str("🎉🎊"); // 2 emoji at the boundary (2000 chars total)
|
||||
let parts = split_message_for_discord(&msg);
|
||||
for part in &parts {
|
||||
// The function splits on character count, not byte count
|
||||
assert!(
|
||||
part.chars().count() <= DISCORD_MAX_MESSAGE_LENGTH,
|
||||
"emoji boundary split must respect limit"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_message_consecutive_newlines_at_boundary() {
|
||||
let mut msg = "a".repeat(1995);
|
||||
msg.push_str("\n\n\n\n\n");
|
||||
msg.push_str(&"b".repeat(100));
|
||||
let parts = split_message_for_discord(&msg);
|
||||
for part in &parts {
|
||||
assert!(part.len() <= DISCORD_MAX_MESSAGE_LENGTH);
|
||||
}
|
||||
}
|
||||
|
||||
// ── channel_id field tests ───────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn channel_id_stored_in_struct() {
|
||||
let ch = DiscordChannel::new(
|
||||
"token".into(),
|
||||
Some("guild1".into()),
|
||||
Some("channel1".into()),
|
||||
vec![],
|
||||
false,
|
||||
false,
|
||||
);
|
||||
assert_eq!(ch.channel_id.as_deref(), Some("channel1"));
|
||||
assert_eq!(ch.guild_id.as_deref(), Some("guild1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_id_defaults_to_none() {
|
||||
let ch = DiscordChannel::new("token".into(), None, None, vec![], false, false);
|
||||
assert!(ch.channel_id.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passes_guild_scope_covers_guild_dm_and_unscoped_cases() {
|
||||
// No configured guild → everything passes (filter inactive).
|
||||
assert!(DiscordChannel::passes_guild_scope(None, Some("g1"), true));
|
||||
assert!(DiscordChannel::passes_guild_scope(None, None, true));
|
||||
// Configured guild: same guild passes, other guild blocked.
|
||||
assert!(DiscordChannel::passes_guild_scope(
|
||||
Some("g1"),
|
||||
Some("g1"),
|
||||
true
|
||||
));
|
||||
assert!(!DiscordChannel::passes_guild_scope(
|
||||
Some("g1"),
|
||||
Some("g2"),
|
||||
true
|
||||
));
|
||||
// #3794 review (Codex P1): DM (no guild_id) under guild scope is blocked
|
||||
// with a blank allowlist, allowed with an explicit one.
|
||||
assert!(!DiscordChannel::passes_guild_scope(Some("g1"), None, true));
|
||||
assert!(DiscordChannel::passes_guild_scope(Some("g1"), None, false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_recipient_prefers_explicit_then_configured_channel() {
|
||||
// #3794 review (Codex P2): recipient-less proactive sends fall back to the
|
||||
// configured channel_id; an explicit recipient always wins.
|
||||
assert_eq!(
|
||||
DiscordChannel::resolve_recipient("123", Some("999")),
|
||||
Some("123")
|
||||
);
|
||||
assert_eq!(
|
||||
DiscordChannel::resolve_recipient("", Some("999")),
|
||||
Some("999")
|
||||
);
|
||||
// Neither available → None, so the caller errors instead of POSTing to "".
|
||||
assert_eq!(DiscordChannel::resolve_recipient("", None), None);
|
||||
assert_eq!(DiscordChannel::resolve_recipient("", Some("")), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proactive_target_uses_configured_channel_id() {
|
||||
use crate::openhuman::channels::traits::Channel;
|
||||
|
||||
// Configured channel_id ⇒ recipient-less proactive sends have a target.
|
||||
let with_channel = DiscordChannel::new(
|
||||
"fake".into(),
|
||||
None,
|
||||
Some("12345".into()),
|
||||
vec![],
|
||||
false,
|
||||
false,
|
||||
);
|
||||
assert_eq!(with_channel.proactive_target(), Some("12345".to_string()));
|
||||
|
||||
// No channel_id ⇒ None, so proactive routing skips Discord (#3794 Codex P2).
|
||||
let no_channel = DiscordChannel::new("fake".into(), None, None, vec![], false, false);
|
||||
assert_eq!(no_channel.proactive_target(), None);
|
||||
|
||||
// Whitespace-only channel_id is treated as unset.
|
||||
let blank_channel = DiscordChannel::new(
|
||||
"fake".into(),
|
||||
None,
|
||||
Some(" ".into()),
|
||||
vec![],
|
||||
false,
|
||||
false,
|
||||
);
|
||||
assert_eq!(blank_channel.proactive_target(), None);
|
||||
}
|
||||
@@ -1,4 +0,0 @@
|
||||
pub mod api;
|
||||
pub mod channel;
|
||||
|
||||
pub use channel::DiscordChannel;
|
||||
@@ -1,553 +1 @@
|
||||
#![allow(clippy::uninlined_format_args)]
|
||||
#![allow(clippy::map_unwrap_or)]
|
||||
#![allow(clippy::redundant_closure_for_method_calls)]
|
||||
#![allow(clippy::cast_lossless)]
|
||||
#![allow(clippy::trim_split_whitespace)]
|
||||
#![allow(clippy::doc_link_with_quotes)]
|
||||
#![allow(clippy::doc_markdown)]
|
||||
#![allow(clippy::too_many_lines)]
|
||||
#![allow(clippy::unnecessary_map_or)]
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_imap::extensions::idle::IdleResponse;
|
||||
use async_imap::types::Fetch;
|
||||
use async_imap::Session;
|
||||
use async_trait::async_trait;
|
||||
use futures::TryStreamExt;
|
||||
use lettre::message::{header::ContentType, Attachment, MultiPart, SinglePart};
|
||||
use lettre::transport::smtp::authentication::Credentials;
|
||||
use lettre::{Message, SmtpTransport, Transport};
|
||||
use mail_parser::{MessageParser, MimeHeaders};
|
||||
use rustls::{ClientConfig, RootCertStore};
|
||||
use rustls_pki_types::DnsName;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::{mpsc, Mutex};
|
||||
use tokio::time::{sleep, timeout};
|
||||
use tokio_rustls::client::TlsStream;
|
||||
use tokio_rustls::TlsConnector;
|
||||
use tracing::{debug, error, info, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
pub use crate::openhuman::config::schema::EmailConfig;
|
||||
|
||||
type ImapSession = Session<TlsStream<TcpStream>>;
|
||||
|
||||
/// Email channel — IMAP IDLE for instant push notifications, SMTP for outbound
|
||||
pub struct EmailChannel {
|
||||
pub config: EmailConfig,
|
||||
seen_messages: Arc<Mutex<HashSet<String>>>,
|
||||
}
|
||||
|
||||
impl EmailChannel {
|
||||
pub fn new(config: EmailConfig) -> Self {
|
||||
Self {
|
||||
config,
|
||||
seen_messages: Arc::new(Mutex::new(HashSet::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a sender email is in the allowlist
|
||||
pub fn is_sender_allowed(&self, email: &str) -> bool {
|
||||
if self.config.allowed_senders.is_empty() {
|
||||
return false; // Empty = deny all
|
||||
}
|
||||
if self.config.allowed_senders.iter().any(|a| a == "*") {
|
||||
return true; // Wildcard = allow all
|
||||
}
|
||||
let email_lower = email.to_lowercase();
|
||||
self.config.allowed_senders.iter().any(|allowed| {
|
||||
if allowed.starts_with('@') {
|
||||
// Domain match with @ prefix: "@example.com"
|
||||
email_lower.ends_with(&allowed.to_lowercase())
|
||||
} else if allowed.contains('@') {
|
||||
// Full email address match
|
||||
allowed.eq_ignore_ascii_case(email)
|
||||
} else {
|
||||
// Domain match without @ prefix: "example.com"
|
||||
email_lower.ends_with(&format!("@{}", allowed.to_lowercase()))
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Strip HTML tags from content (basic)
|
||||
pub fn strip_html(html: &str) -> String {
|
||||
let mut result = String::new();
|
||||
let mut in_tag = false;
|
||||
for ch in html.chars() {
|
||||
match ch {
|
||||
'<' => in_tag = true,
|
||||
'>' => in_tag = false,
|
||||
_ if !in_tag => result.push(ch),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
let mut normalized = String::with_capacity(result.len());
|
||||
for word in result.split_whitespace() {
|
||||
if !normalized.is_empty() {
|
||||
normalized.push(' ');
|
||||
}
|
||||
normalized.push_str(word);
|
||||
}
|
||||
normalized
|
||||
}
|
||||
|
||||
/// Extract the sender address from a parsed email
|
||||
fn extract_sender(parsed: &mail_parser::Message) -> String {
|
||||
parsed
|
||||
.from()
|
||||
.and_then(|addr| addr.first())
|
||||
.and_then(|a| a.address())
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| "unknown".into())
|
||||
}
|
||||
|
||||
/// Extract readable text from a parsed email
|
||||
fn extract_text(parsed: &mail_parser::Message) -> String {
|
||||
if let Some(text) = parsed.body_text(0) {
|
||||
return text.to_string();
|
||||
}
|
||||
if let Some(html) = parsed.body_html(0) {
|
||||
return Self::strip_html(html.as_ref());
|
||||
}
|
||||
for part in parsed.attachments() {
|
||||
let part: &mail_parser::MessagePart = part;
|
||||
if let Some(ct) = MimeHeaders::content_type(part) {
|
||||
if ct.ctype() == "text" {
|
||||
if let Ok(text) = std::str::from_utf8(part.contents()) {
|
||||
let name = MimeHeaders::attachment_name(part).unwrap_or("file");
|
||||
return format!("[Attachment: {}]\n{}", name, text);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"(no readable content)".to_string()
|
||||
}
|
||||
|
||||
/// Connect to IMAP server with TLS and authenticate
|
||||
async fn connect_imap(&self) -> Result<ImapSession> {
|
||||
let addr = format!("{}:{}", self.config.imap_host, self.config.imap_port);
|
||||
debug!("Connecting to IMAP server at {}", addr);
|
||||
|
||||
// Connect TCP
|
||||
let tcp = TcpStream::connect(&addr).await?;
|
||||
|
||||
// Establish TLS using rustls
|
||||
let certs = RootCertStore {
|
||||
roots: webpki_roots::TLS_SERVER_ROOTS.into(),
|
||||
};
|
||||
let config = ClientConfig::builder()
|
||||
.with_root_certificates(certs)
|
||||
.with_no_client_auth();
|
||||
let tls_stream: TlsConnector = Arc::new(config).into();
|
||||
let sni: DnsName = self.config.imap_host.clone().try_into()?;
|
||||
let stream = tls_stream.connect(sni.into(), tcp).await?;
|
||||
|
||||
// Create IMAP client
|
||||
let client = async_imap::Client::new(stream);
|
||||
|
||||
// Login
|
||||
let session = client
|
||||
.login(&self.config.username, &self.config.password)
|
||||
.await
|
||||
.map_err(|(e, _)| anyhow!("IMAP login failed: {}", e))?;
|
||||
|
||||
debug!("IMAP login successful");
|
||||
Ok(session)
|
||||
}
|
||||
|
||||
/// Fetch and process unseen messages from the selected mailbox
|
||||
async fn fetch_unseen(&self, session: &mut ImapSession) -> Result<Vec<ParsedEmail>> {
|
||||
// Search for unseen messages
|
||||
let uids = session.uid_search("UNSEEN").await?;
|
||||
if uids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
debug!("Found {} unseen messages", uids.len());
|
||||
|
||||
let mut results = Vec::new();
|
||||
let uid_set: String = uids
|
||||
.iter()
|
||||
.map(|u| u.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
|
||||
// Fetch message bodies
|
||||
let messages = session.uid_fetch(&uid_set, "RFC822").await?;
|
||||
let messages: Vec<Fetch> = messages.try_collect().await?;
|
||||
|
||||
for msg in messages {
|
||||
let uid = msg.uid.unwrap_or(0);
|
||||
if let Some(body) = msg.body() {
|
||||
if let Some(parsed) = MessageParser::default().parse(body) {
|
||||
let sender = Self::extract_sender(&parsed);
|
||||
let subject = parsed.subject().unwrap_or("(no subject)").to_string();
|
||||
let body_text = Self::extract_text(&parsed);
|
||||
let content = format!("Subject: {}\n\n{}", subject, body_text);
|
||||
let msg_id = parsed
|
||||
.message_id()
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| format!("gen-{}", Uuid::new_v4()));
|
||||
|
||||
#[allow(clippy::cast_sign_loss)]
|
||||
let ts = parsed
|
||||
.date()
|
||||
.map(|d| {
|
||||
let naive = chrono::NaiveDate::from_ymd_opt(
|
||||
d.year as i32,
|
||||
u32::from(d.month),
|
||||
u32::from(d.day),
|
||||
)
|
||||
.and_then(|date| {
|
||||
date.and_hms_opt(
|
||||
u32::from(d.hour),
|
||||
u32::from(d.minute),
|
||||
u32::from(d.second),
|
||||
)
|
||||
});
|
||||
naive.map_or(0, |n| n.and_utc().timestamp() as u64)
|
||||
})
|
||||
.unwrap_or_else(|| {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
});
|
||||
|
||||
results.push(ParsedEmail {
|
||||
_uid: uid,
|
||||
msg_id,
|
||||
sender,
|
||||
content,
|
||||
timestamp: ts,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Mark fetched messages as seen
|
||||
if !results.is_empty() {
|
||||
let _ = session
|
||||
.uid_store(&uid_set, "+FLAGS (\\Seen)")
|
||||
.await?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await;
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Run the IDLE loop, returning when a new message arrives or timeout
|
||||
/// Note: IDLE consumes the session and returns it via done()
|
||||
async fn wait_for_changes(
|
||||
&self,
|
||||
session: ImapSession,
|
||||
) -> Result<(IdleWaitResult, ImapSession)> {
|
||||
let idle_timeout = Duration::from_secs(self.config.idle_timeout_secs);
|
||||
|
||||
// Start IDLE mode - this consumes the session
|
||||
let mut idle = session.idle();
|
||||
idle.init().await?;
|
||||
|
||||
debug!("Entering IMAP IDLE mode");
|
||||
|
||||
// wait() returns (future, stop_source) - we only need the future
|
||||
let (wait_future, _stop_source) = idle.wait();
|
||||
|
||||
// Wait for server notification or timeout
|
||||
let result = timeout(idle_timeout, wait_future).await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(response)) => {
|
||||
debug!("IDLE response: {:?}", response);
|
||||
// Done with IDLE, return session to normal mode
|
||||
let session = idle.done().await?;
|
||||
let wait_result = match response {
|
||||
IdleResponse::NewData(_) => IdleWaitResult::NewMail,
|
||||
IdleResponse::Timeout => IdleWaitResult::Timeout,
|
||||
IdleResponse::ManualInterrupt => IdleWaitResult::Interrupted,
|
||||
};
|
||||
Ok((wait_result, session))
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
// Try to clean up IDLE state
|
||||
let _ = idle.done().await;
|
||||
Err(anyhow!("IDLE error: {}", e))
|
||||
}
|
||||
Err(_) => {
|
||||
// Timeout - RFC 2177 recommends restarting IDLE every 29 minutes
|
||||
debug!("IDLE timeout reached, will re-establish");
|
||||
let session = idle.done().await?;
|
||||
Ok((IdleWaitResult::Timeout, session))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Main IDLE-based listen loop with automatic reconnection
|
||||
async fn listen_with_idle(&self, tx: mpsc::Sender<ChannelMessage>) -> Result<()> {
|
||||
let mut backoff = Duration::from_secs(1);
|
||||
let max_backoff = Duration::from_secs(60);
|
||||
|
||||
loop {
|
||||
match self.run_idle_session(&tx).await {
|
||||
Ok(()) => {
|
||||
// Clean exit (channel closed)
|
||||
return Ok(());
|
||||
}
|
||||
Err(e) => {
|
||||
error!(
|
||||
"IMAP session error: {}. Reconnecting in {:?}...",
|
||||
e, backoff
|
||||
);
|
||||
sleep(backoff).await;
|
||||
// Exponential backoff with cap
|
||||
backoff = std::cmp::min(backoff * 2, max_backoff);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Run a single IDLE session until error or clean shutdown
|
||||
async fn run_idle_session(&self, tx: &mpsc::Sender<ChannelMessage>) -> Result<()> {
|
||||
// Connect and authenticate
|
||||
let mut session = self.connect_imap().await?;
|
||||
|
||||
// Select the mailbox
|
||||
session.select(&self.config.imap_folder).await?;
|
||||
info!(
|
||||
"Email IDLE listening on {} (instant push enabled)",
|
||||
self.config.imap_folder
|
||||
);
|
||||
|
||||
// Check for existing unseen messages first
|
||||
self.process_unseen(&mut session, tx).await?;
|
||||
|
||||
loop {
|
||||
// Enter IDLE and wait for changes (consumes session, returns it via result)
|
||||
match self.wait_for_changes(session).await {
|
||||
Ok((IdleWaitResult::NewMail, returned_session)) => {
|
||||
debug!("New mail notification received");
|
||||
session = returned_session;
|
||||
self.process_unseen(&mut session, tx).await?;
|
||||
}
|
||||
Ok((IdleWaitResult::Timeout, returned_session)) => {
|
||||
// Re-check for mail after IDLE timeout (defensive)
|
||||
session = returned_session;
|
||||
self.process_unseen(&mut session, tx).await?;
|
||||
}
|
||||
Ok((IdleWaitResult::Interrupted, _)) => {
|
||||
info!("IDLE interrupted, exiting");
|
||||
return Ok(());
|
||||
}
|
||||
Err(e) => {
|
||||
// Connection likely broken, need to reconnect
|
||||
return Err(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch unseen messages and send to channel
|
||||
async fn process_unseen(
|
||||
&self,
|
||||
session: &mut ImapSession,
|
||||
tx: &mpsc::Sender<ChannelMessage>,
|
||||
) -> Result<()> {
|
||||
let messages = self.fetch_unseen(session).await?;
|
||||
|
||||
for email in messages {
|
||||
// Check allowlist
|
||||
if !self.is_sender_allowed(&email.sender) {
|
||||
warn!("Blocked email from {}", email.sender);
|
||||
continue;
|
||||
}
|
||||
|
||||
let is_new = {
|
||||
let mut seen = self.seen_messages.lock().await;
|
||||
seen.insert(email.msg_id.clone())
|
||||
};
|
||||
if !is_new {
|
||||
continue;
|
||||
}
|
||||
|
||||
let msg = ChannelMessage {
|
||||
id: email.msg_id,
|
||||
reply_target: email.sender.clone(),
|
||||
sender: email.sender,
|
||||
content: email.content,
|
||||
channel: "email".to_string(),
|
||||
timestamp: email.timestamp,
|
||||
thread_ts: None,
|
||||
};
|
||||
|
||||
if tx.send(msg).await.is_err() {
|
||||
// Channel closed, exit cleanly
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn create_smtp_transport(&self) -> Result<SmtpTransport> {
|
||||
let creds = Credentials::new(self.config.username.clone(), self.config.password.clone());
|
||||
let transport = if self.config.smtp_tls {
|
||||
SmtpTransport::relay(&self.config.smtp_host)?
|
||||
.port(self.config.smtp_port)
|
||||
.credentials(creds)
|
||||
.build()
|
||||
} else {
|
||||
SmtpTransport::builder_dangerous(&self.config.smtp_host)
|
||||
.port(self.config.smtp_port)
|
||||
.credentials(creds)
|
||||
.build()
|
||||
};
|
||||
Ok(transport)
|
||||
}
|
||||
|
||||
pub fn send_message(&self, email: Message) -> Result<()> {
|
||||
let transport = self.create_smtp_transport()?;
|
||||
transport.send(&email)?;
|
||||
info!("Email sent");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn build_plain_message(
|
||||
&self,
|
||||
recipient: &str,
|
||||
subject: &str,
|
||||
body: &str,
|
||||
) -> Result<Message> {
|
||||
Message::builder()
|
||||
.from(self.config.from_address.parse()?)
|
||||
.to(recipient.parse()?)
|
||||
.subject(subject)
|
||||
.singlepart(SinglePart::plain(body.to_string()))
|
||||
.map_err(Into::into)
|
||||
}
|
||||
|
||||
pub fn build_message_with_attachment(
|
||||
&self,
|
||||
recipient: &str,
|
||||
subject: &str,
|
||||
body: &str,
|
||||
attachment_name: &str,
|
||||
content_type: ContentType,
|
||||
attachment_bytes: Vec<u8>,
|
||||
) -> Result<Message> {
|
||||
let attachment =
|
||||
Attachment::new(attachment_name.to_string()).body(attachment_bytes, content_type);
|
||||
Message::builder()
|
||||
.from(self.config.from_address.parse()?)
|
||||
.to(recipient.parse()?)
|
||||
.subject(subject)
|
||||
.multipart(
|
||||
MultiPart::mixed()
|
||||
.singlepart(SinglePart::plain(body.to_string()))
|
||||
.singlepart(attachment),
|
||||
)
|
||||
.map_err(Into::into)
|
||||
}
|
||||
}
|
||||
|
||||
/// Internal struct for parsed email data
|
||||
struct ParsedEmail {
|
||||
_uid: u32,
|
||||
msg_id: String,
|
||||
sender: String,
|
||||
content: String,
|
||||
timestamp: u64,
|
||||
}
|
||||
|
||||
/// Result from waiting on IDLE
|
||||
enum IdleWaitResult {
|
||||
NewMail,
|
||||
Timeout,
|
||||
Interrupted,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for EmailChannel {
|
||||
fn name(&self) -> &str {
|
||||
"email"
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> Result<()> {
|
||||
// Use explicit subject if provided, otherwise fall back to legacy parsing or default
|
||||
let (subject, body) = if let Some(ref subj) = message.subject {
|
||||
(subj.as_str(), message.content.as_str())
|
||||
} else if message.content.starts_with("Subject: ") {
|
||||
if let Some(pos) = message.content.find('\n') {
|
||||
(&message.content[9..pos], message.content[pos + 1..].trim())
|
||||
} else {
|
||||
("OpenHuman Message", message.content.as_str())
|
||||
}
|
||||
} else {
|
||||
("OpenHuman Message", message.content.as_str())
|
||||
};
|
||||
|
||||
let email = self.build_plain_message(message.recipient.as_str(), subject, body)?;
|
||||
self.send_message(email)?;
|
||||
info!("Email sent to {}", message.recipient);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn listen(&self, tx: mpsc::Sender<ChannelMessage>) -> Result<()> {
|
||||
info!(
|
||||
"Starting email channel with IDLE support on {}",
|
||||
self.config.imap_folder
|
||||
);
|
||||
self.listen_with_idle(tx).await
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
// Fully async health check - attempt IMAP connection
|
||||
match timeout(Duration::from_secs(10), self.connect_imap()).await {
|
||||
Ok(Ok(mut session)) => {
|
||||
// Try to logout cleanly
|
||||
let _ = session.logout().await;
|
||||
true
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
debug!("Health check failed: {}", e);
|
||||
false
|
||||
}
|
||||
Err(_) => {
|
||||
debug!("Health check timed out");
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "email_channel_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
#[cfg(any(test, debug_assertions))]
|
||||
pub mod test_support {
|
||||
//! Debug-build helpers for raw integration tests. They exercise the email
|
||||
//! parser without opening IMAP or SMTP sockets.
|
||||
|
||||
use super::*;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ParsedEmailFixture {
|
||||
pub sender: String,
|
||||
pub text: String,
|
||||
pub subject: Option<String>,
|
||||
}
|
||||
|
||||
pub fn parse_email_fixture(raw: &[u8]) -> Option<ParsedEmailFixture> {
|
||||
let parsed = MessageParser::default().parse(raw)?;
|
||||
Some(ParsedEmailFixture {
|
||||
sender: EmailChannel::extract_sender(&parsed),
|
||||
text: EmailChannel::extract_text(&parsed),
|
||||
subject: parsed.subject().map(str::to_string),
|
||||
})
|
||||
}
|
||||
}
|
||||
pub use tinychannels::providers::email_channel::*;
|
||||
|
||||
@@ -1,535 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn default_smtp_port_uses_tls_port() {
|
||||
assert_eq!(EmailConfig::default().smtp_port, 465);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_config_default_uses_tls_smtp_defaults() {
|
||||
let config = EmailConfig::default();
|
||||
assert_eq!(config.smtp_port, 465);
|
||||
assert!(config.smtp_tls);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_idle_timeout_is_29_minutes() {
|
||||
assert_eq!(EmailConfig::default().idle_timeout_secs, 1740);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn seen_messages_starts_empty() {
|
||||
let channel = EmailChannel::new(EmailConfig::default());
|
||||
let seen = channel.seen_messages.lock().await;
|
||||
assert!(seen.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn seen_messages_tracks_unique_ids() {
|
||||
let channel = EmailChannel::new(EmailConfig::default());
|
||||
let mut seen = channel.seen_messages.lock().await;
|
||||
|
||||
assert!(seen.insert("first-id".to_string()));
|
||||
assert!(!seen.insert("first-id".to_string()));
|
||||
assert!(seen.insert("second-id".to_string()));
|
||||
assert_eq!(seen.len(), 2);
|
||||
}
|
||||
|
||||
// EmailConfig tests
|
||||
|
||||
#[test]
|
||||
fn email_config_default() {
|
||||
let config = EmailConfig::default();
|
||||
assert_eq!(config.imap_host, "");
|
||||
assert_eq!(config.imap_port, 993);
|
||||
assert_eq!(config.imap_folder, "INBOX");
|
||||
assert_eq!(config.smtp_host, "");
|
||||
assert_eq!(config.smtp_port, 465);
|
||||
assert!(config.smtp_tls);
|
||||
assert_eq!(config.username, "");
|
||||
assert_eq!(config.password, "");
|
||||
assert_eq!(config.from_address, "");
|
||||
assert_eq!(config.idle_timeout_secs, 1740);
|
||||
assert!(config.allowed_senders.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_config_custom() {
|
||||
let config = EmailConfig {
|
||||
imap_host: "imap.example.com".to_string(),
|
||||
imap_port: 993,
|
||||
imap_folder: "Archive".to_string(),
|
||||
smtp_host: "smtp.example.com".to_string(),
|
||||
smtp_port: 465,
|
||||
smtp_tls: true,
|
||||
username: "user@example.com".to_string(),
|
||||
password: "pass123".to_string(),
|
||||
from_address: "bot@example.com".to_string(),
|
||||
idle_timeout_secs: 1200,
|
||||
allowed_senders: vec!["allowed@example.com".to_string()],
|
||||
};
|
||||
assert_eq!(config.imap_host, "imap.example.com");
|
||||
assert_eq!(config.imap_folder, "Archive");
|
||||
assert_eq!(config.idle_timeout_secs, 1200);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_config_clone() {
|
||||
let config = EmailConfig {
|
||||
imap_host: "imap.test.com".to_string(),
|
||||
imap_port: 993,
|
||||
imap_folder: "INBOX".to_string(),
|
||||
smtp_host: "smtp.test.com".to_string(),
|
||||
smtp_port: 587,
|
||||
smtp_tls: true,
|
||||
username: "user@test.com".to_string(),
|
||||
password: "secret".to_string(),
|
||||
from_address: "bot@test.com".to_string(),
|
||||
idle_timeout_secs: 1740,
|
||||
allowed_senders: vec!["*".to_string()],
|
||||
};
|
||||
let cloned = config.clone();
|
||||
assert_eq!(cloned.imap_host, config.imap_host);
|
||||
assert_eq!(cloned.smtp_port, config.smtp_port);
|
||||
assert_eq!(cloned.allowed_senders, config.allowed_senders);
|
||||
}
|
||||
|
||||
// EmailChannel tests
|
||||
|
||||
#[tokio::test]
|
||||
async fn email_channel_new() {
|
||||
let config = EmailConfig::default();
|
||||
let channel = EmailChannel::new(config.clone());
|
||||
assert_eq!(channel.config.imap_host, config.imap_host);
|
||||
|
||||
let seen_guard = channel.seen_messages.lock().await;
|
||||
assert_eq!(seen_guard.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_channel_name() {
|
||||
let channel = EmailChannel::new(EmailConfig::default());
|
||||
assert_eq!(channel.name(), "email");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_plain_message_uses_subject_and_body() {
|
||||
let channel = EmailChannel::new(EmailConfig {
|
||||
from_address: "bot@example.com".to_string(),
|
||||
..Default::default()
|
||||
});
|
||||
let message = channel
|
||||
.build_plain_message("listener@example.com", "Podcast", "Hello there")
|
||||
.expect("plain message");
|
||||
let wire = String::from_utf8_lossy(&message.formatted()).to_string();
|
||||
assert!(wire.contains("Subject: Podcast"));
|
||||
assert!(wire.contains("Hello there"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_message_with_attachment_adds_audio_part() {
|
||||
let channel = EmailChannel::new(EmailConfig {
|
||||
from_address: "bot@example.com".to_string(),
|
||||
..Default::default()
|
||||
});
|
||||
let message = channel
|
||||
.build_message_with_attachment(
|
||||
"listener@example.com",
|
||||
"Weekly briefing",
|
||||
"Attached.",
|
||||
"briefing.mp3",
|
||||
"audio/mpeg".parse().expect("content type"),
|
||||
vec![1, 2, 3, 4],
|
||||
)
|
||||
.expect("attachment message");
|
||||
let wire = String::from_utf8_lossy(&message.formatted()).to_string();
|
||||
assert!(wire.contains("Subject: Weekly briefing"));
|
||||
assert!(wire.contains("filename=\"briefing.mp3\""));
|
||||
assert!(wire.contains("Content-Type: audio/mpeg"));
|
||||
}
|
||||
|
||||
// is_sender_allowed tests
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_empty_list_denies_all() {
|
||||
let config = EmailConfig {
|
||||
allowed_senders: vec![],
|
||||
..Default::default()
|
||||
};
|
||||
let channel = EmailChannel::new(config);
|
||||
assert!(!channel.is_sender_allowed("anyone@example.com"));
|
||||
assert!(!channel.is_sender_allowed("user@test.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_wildcard_allows_all() {
|
||||
let config = EmailConfig {
|
||||
allowed_senders: vec!["*".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let channel = EmailChannel::new(config);
|
||||
assert!(channel.is_sender_allowed("anyone@example.com"));
|
||||
assert!(channel.is_sender_allowed("user@test.com"));
|
||||
assert!(channel.is_sender_allowed("random@domain.org"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_specific_email() {
|
||||
let config = EmailConfig {
|
||||
allowed_senders: vec!["allowed@example.com".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let channel = EmailChannel::new(config);
|
||||
assert!(channel.is_sender_allowed("allowed@example.com"));
|
||||
assert!(!channel.is_sender_allowed("other@example.com"));
|
||||
assert!(!channel.is_sender_allowed("allowed@other.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_domain_with_at_prefix() {
|
||||
let config = EmailConfig {
|
||||
allowed_senders: vec!["@example.com".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let channel = EmailChannel::new(config);
|
||||
assert!(channel.is_sender_allowed("user@example.com"));
|
||||
assert!(channel.is_sender_allowed("admin@example.com"));
|
||||
assert!(!channel.is_sender_allowed("user@other.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_domain_without_at_prefix() {
|
||||
let config = EmailConfig {
|
||||
allowed_senders: vec!["example.com".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let channel = EmailChannel::new(config);
|
||||
assert!(channel.is_sender_allowed("user@example.com"));
|
||||
assert!(channel.is_sender_allowed("admin@example.com"));
|
||||
assert!(!channel.is_sender_allowed("user@other.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_case_insensitive() {
|
||||
let config = EmailConfig {
|
||||
allowed_senders: vec!["Allowed@Example.COM".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let channel = EmailChannel::new(config);
|
||||
assert!(channel.is_sender_allowed("allowed@example.com"));
|
||||
assert!(channel.is_sender_allowed("ALLOWED@EXAMPLE.COM"));
|
||||
assert!(channel.is_sender_allowed("AlLoWeD@eXaMpLe.cOm"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_multiple_senders() {
|
||||
let config = EmailConfig {
|
||||
allowed_senders: vec![
|
||||
"user1@example.com".to_string(),
|
||||
"user2@test.com".to_string(),
|
||||
"@allowed.com".to_string(),
|
||||
],
|
||||
..Default::default()
|
||||
};
|
||||
let channel = EmailChannel::new(config);
|
||||
assert!(channel.is_sender_allowed("user1@example.com"));
|
||||
assert!(channel.is_sender_allowed("user2@test.com"));
|
||||
assert!(channel.is_sender_allowed("anyone@allowed.com"));
|
||||
assert!(!channel.is_sender_allowed("user3@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_wildcard_with_specific() {
|
||||
let config = EmailConfig {
|
||||
allowed_senders: vec!["*".to_string(), "specific@example.com".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let channel = EmailChannel::new(config);
|
||||
assert!(channel.is_sender_allowed("anyone@example.com"));
|
||||
assert!(channel.is_sender_allowed("specific@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_empty_sender() {
|
||||
let config = EmailConfig {
|
||||
allowed_senders: vec!["@example.com".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let channel = EmailChannel::new(config);
|
||||
assert!(!channel.is_sender_allowed(""));
|
||||
// "@example.com" ends with "@example.com" so it's allowed
|
||||
assert!(channel.is_sender_allowed("@example.com"));
|
||||
}
|
||||
|
||||
// strip_html tests
|
||||
|
||||
#[test]
|
||||
fn strip_html_basic() {
|
||||
assert_eq!(EmailChannel::strip_html("<p>Hello</p>"), "Hello");
|
||||
assert_eq!(EmailChannel::strip_html("<div>World</div>"), "World");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_nested_tags() {
|
||||
assert_eq!(
|
||||
EmailChannel::strip_html("<div><p>Hello <strong>World</strong></p></div>"),
|
||||
"Hello World"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_multiple_lines() {
|
||||
let html = "<div>\n <p>Line 1</p>\n <p>Line 2</p>\n</div>";
|
||||
assert_eq!(EmailChannel::strip_html(html), "Line 1 Line 2");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_preserves_text() {
|
||||
assert_eq!(EmailChannel::strip_html("No tags here"), "No tags here");
|
||||
assert_eq!(EmailChannel::strip_html(""), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_handles_malformed() {
|
||||
assert_eq!(EmailChannel::strip_html("<p>Unclosed"), "Unclosed");
|
||||
// The function removes everything between < and >, so "Text>with>brackets" becomes "Textwithbrackets"
|
||||
assert_eq!(
|
||||
EmailChannel::strip_html("Text>with>brackets"),
|
||||
"Textwithbrackets"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_self_closing_tags() {
|
||||
// Self-closing tags are removed but don't add spaces
|
||||
assert_eq!(EmailChannel::strip_html("Hello<br/>World"), "HelloWorld");
|
||||
assert_eq!(EmailChannel::strip_html("Text<hr/>More"), "TextMore");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_attributes_preserved() {
|
||||
assert_eq!(
|
||||
EmailChannel::strip_html("<a href=\"http://example.com\">Link</a>"),
|
||||
"Link"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_multiple_spaces_collapsed() {
|
||||
assert_eq!(
|
||||
EmailChannel::strip_html("<p>Word</p> <p>Word</p>"),
|
||||
"Word Word"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_special_characters() {
|
||||
assert_eq!(
|
||||
EmailChannel::strip_html("<span><tag></span>"),
|
||||
"<tag>"
|
||||
);
|
||||
}
|
||||
|
||||
// Default function tests
|
||||
|
||||
#[test]
|
||||
fn default_imap_port_returns_993() {
|
||||
assert_eq!(EmailConfig::default().imap_port, 993);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_smtp_port_returns_465() {
|
||||
assert_eq!(EmailConfig::default().smtp_port, 465);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_imap_folder_returns_inbox() {
|
||||
assert_eq!(EmailConfig::default().imap_folder, "INBOX");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_true_returns_true() {
|
||||
assert!(EmailConfig::default().smtp_tls);
|
||||
}
|
||||
|
||||
// EmailConfig serialization tests
|
||||
|
||||
#[test]
|
||||
fn email_config_serialize_deserialize() {
|
||||
let config = EmailConfig {
|
||||
imap_host: "imap.example.com".to_string(),
|
||||
imap_port: 993,
|
||||
imap_folder: "INBOX".to_string(),
|
||||
smtp_host: "smtp.example.com".to_string(),
|
||||
smtp_port: 587,
|
||||
smtp_tls: true,
|
||||
username: "user@example.com".to_string(),
|
||||
password: "password123".to_string(),
|
||||
from_address: "bot@example.com".to_string(),
|
||||
idle_timeout_secs: 1740,
|
||||
allowed_senders: vec!["allowed@example.com".to_string()],
|
||||
};
|
||||
|
||||
let json = serde_json::to_string(&config).unwrap();
|
||||
let deserialized: EmailConfig = serde_json::from_str(&json).unwrap();
|
||||
|
||||
assert_eq!(deserialized.imap_host, config.imap_host);
|
||||
assert_eq!(deserialized.smtp_port, config.smtp_port);
|
||||
assert_eq!(deserialized.allowed_senders, config.allowed_senders);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_config_deserialize_with_defaults() {
|
||||
let json = r#"{
|
||||
"imap_host": "imap.test.com",
|
||||
"smtp_host": "smtp.test.com",
|
||||
"username": "user",
|
||||
"password": "pass",
|
||||
"from_address": "bot@test.com"
|
||||
}"#;
|
||||
|
||||
let config: EmailConfig = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(config.imap_port, 993); // default
|
||||
assert_eq!(config.smtp_port, 465); // default
|
||||
assert!(config.smtp_tls); // default
|
||||
assert_eq!(config.idle_timeout_secs, 1740); // default
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idle_timeout_deserializes_explicit_value() {
|
||||
let json = r#"{
|
||||
"imap_host": "imap.test.com",
|
||||
"smtp_host": "smtp.test.com",
|
||||
"username": "user",
|
||||
"password": "pass",
|
||||
"from_address": "bot@test.com",
|
||||
"idle_timeout_secs": 900
|
||||
}"#;
|
||||
let config: EmailConfig = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(config.idle_timeout_secs, 900);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idle_timeout_deserializes_legacy_poll_interval_alias() {
|
||||
let json = r#"{
|
||||
"imap_host": "imap.test.com",
|
||||
"smtp_host": "smtp.test.com",
|
||||
"username": "user",
|
||||
"password": "pass",
|
||||
"from_address": "bot@test.com",
|
||||
"poll_interval_secs": 120
|
||||
}"#;
|
||||
let config: EmailConfig = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(config.idle_timeout_secs, 120);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn idle_timeout_propagates_to_channel() {
|
||||
let config = EmailConfig {
|
||||
idle_timeout_secs: 600,
|
||||
..Default::default()
|
||||
};
|
||||
let channel = EmailChannel::new(config);
|
||||
assert_eq!(channel.config.idle_timeout_secs, 600);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn email_config_debug_output() {
|
||||
let config = EmailConfig {
|
||||
imap_host: "imap.debug.com".to_string(),
|
||||
..Default::default()
|
||||
};
|
||||
let debug_str = format!("{:?}", config);
|
||||
assert!(debug_str.contains("imap.debug.com"));
|
||||
}
|
||||
|
||||
// ── is_sender_allowed comprehensive matrix ─────────────────────
|
||||
|
||||
fn channel_with_allowlist(allowlist: Vec<String>) -> EmailChannel {
|
||||
let cfg = EmailConfig {
|
||||
imap_host: "imap.x".into(),
|
||||
imap_port: 993,
|
||||
imap_folder: "INBOX".into(),
|
||||
smtp_host: "smtp.x".into(),
|
||||
smtp_port: 465,
|
||||
smtp_tls: true,
|
||||
username: "u".into(),
|
||||
password: "p".into(),
|
||||
from_address: "me@x".into(),
|
||||
idle_timeout_secs: 300,
|
||||
allowed_senders: allowlist,
|
||||
};
|
||||
EmailChannel::new(cfg)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_empty_denies_all() {
|
||||
let ch = channel_with_allowlist(vec![]);
|
||||
assert!(!ch.is_sender_allowed("anyone@any.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_wildcard_allows_everyone() {
|
||||
let ch = channel_with_allowlist(vec!["*".into()]);
|
||||
assert!(ch.is_sender_allowed("anyone@any.com"));
|
||||
assert!(ch.is_sender_allowed("other@different.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_full_email_exact_match_case_insensitive() {
|
||||
let ch = channel_with_allowlist(vec!["alice@example.com".into()]);
|
||||
assert!(ch.is_sender_allowed("alice@example.com"));
|
||||
assert!(ch.is_sender_allowed("ALICE@EXAMPLE.COM"));
|
||||
assert!(!ch.is_sender_allowed("bob@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_at_prefix_domain_match() {
|
||||
let ch = channel_with_allowlist(vec!["@trusted.com".into()]);
|
||||
assert!(ch.is_sender_allowed("user@trusted.com"));
|
||||
assert!(ch.is_sender_allowed("other@Trusted.com"));
|
||||
assert!(!ch.is_sender_allowed("user@untrusted.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_bare_domain_match_is_case_insensitive() {
|
||||
let ch = channel_with_allowlist(vec!["trusted.com".into()]);
|
||||
assert!(ch.is_sender_allowed("user@trusted.com"));
|
||||
assert!(ch.is_sender_allowed("USER@TRUSTED.COM"));
|
||||
assert!(!ch.is_sender_allowed("user@other.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_sender_allowed_prevents_subdomain_confusion() {
|
||||
// "trusted.com" must NOT match "user@malicioustrusted.com"
|
||||
let ch = channel_with_allowlist(vec!["trusted.com".into()]);
|
||||
assert!(!ch.is_sender_allowed("user@notmytrusted.com"));
|
||||
assert!(!ch.is_sender_allowed("user@trusted.com.evil.com"));
|
||||
}
|
||||
|
||||
// ── strip_html edge cases ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn strip_html_empty_string() {
|
||||
assert_eq!(EmailChannel::strip_html(""), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_only_tags() {
|
||||
assert_eq!(EmailChannel::strip_html("<p></p><br/>"), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_unclosed_tag_eats_rest_until_gt() {
|
||||
// A '<' without '>' enters tag mode; anything after until a '>' is
|
||||
// discarded. This is the implementation's behaviour — lock it in.
|
||||
assert_eq!(EmailChannel::strip_html("before<never closed"), "before");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_html_collapses_whitespace_runs() {
|
||||
assert_eq!(
|
||||
EmailChannel::strip_html("<p>hello</p>\n\n\n <p>world</p>"),
|
||||
"hello world"
|
||||
);
|
||||
}
|
||||
@@ -1,318 +1 @@
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
use async_trait::async_trait;
|
||||
use directories::UserDirs;
|
||||
use rusqlite::{Connection, OpenFlags};
|
||||
use std::path::Path;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// iMessage channel using macOS `AppleScript` bridge.
|
||||
/// Polls the Messages database for new messages and sends replies via `osascript`.
|
||||
#[derive(Clone)]
|
||||
pub struct IMessageChannel {
|
||||
allowed_contacts: Vec<String>,
|
||||
poll_interval_secs: u64,
|
||||
}
|
||||
|
||||
impl IMessageChannel {
|
||||
pub fn new(allowed_contacts: Vec<String>) -> Self {
|
||||
Self {
|
||||
allowed_contacts,
|
||||
poll_interval_secs: 3,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_contact_allowed(&self, sender: &str) -> bool {
|
||||
if self.allowed_contacts.iter().any(|u| u == "*") {
|
||||
return true;
|
||||
}
|
||||
self.allowed_contacts
|
||||
.iter()
|
||||
.any(|u| u.eq_ignore_ascii_case(sender))
|
||||
}
|
||||
}
|
||||
|
||||
/// Escape a string for safe interpolation into `AppleScript`.
|
||||
///
|
||||
/// This prevents injection attacks by escaping:
|
||||
/// - Backslashes (`\` → `\\`)
|
||||
/// - Double quotes (`"` → `\"`)
|
||||
/// - Newlines (`\n` → `\\n`, `\r` → `\\r`) to prevent code injection via line breaks
|
||||
fn escape_applescript(s: &str) -> String {
|
||||
s.replace('\\', "\\\\")
|
||||
.replace('"', "\\\"")
|
||||
.replace('\n', "\\n")
|
||||
.replace('\r', "\\r")
|
||||
}
|
||||
|
||||
/// Validate that a target looks like a valid phone number or email address.
|
||||
///
|
||||
/// This is a defense-in-depth measure to reject obviously malicious targets
|
||||
/// before they reach `AppleScript` interpolation.
|
||||
///
|
||||
/// Valid patterns:
|
||||
/// - Phone: starts with `+` followed by digits (with optional spaces/dashes)
|
||||
/// - Email: contains `@` with alphanumeric chars on both sides
|
||||
fn is_valid_imessage_target(target: &str) -> bool {
|
||||
let target = target.trim();
|
||||
if target.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
// Phone number: +1234567890 or +1 234-567-8900
|
||||
if target.starts_with('+') {
|
||||
let digits_only: String = target.chars().filter(char::is_ascii_digit).collect();
|
||||
// Must have at least 7 digits (shortest valid phone numbers)
|
||||
return digits_only.len() >= 7 && digits_only.len() <= 15;
|
||||
}
|
||||
|
||||
// Email: simple validation (contains @ with chars on both sides)
|
||||
if let Some(at_pos) = target.find('@') {
|
||||
let local = &target[..at_pos];
|
||||
let domain = &target[at_pos + 1..];
|
||||
|
||||
// Local part: non-empty, alphanumeric + common email chars
|
||||
let local_valid = !local.is_empty()
|
||||
&& local
|
||||
.chars()
|
||||
.all(|c| c.is_alphanumeric() || "._+-".contains(c));
|
||||
|
||||
// Domain: non-empty, contains a dot, alphanumeric + dots/hyphens
|
||||
let domain_valid = !domain.is_empty()
|
||||
&& domain.contains('.')
|
||||
&& domain
|
||||
.chars()
|
||||
.all(|c| c.is_alphanumeric() || ".-".contains(c));
|
||||
|
||||
return local_valid && domain_valid;
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for IMessageChannel {
|
||||
fn name(&self) -> &str {
|
||||
"imessage"
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> anyhow::Result<()> {
|
||||
// Defense-in-depth: validate target format before any interpolation
|
||||
if !is_valid_imessage_target(&message.recipient) {
|
||||
anyhow::bail!(
|
||||
"Invalid iMessage target: must be a phone number (+1234567890) or email (user@example.com)"
|
||||
);
|
||||
}
|
||||
|
||||
// SECURITY: Escape both message AND target to prevent AppleScript injection
|
||||
// See: CWE-78 (OS Command Injection)
|
||||
let escaped_msg = escape_applescript(&message.content);
|
||||
let escaped_target = escape_applescript(&message.recipient);
|
||||
|
||||
let script = format!(
|
||||
r#"tell application "Messages"
|
||||
set targetService to 1st account whose service type = iMessage
|
||||
set targetBuddy to participant "{escaped_target}" of targetService
|
||||
send "{escaped_msg}" to targetBuddy
|
||||
end tell"#
|
||||
);
|
||||
|
||||
let output = tokio::process::Command::new("osascript")
|
||||
.arg("-e")
|
||||
.arg(&script)
|
||||
.output()
|
||||
.await?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
anyhow::bail!("iMessage send failed: {stderr}");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn listen(&self, tx: mpsc::Sender<ChannelMessage>) -> anyhow::Result<()> {
|
||||
tracing::info!("iMessage channel listening (AppleScript bridge)...");
|
||||
|
||||
// Query the Messages SQLite database for new messages
|
||||
// The database is at ~/Library/Messages/chat.db
|
||||
let db_path = UserDirs::new()
|
||||
.map(|u| u.home_dir().join("Library/Messages/chat.db"))
|
||||
.ok_or_else(|| anyhow::anyhow!("Cannot find home directory"))?;
|
||||
|
||||
if !db_path.exists() {
|
||||
anyhow::bail!(
|
||||
"Messages database not found at {}. Ensure Messages.app is set up and Full Disk Access is granted.",
|
||||
db_path.display()
|
||||
);
|
||||
}
|
||||
|
||||
// Open a persistent read-only connection instead of creating
|
||||
// a new one on every 3-second poll cycle.
|
||||
let path = db_path.to_path_buf();
|
||||
let conn = tokio::task::spawn_blocking(move || -> anyhow::Result<Connection> {
|
||||
Ok(Connection::open_with_flags(
|
||||
&path,
|
||||
OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
|
||||
)?)
|
||||
})
|
||||
.await??;
|
||||
|
||||
// Track the last ROWID we've seen (shuttle conn in and out)
|
||||
let (mut conn, initial_rowid) =
|
||||
tokio::task::spawn_blocking(move || -> anyhow::Result<(Connection, i64)> {
|
||||
let rowid = {
|
||||
let mut stmt =
|
||||
conn.prepare("SELECT MAX(ROWID) FROM message WHERE is_from_me = 0")?;
|
||||
let rowid: Option<i64> = stmt.query_row([], |row| row.get(0))?;
|
||||
rowid.unwrap_or(0)
|
||||
};
|
||||
Ok((conn, rowid))
|
||||
})
|
||||
.await??;
|
||||
let mut last_rowid = initial_rowid;
|
||||
|
||||
loop {
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(self.poll_interval_secs)).await;
|
||||
|
||||
let since = last_rowid;
|
||||
let (returned_conn, poll_result) = tokio::task::spawn_blocking(
|
||||
move || -> (Connection, anyhow::Result<Vec<(i64, String, String)>>) {
|
||||
let result = (|| -> anyhow::Result<Vec<(i64, String, String)>> {
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT m.ROWID, h.id, m.text \
|
||||
FROM message m \
|
||||
JOIN handle h ON m.handle_id = h.ROWID \
|
||||
WHERE m.ROWID > ?1 \
|
||||
AND m.is_from_me = 0 \
|
||||
AND m.text IS NOT NULL \
|
||||
ORDER BY m.ROWID ASC \
|
||||
LIMIT 20",
|
||||
)?;
|
||||
let rows = stmt.query_map([since], |row| {
|
||||
Ok((
|
||||
row.get::<_, i64>(0)?,
|
||||
row.get::<_, String>(1)?,
|
||||
row.get::<_, String>(2)?,
|
||||
))
|
||||
})?;
|
||||
let results = rows.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(results)
|
||||
})();
|
||||
|
||||
(conn, result)
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("iMessage poll worker join error: {e}"))?;
|
||||
conn = returned_conn;
|
||||
|
||||
match poll_result {
|
||||
Ok(messages) => {
|
||||
for (rowid, sender, text) in messages {
|
||||
if rowid > last_rowid {
|
||||
last_rowid = rowid;
|
||||
}
|
||||
|
||||
if !self.is_contact_allowed(&sender) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if text.trim().is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let msg = ChannelMessage {
|
||||
id: rowid.to_string(),
|
||||
sender: sender.clone(),
|
||||
reply_target: sender.clone(),
|
||||
content: text,
|
||||
channel: "imessage".to_string(),
|
||||
timestamp: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
thread_ts: None,
|
||||
};
|
||||
|
||||
if tx.send(msg).await.is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("iMessage poll error: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
if std::env::consts::OS != "macos" {
|
||||
return false;
|
||||
}
|
||||
|
||||
let db_path = UserDirs::new()
|
||||
.map(|u| u.home_dir().join("Library/Messages/chat.db"))
|
||||
.unwrap_or_default();
|
||||
|
||||
db_path.exists()
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the current max ROWID from the messages table.
|
||||
/// Uses rusqlite with parameterized queries for security (CWE-89 prevention).
|
||||
async fn get_max_rowid(db_path: &Path) -> anyhow::Result<i64> {
|
||||
let path = db_path.to_path_buf();
|
||||
let result = tokio::task::spawn_blocking(move || -> anyhow::Result<i64> {
|
||||
let conn = Connection::open_with_flags(
|
||||
&path,
|
||||
OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
|
||||
)?;
|
||||
let mut stmt = conn.prepare("SELECT MAX(ROWID) FROM message WHERE is_from_me = 0")?;
|
||||
let rowid: Option<i64> = stmt.query_row([], |row| row.get(0))?;
|
||||
Ok(rowid.unwrap_or(0))
|
||||
})
|
||||
.await??;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Fetch messages newer than `since_rowid`.
|
||||
/// Uses rusqlite with parameterized queries for security (CWE-89 prevention).
|
||||
/// The `since_rowid` parameter is bound safely, preventing SQL injection.
|
||||
async fn fetch_new_messages(
|
||||
db_path: &Path,
|
||||
since_rowid: i64,
|
||||
) -> anyhow::Result<Vec<(i64, String, String)>> {
|
||||
let path = db_path.to_path_buf();
|
||||
let results =
|
||||
tokio::task::spawn_blocking(move || -> anyhow::Result<Vec<(i64, String, String)>> {
|
||||
let conn = Connection::open_with_flags(
|
||||
&path,
|
||||
OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
|
||||
)?;
|
||||
let mut stmt = conn.prepare(
|
||||
"SELECT m.ROWID, h.id, m.text \
|
||||
FROM message m \
|
||||
JOIN handle h ON m.handle_id = h.ROWID \
|
||||
WHERE m.ROWID > ?1 \
|
||||
AND m.is_from_me = 0 \
|
||||
AND m.text IS NOT NULL \
|
||||
ORDER BY m.ROWID ASC \
|
||||
LIMIT 20",
|
||||
)?;
|
||||
let rows = stmt.query_map([since_rowid], |row| {
|
||||
Ok((
|
||||
row.get::<_, i64>(0)?,
|
||||
row.get::<_, String>(1)?,
|
||||
row.get::<_, String>(2)?,
|
||||
))
|
||||
})?;
|
||||
rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
|
||||
})
|
||||
.await??;
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "imessage_tests.rs"]
|
||||
mod tests;
|
||||
pub use tinychannels::providers::imessage::*;
|
||||
|
||||
@@ -1,672 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn creates_with_contacts() {
|
||||
let ch = IMessageChannel::new(vec!["+1234567890".into()]);
|
||||
assert_eq!(ch.allowed_contacts.len(), 1);
|
||||
assert_eq!(ch.poll_interval_secs, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn creates_with_empty_contacts() {
|
||||
let ch = IMessageChannel::new(vec![]);
|
||||
assert!(ch.allowed_contacts.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wildcard_allows_anyone() {
|
||||
let ch = IMessageChannel::new(vec!["*".into()]);
|
||||
assert!(ch.is_contact_allowed("+1234567890"));
|
||||
assert!(ch.is_contact_allowed("random@icloud.com"));
|
||||
assert!(ch.is_contact_allowed(""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn specific_contact_allowed() {
|
||||
let ch = IMessageChannel::new(vec!["+1234567890".into(), "user@icloud.com".into()]);
|
||||
assert!(ch.is_contact_allowed("+1234567890"));
|
||||
assert!(ch.is_contact_allowed("user@icloud.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_contact_denied() {
|
||||
let ch = IMessageChannel::new(vec!["+1234567890".into()]);
|
||||
assert!(!ch.is_contact_allowed("+9999999999"));
|
||||
assert!(!ch.is_contact_allowed("hacker@evil.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn contact_case_insensitive() {
|
||||
let ch = IMessageChannel::new(vec!["User@iCloud.com".into()]);
|
||||
assert!(ch.is_contact_allowed("user@icloud.com"));
|
||||
assert!(ch.is_contact_allowed("USER@ICLOUD.COM"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_allowlist_denies_all() {
|
||||
let ch = IMessageChannel::new(vec![]);
|
||||
assert!(!ch.is_contact_allowed("+1234567890"));
|
||||
assert!(!ch.is_contact_allowed("anyone"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn name_returns_imessage() {
|
||||
let ch = IMessageChannel::new(vec![]);
|
||||
assert_eq!(ch.name(), "imessage");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wildcard_among_others_still_allows_all() {
|
||||
let ch = IMessageChannel::new(vec!["+111".into(), "*".into(), "+222".into()]);
|
||||
assert!(ch.is_contact_allowed("totally-unknown"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn contact_with_spaces_exact_match() {
|
||||
let ch = IMessageChannel::new(vec![" spaced ".into()]);
|
||||
assert!(ch.is_contact_allowed(" spaced "));
|
||||
assert!(!ch.is_contact_allowed("spaced"));
|
||||
}
|
||||
|
||||
// ══════════════════════════════════════════════════════════
|
||||
// AppleScript Escaping Tests (CWE-78 Prevention)
|
||||
// ══════════════════════════════════════════════════════════
|
||||
|
||||
#[test]
|
||||
fn escape_applescript_double_quotes() {
|
||||
assert_eq!(escape_applescript(r#"hello "world""#), r#"hello \"world\""#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_applescript_backslashes() {
|
||||
assert_eq!(escape_applescript(r"path\to\file"), r"path\\to\\file");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_applescript_mixed() {
|
||||
assert_eq!(
|
||||
escape_applescript(r#"say "hello\" world"#),
|
||||
r#"say \"hello\\\" world"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_applescript_injection_attempt() {
|
||||
// This is the exact attack vector from the security report
|
||||
let malicious = r#"" & do shell script "id" & ""#;
|
||||
let escaped = escape_applescript(malicious);
|
||||
// After escaping, the quotes should be escaped and not break out
|
||||
assert_eq!(escaped, r#"\" & do shell script \"id\" & \""#);
|
||||
// Verify all quotes are now escaped (preceded by backslash)
|
||||
// The escaped string should not have any unescaped quotes (quote not preceded by backslash)
|
||||
let chars: Vec<char> = escaped.chars().collect();
|
||||
for (i, &c) in chars.iter().enumerate() {
|
||||
if c == '"' {
|
||||
// Every quote must be preceded by a backslash
|
||||
assert!(
|
||||
i > 0 && chars[i - 1] == '\\',
|
||||
"Found unescaped quote at position {i}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_applescript_empty_string() {
|
||||
assert_eq!(escape_applescript(""), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_applescript_no_special_chars() {
|
||||
assert_eq!(escape_applescript("hello world"), "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_applescript_unicode() {
|
||||
assert_eq!(escape_applescript("hello 🦀 world"), "hello 🦀 world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_applescript_newlines_escaped() {
|
||||
assert_eq!(escape_applescript("line1\nline2"), "line1\\nline2");
|
||||
assert_eq!(escape_applescript("line1\rline2"), "line1\\rline2");
|
||||
assert_eq!(escape_applescript("line1\r\nline2"), "line1\\r\\nline2");
|
||||
}
|
||||
|
||||
// ══════════════════════════════════════════════════════════
|
||||
// Target Validation Tests
|
||||
// ══════════════════════════════════════════════════════════
|
||||
|
||||
#[test]
|
||||
fn valid_phone_number_simple() {
|
||||
assert!(is_valid_imessage_target("+1234567890"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_phone_number_with_country_code() {
|
||||
assert!(is_valid_imessage_target("+14155551234"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_phone_number_with_spaces() {
|
||||
assert!(is_valid_imessage_target("+1 415 555 1234"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_phone_number_with_dashes() {
|
||||
assert!(is_valid_imessage_target("+1-415-555-1234"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_phone_number_international() {
|
||||
assert!(is_valid_imessage_target("+447911123456")); // UK
|
||||
assert!(is_valid_imessage_target("+81312345678")); // Japan
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_email_simple() {
|
||||
assert!(is_valid_imessage_target("user@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_email_with_subdomain() {
|
||||
assert!(is_valid_imessage_target("user@mail.example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_email_with_plus() {
|
||||
assert!(is_valid_imessage_target("user+tag@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_email_with_dots() {
|
||||
assert!(is_valid_imessage_target("first.last@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_email_icloud() {
|
||||
assert!(is_valid_imessage_target("user@icloud.com"));
|
||||
assert!(is_valid_imessage_target("user@me.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_empty() {
|
||||
assert!(!is_valid_imessage_target(""));
|
||||
assert!(!is_valid_imessage_target(" "));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_no_plus_prefix() {
|
||||
// Phone numbers must start with +
|
||||
assert!(!is_valid_imessage_target("1234567890"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_too_short_phone() {
|
||||
// Less than 7 digits
|
||||
assert!(!is_valid_imessage_target("+123456"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_too_long_phone() {
|
||||
// More than 15 digits
|
||||
assert!(!is_valid_imessage_target("+1234567890123456"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_email_no_at() {
|
||||
assert!(!is_valid_imessage_target("userexample.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_email_no_domain() {
|
||||
assert!(!is_valid_imessage_target("user@"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_email_no_local() {
|
||||
assert!(!is_valid_imessage_target("@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_email_no_dot_in_domain() {
|
||||
assert!(!is_valid_imessage_target("user@localhost"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_injection_attempt() {
|
||||
// The exact attack vector from the security report
|
||||
assert!(!is_valid_imessage_target(r#"" & do shell script "id" & ""#));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_applescript_injection() {
|
||||
// Various injection attempts
|
||||
assert!(!is_valid_imessage_target(r#"test" & quit"#));
|
||||
assert!(!is_valid_imessage_target(r"test\ndo shell script"));
|
||||
assert!(!is_valid_imessage_target("test\"; malicious code; \""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_special_chars() {
|
||||
assert!(!is_valid_imessage_target("user<script>@example.com"));
|
||||
assert!(!is_valid_imessage_target("user@example.com; rm -rf /"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_null_byte() {
|
||||
assert!(!is_valid_imessage_target("user\0@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_target_newline() {
|
||||
assert!(!is_valid_imessage_target("user\n@example.com"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn target_with_leading_trailing_whitespace_trimmed() {
|
||||
// Should trim and validate
|
||||
assert!(is_valid_imessage_target(" +1234567890 "));
|
||||
assert!(is_valid_imessage_target(" user@example.com "));
|
||||
}
|
||||
|
||||
// ══════════════════════════════════════════════════════════
|
||||
// SQLite/rusqlite Database Tests (CWE-89 Prevention)
|
||||
// ══════════════════════════════════════════════════════════
|
||||
|
||||
/// Helper to create a temporary test database with Messages schema
|
||||
fn create_test_db() -> (tempfile::TempDir, std::path::PathBuf) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db_path = dir.path().join("chat.db");
|
||||
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
|
||||
// Create minimal schema matching macOS Messages.app
|
||||
conn.execute_batch(
|
||||
"CREATE TABLE handle (
|
||||
ROWID INTEGER PRIMARY KEY,
|
||||
id TEXT NOT NULL
|
||||
);
|
||||
CREATE TABLE message (
|
||||
ROWID INTEGER PRIMARY KEY,
|
||||
handle_id INTEGER,
|
||||
text TEXT,
|
||||
is_from_me INTEGER DEFAULT 0,
|
||||
FOREIGN KEY (handle_id) REFERENCES handle(ROWID)
|
||||
);",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
(dir, db_path)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_max_rowid_empty_database() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
let result = get_max_rowid(&db_path).await;
|
||||
assert!(result.is_ok());
|
||||
// Empty table returns 0 (NULL coalesced)
|
||||
assert_eq!(result.unwrap(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_max_rowid_with_messages() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
// Insert test data
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (100, 1, 'Hello', 0)",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (200, 1, 'World', 0)",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
// This one is from_me=1, should be ignored
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (300, 1, 'Sent', 1)",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let result = get_max_rowid(&db_path).await.unwrap();
|
||||
// Should return 200, not 300 (ignores is_from_me=1)
|
||||
assert_eq!(result, 200);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_max_rowid_nonexistent_database() {
|
||||
let path = std::path::Path::new("/nonexistent/path/chat.db");
|
||||
let result = get_max_rowid(path).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_empty_database() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
let result = fetch_new_messages(&db_path, 0).await;
|
||||
assert!(result.is_ok());
|
||||
assert!(result.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_returns_correct_data() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
// Insert test data
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (2, 'user@example.com')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (10, 1, 'First message', 0)",
|
||||
[]
|
||||
).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (20, 2, 'Second message', 0)",
|
||||
[]
|
||||
).unwrap();
|
||||
}
|
||||
|
||||
let result = fetch_new_messages(&db_path, 0).await.unwrap();
|
||||
assert_eq!(result.len(), 2);
|
||||
assert_eq!(
|
||||
result[0],
|
||||
(10, "+1234567890".to_string(), "First message".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
result[1],
|
||||
(
|
||||
20,
|
||||
"user@example.com".to_string(),
|
||||
"Second message".to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_filters_by_rowid() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
// Insert test data
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (10, 1, 'Old message', 0)",
|
||||
[]
|
||||
).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (20, 1, 'New message', 0)",
|
||||
[]
|
||||
).unwrap();
|
||||
}
|
||||
|
||||
// Fetch only messages after ROWID 15
|
||||
let result = fetch_new_messages(&db_path, 15).await.unwrap();
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].0, 20);
|
||||
assert_eq!(result[0].2, "New message");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_excludes_sent_messages() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
// Insert test data
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (10, 1, 'Received', 0)",
|
||||
[]
|
||||
).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (20, 1, 'Sent by me', 1)",
|
||||
[]
|
||||
).unwrap();
|
||||
}
|
||||
|
||||
let result = fetch_new_messages(&db_path, 0).await.unwrap();
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].2, "Received");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_excludes_null_text() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
// Insert test data
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (10, 1, 'Has text', 0)",
|
||||
[]
|
||||
).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (20, 1, NULL, 0)",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let result = fetch_new_messages(&db_path, 0).await.unwrap();
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].2, "Has text");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_respects_limit() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
// Insert 25 messages (limit is 20)
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
for i in 1..=25 {
|
||||
conn.execute(
|
||||
&format!("INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES ({i}, 1, 'Message {i}', 0)"),
|
||||
[]
|
||||
).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
let result = fetch_new_messages(&db_path, 0).await.unwrap();
|
||||
assert_eq!(result.len(), 20); // Limited to 20
|
||||
assert_eq!(result[0].0, 1); // First message
|
||||
assert_eq!(result[19].0, 20); // 20th message
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_ordered_by_rowid_asc() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
// Insert messages out of order
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (30, 1, 'Third', 0)",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (10, 1, 'First', 0)",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (20, 1, 'Second', 0)",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let result = fetch_new_messages(&db_path, 0).await.unwrap();
|
||||
assert_eq!(result.len(), 3);
|
||||
assert_eq!(result[0].0, 10);
|
||||
assert_eq!(result[1].0, 20);
|
||||
assert_eq!(result[2].0, 30);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_nonexistent_database() {
|
||||
let path = std::path::Path::new("/nonexistent/path/chat.db");
|
||||
let result = fetch_new_messages(path, 0).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_handles_special_characters() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
// Insert message with special characters (potential SQL injection patterns)
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (10, 1, 'Hello \"world'' OR 1=1; DROP TABLE message;--', 0)",
|
||||
[]
|
||||
).unwrap();
|
||||
}
|
||||
|
||||
let result = fetch_new_messages(&db_path, 0).await.unwrap();
|
||||
assert_eq!(result.len(), 1);
|
||||
// The special characters should be preserved, not interpreted as SQL
|
||||
assert!(result[0].2.contains("DROP TABLE"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_handles_unicode() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (10, 1, 'Hello 🦀 世界 مرحبا', 0)",
|
||||
[]
|
||||
).unwrap();
|
||||
}
|
||||
|
||||
let result = fetch_new_messages(&db_path, 0).await.unwrap();
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].2, "Hello 🦀 世界 مرحبا");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_handles_empty_text() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (10, 1, '', 0)",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let result = fetch_new_messages(&db_path, 0).await.unwrap();
|
||||
// Empty string is NOT NULL, so it's included
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].2, "");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_negative_rowid_edge_case() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (10, 1, 'Test', 0)",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Negative rowid should still work (fetch all messages with ROWID > -1)
|
||||
let result = fetch_new_messages(&db_path, -1).await.unwrap();
|
||||
assert_eq!(result.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fetch_new_messages_large_rowid_edge_case() {
|
||||
let (_dir, db_path) = create_test_db();
|
||||
|
||||
{
|
||||
let conn = Connection::open(&db_path).unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO handle (ROWID, id) VALUES (1, '+1234567890')",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO message (ROWID, handle_id, text, is_from_me) VALUES (10, 1, 'Test', 0)",
|
||||
[],
|
||||
)
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Very large rowid should return empty (no messages after this)
|
||||
let result = fetch_new_messages(&db_path, i64::MAX - 1).await.unwrap();
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
@@ -1,656 +1 @@
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
use async_trait::async_trait;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::Arc;
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::sync::{mpsc, Mutex};
|
||||
|
||||
// Use tokio_rustls's re-export of rustls types
|
||||
use tokio_rustls::rustls;
|
||||
|
||||
/// Read timeout for IRC — if no data arrives within this duration, the
|
||||
/// connection is considered dead. IRC servers typically PING every 60-120s.
|
||||
const READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
|
||||
|
||||
/// Monotonic counter to ensure unique message IDs under burst traffic.
|
||||
static MSG_SEQ: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
/// IRC over TLS channel.
|
||||
///
|
||||
/// Connects to an IRC server using TLS, joins configured channels,
|
||||
/// and forwards PRIVMSG messages to the `OpenHuman` message bus.
|
||||
/// Supports both channel messages and private messages (DMs).
|
||||
pub struct IrcChannel {
|
||||
server: String,
|
||||
port: u16,
|
||||
nickname: String,
|
||||
username: String,
|
||||
channels: Vec<String>,
|
||||
allowed_users: Vec<String>,
|
||||
server_password: Option<String>,
|
||||
nickserv_password: Option<String>,
|
||||
sasl_password: Option<String>,
|
||||
verify_tls: bool,
|
||||
/// Shared write half of the TLS stream for sending messages.
|
||||
writer: Arc<Mutex<Option<WriteHalf>>>,
|
||||
}
|
||||
|
||||
type WriteHalf = tokio::io::WriteHalf<tokio_rustls::client::TlsStream<tokio::net::TcpStream>>;
|
||||
|
||||
/// Style instruction prepended to every IRC message before it reaches the LLM.
|
||||
/// IRC clients render plain text only — no markdown, no HTML, no XML.
|
||||
const IRC_STYLE_PREFIX: &str = "\
|
||||
[context: you are responding over IRC. \
|
||||
Plain text only. No markdown, no tables, no XML/HTML tags. \
|
||||
Never use triple backtick code fences. Use a single blank line to separate blocks instead. \
|
||||
Be terse and concise. \
|
||||
Use short lines. Avoid walls of text.]\n";
|
||||
|
||||
/// Reserved bytes for the server-prepended sender prefix (`:nick!user@host `).
|
||||
const SENDER_PREFIX_RESERVE: usize = 64;
|
||||
|
||||
/// A parsed IRC message.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct IrcMessage {
|
||||
prefix: Option<String>,
|
||||
command: String,
|
||||
params: Vec<String>,
|
||||
}
|
||||
|
||||
impl IrcMessage {
|
||||
/// Parse a raw IRC line into an `IrcMessage`.
|
||||
///
|
||||
/// IRC format: `[:<prefix>] <command> [<params>] [:<trailing>]`
|
||||
fn parse(line: &str) -> Option<Self> {
|
||||
let line = line.trim_end_matches(['\r', '\n']);
|
||||
if line.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let (prefix, rest) = if let Some(stripped) = line.strip_prefix(':') {
|
||||
let space = stripped.find(' ')?;
|
||||
(Some(stripped[..space].to_string()), &stripped[space + 1..])
|
||||
} else {
|
||||
(None, line)
|
||||
};
|
||||
|
||||
// Split at trailing (first `:` after command/params)
|
||||
let (params_part, trailing) = if let Some(colon_pos) = rest.find(" :") {
|
||||
(&rest[..colon_pos], Some(&rest[colon_pos + 2..]))
|
||||
} else {
|
||||
(rest, None)
|
||||
};
|
||||
|
||||
let mut parts: Vec<&str> = params_part.split_whitespace().collect();
|
||||
if parts.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let command = parts.remove(0).to_uppercase();
|
||||
let mut params: Vec<String> = parts.iter().map(std::string::ToString::to_string).collect();
|
||||
if let Some(t) = trailing {
|
||||
params.push(t.to_string());
|
||||
}
|
||||
|
||||
Some(IrcMessage {
|
||||
prefix,
|
||||
command,
|
||||
params,
|
||||
})
|
||||
}
|
||||
|
||||
/// Extract the nickname from the prefix (nick!user@host → nick).
|
||||
fn nick(&self) -> Option<&str> {
|
||||
self.prefix.as_ref().and_then(|p| {
|
||||
let end = p.find('!').unwrap_or(p.len());
|
||||
let nick = &p[..end];
|
||||
if nick.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(nick)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Encode SASL PLAIN credentials: base64(\0nick\0password).
|
||||
fn encode_sasl_plain(nick: &str, password: &str) -> String {
|
||||
// Simple base64 encoder — avoids adding a base64 crate dependency.
|
||||
// The project's Discord channel uses a similar inline approach.
|
||||
const CHARS: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
|
||||
|
||||
let input = format!("\0{nick}\0{password}");
|
||||
let bytes = input.as_bytes();
|
||||
let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);
|
||||
|
||||
for chunk in bytes.chunks(3) {
|
||||
let b0 = u32::from(chunk[0]);
|
||||
let b1 = u32::from(chunk.get(1).copied().unwrap_or(0));
|
||||
let b2 = u32::from(chunk.get(2).copied().unwrap_or(0));
|
||||
let triple = (b0 << 16) | (b1 << 8) | b2;
|
||||
|
||||
out.push(CHARS[(triple >> 18 & 0x3F) as usize] as char);
|
||||
out.push(CHARS[(triple >> 12 & 0x3F) as usize] as char);
|
||||
|
||||
if chunk.len() > 1 {
|
||||
out.push(CHARS[(triple >> 6 & 0x3F) as usize] as char);
|
||||
} else {
|
||||
out.push('=');
|
||||
}
|
||||
|
||||
if chunk.len() > 2 {
|
||||
out.push(CHARS[(triple & 0x3F) as usize] as char);
|
||||
} else {
|
||||
out.push('=');
|
||||
}
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
/// Split a message into lines safe for IRC transmission.
|
||||
///
|
||||
/// IRC is a line-based protocol — `\r\n` terminates each command, so any
|
||||
/// newline inside a PRIVMSG payload would truncate the message and turn the
|
||||
/// remainder into garbled/invalid IRC commands.
|
||||
///
|
||||
/// This function:
|
||||
/// 1. Splits on `\n` (and strips `\r`) so each logical line becomes its own PRIVMSG.
|
||||
/// 2. Splits any line that exceeds `max_bytes` at a safe UTF-8 boundary.
|
||||
/// 3. Skips empty lines to avoid sending blank PRIVMSGs.
|
||||
fn split_message(message: &str, max_bytes: usize) -> Vec<String> {
|
||||
let mut chunks = Vec::new();
|
||||
|
||||
// Guard against max_bytes == 0 to prevent infinite loop
|
||||
if max_bytes == 0 {
|
||||
let mut full = String::new();
|
||||
for l in message
|
||||
.lines()
|
||||
.map(|l| l.trim_end_matches('\r'))
|
||||
.filter(|l| !l.is_empty())
|
||||
{
|
||||
if !full.is_empty() {
|
||||
full.push(' ');
|
||||
}
|
||||
full.push_str(l);
|
||||
}
|
||||
if full.is_empty() {
|
||||
chunks.push(String::new());
|
||||
} else {
|
||||
chunks.push(full);
|
||||
}
|
||||
return chunks;
|
||||
}
|
||||
|
||||
for line in message.split('\n') {
|
||||
let line = line.trim_end_matches('\r');
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
if line.len() <= max_bytes {
|
||||
chunks.push(line.to_string());
|
||||
continue;
|
||||
}
|
||||
|
||||
// Line exceeds max_bytes — split at safe UTF-8 boundaries
|
||||
let mut remaining = line;
|
||||
while !remaining.is_empty() {
|
||||
if remaining.len() <= max_bytes {
|
||||
chunks.push(remaining.to_string());
|
||||
break;
|
||||
}
|
||||
|
||||
let mut split_at = max_bytes;
|
||||
while split_at > 0 && !remaining.is_char_boundary(split_at) {
|
||||
split_at -= 1;
|
||||
}
|
||||
if split_at == 0 {
|
||||
// No valid boundary found going backward — advance forward instead
|
||||
split_at = max_bytes;
|
||||
while split_at < remaining.len() && !remaining.is_char_boundary(split_at) {
|
||||
split_at += 1;
|
||||
}
|
||||
}
|
||||
|
||||
chunks.push(remaining[..split_at].to_string());
|
||||
remaining = &remaining[split_at..];
|
||||
}
|
||||
}
|
||||
|
||||
if chunks.is_empty() {
|
||||
chunks.push(String::new());
|
||||
}
|
||||
|
||||
chunks
|
||||
}
|
||||
|
||||
/// Configuration for constructing an `IrcChannel`.
|
||||
pub struct IrcChannelConfig {
|
||||
pub server: String,
|
||||
pub port: u16,
|
||||
pub nickname: String,
|
||||
pub username: Option<String>,
|
||||
pub channels: Vec<String>,
|
||||
pub allowed_users: Vec<String>,
|
||||
pub server_password: Option<String>,
|
||||
pub nickserv_password: Option<String>,
|
||||
pub sasl_password: Option<String>,
|
||||
pub verify_tls: bool,
|
||||
}
|
||||
|
||||
impl IrcChannel {
|
||||
pub fn new(cfg: IrcChannelConfig) -> Self {
|
||||
let username = cfg.username.unwrap_or_else(|| cfg.nickname.clone());
|
||||
Self {
|
||||
server: cfg.server,
|
||||
port: cfg.port,
|
||||
nickname: cfg.nickname,
|
||||
username,
|
||||
channels: cfg.channels,
|
||||
allowed_users: cfg.allowed_users,
|
||||
server_password: cfg.server_password,
|
||||
nickserv_password: cfg.nickserv_password,
|
||||
sasl_password: cfg.sasl_password,
|
||||
verify_tls: cfg.verify_tls,
|
||||
writer: Arc::new(Mutex::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_user_allowed(&self, nick: &str) -> bool {
|
||||
if self.allowed_users.iter().any(|u| u == "*") {
|
||||
return true;
|
||||
}
|
||||
self.allowed_users
|
||||
.iter()
|
||||
.any(|u| u.eq_ignore_ascii_case(nick))
|
||||
}
|
||||
|
||||
/// Create a TLS connection to the IRC server.
|
||||
async fn connect(
|
||||
&self,
|
||||
) -> anyhow::Result<tokio_rustls::client::TlsStream<tokio::net::TcpStream>> {
|
||||
let addr = format!("{}:{}", self.server, self.port);
|
||||
let tcp = tokio::net::TcpStream::connect(&addr).await?;
|
||||
|
||||
let tls_config = if self.verify_tls {
|
||||
let root_store: rustls::RootCertStore =
|
||||
webpki_roots::TLS_SERVER_ROOTS.iter().cloned().collect();
|
||||
rustls::ClientConfig::builder()
|
||||
.with_root_certificates(root_store)
|
||||
.with_no_client_auth()
|
||||
} else {
|
||||
rustls::ClientConfig::builder()
|
||||
.dangerous()
|
||||
.with_custom_certificate_verifier(Arc::new(NoVerify))
|
||||
.with_no_client_auth()
|
||||
};
|
||||
|
||||
let connector = tokio_rustls::TlsConnector::from(Arc::new(tls_config));
|
||||
let domain = rustls::pki_types::ServerName::try_from(self.server.clone())?;
|
||||
let tls = connector.connect(domain, tcp).await?;
|
||||
|
||||
Ok(tls)
|
||||
}
|
||||
|
||||
/// Send a raw IRC line (appends \r\n).
|
||||
async fn send_raw(writer: &mut WriteHalf, line: &str) -> anyhow::Result<()> {
|
||||
let data = format!("{line}\r\n");
|
||||
writer.write_all(data.as_bytes()).await?;
|
||||
writer.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Certificate verifier that accepts any certificate (for `verify_tls=false`).
|
||||
#[derive(Debug)]
|
||||
struct NoVerify;
|
||||
|
||||
impl rustls::client::danger::ServerCertVerifier for NoVerify {
|
||||
fn verify_server_cert(
|
||||
&self,
|
||||
_end_entity: &rustls::pki_types::CertificateDer<'_>,
|
||||
_intermediates: &[rustls::pki_types::CertificateDer<'_>],
|
||||
_server_name: &rustls::pki_types::ServerName<'_>,
|
||||
_ocsp_response: &[u8],
|
||||
_now: rustls::pki_types::UnixTime,
|
||||
) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
|
||||
Ok(rustls::client::danger::ServerCertVerified::assertion())
|
||||
}
|
||||
|
||||
fn verify_tls12_signature(
|
||||
&self,
|
||||
_message: &[u8],
|
||||
_cert: &rustls::pki_types::CertificateDer<'_>,
|
||||
_dss: &rustls::DigitallySignedStruct,
|
||||
) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
|
||||
Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
|
||||
}
|
||||
|
||||
fn verify_tls13_signature(
|
||||
&self,
|
||||
_message: &[u8],
|
||||
_cert: &rustls::pki_types::CertificateDer<'_>,
|
||||
_dss: &rustls::DigitallySignedStruct,
|
||||
) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
|
||||
Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
|
||||
}
|
||||
|
||||
fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
|
||||
rustls::crypto::ring::default_provider()
|
||||
.signature_verification_algorithms
|
||||
.supported_schemes()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
#[allow(clippy::too_many_lines)]
|
||||
impl Channel for IrcChannel {
|
||||
fn name(&self) -> &str {
|
||||
"irc"
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> anyhow::Result<()> {
|
||||
let mut guard = self.writer.lock().await;
|
||||
let writer = guard
|
||||
.as_mut()
|
||||
.ok_or_else(|| anyhow::anyhow!("IRC not connected"))?;
|
||||
|
||||
// Calculate safe payload size:
|
||||
// 512 - sender prefix (~64 bytes for :nick!user@host) - "PRIVMSG " - target - " :" - "\r\n"
|
||||
let overhead = SENDER_PREFIX_RESERVE + 10 + message.recipient.len() + 2;
|
||||
let max_payload = 512_usize.saturating_sub(overhead);
|
||||
let chunks = split_message(&message.content, max_payload);
|
||||
|
||||
for chunk in chunks {
|
||||
Self::send_raw(writer, &format!("PRIVMSG {} :{chunk}", message.recipient)).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn listen(&self, tx: mpsc::Sender<ChannelMessage>) -> anyhow::Result<()> {
|
||||
let mut current_nick = self.nickname.clone();
|
||||
tracing::info!(
|
||||
"IRC channel connecting to {}:{} as {}...",
|
||||
self.server,
|
||||
self.port,
|
||||
current_nick
|
||||
);
|
||||
|
||||
let tls = self.connect().await?;
|
||||
let (reader, mut writer) = tokio::io::split(tls);
|
||||
|
||||
// --- SASL negotiation ---
|
||||
if self.sasl_password.is_some() {
|
||||
Self::send_raw(&mut writer, "CAP REQ :sasl").await?;
|
||||
}
|
||||
|
||||
// --- Server password ---
|
||||
if let Some(ref pass) = self.server_password {
|
||||
Self::send_raw(&mut writer, &format!("PASS {pass}")).await?;
|
||||
}
|
||||
|
||||
// --- Nick/User registration ---
|
||||
Self::send_raw(&mut writer, &format!("NICK {current_nick}")).await?;
|
||||
Self::send_raw(
|
||||
&mut writer,
|
||||
&format!("USER {} 0 * :OpenHuman", self.username),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Store writer for send()
|
||||
{
|
||||
let mut guard = self.writer.lock().await;
|
||||
*guard = Some(writer);
|
||||
}
|
||||
|
||||
let mut buf_reader = BufReader::new(reader);
|
||||
let mut line = String::new();
|
||||
let mut registered = false;
|
||||
let mut sasl_pending = self.sasl_password.is_some();
|
||||
|
||||
loop {
|
||||
line.clear();
|
||||
let n = tokio::time::timeout(READ_TIMEOUT, buf_reader.read_line(&mut line))
|
||||
.await
|
||||
.map_err(|_| {
|
||||
anyhow::anyhow!("IRC read timed out (no data for {READ_TIMEOUT:?})")
|
||||
})??;
|
||||
if n == 0 {
|
||||
anyhow::bail!("IRC connection closed by server");
|
||||
}
|
||||
|
||||
let Some(msg) = IrcMessage::parse(&line) else {
|
||||
continue;
|
||||
};
|
||||
|
||||
match msg.command.as_str() {
|
||||
"PING" => {
|
||||
let token = msg.params.first().map_or("", String::as_str);
|
||||
let mut guard = self.writer.lock().await;
|
||||
if let Some(ref mut w) = *guard {
|
||||
Self::send_raw(w, &format!("PONG :{token}")).await?;
|
||||
}
|
||||
}
|
||||
|
||||
// CAP responses for SASL
|
||||
"CAP" => {
|
||||
if sasl_pending && msg.params.iter().any(|p| p.contains("sasl")) {
|
||||
if msg.params.iter().any(|p| p.contains("ACK")) {
|
||||
// CAP * ACK :sasl — server accepted, start SASL auth
|
||||
let mut guard = self.writer.lock().await;
|
||||
if let Some(ref mut w) = *guard {
|
||||
Self::send_raw(w, "AUTHENTICATE PLAIN").await?;
|
||||
}
|
||||
} else if msg.params.iter().any(|p| p.contains("NAK")) {
|
||||
// CAP * NAK :sasl — server rejected SASL, proceed without it
|
||||
tracing::warn!(
|
||||
"IRC server does not support SASL, continuing without it"
|
||||
);
|
||||
sasl_pending = false;
|
||||
let mut guard = self.writer.lock().await;
|
||||
if let Some(ref mut w) = *guard {
|
||||
Self::send_raw(w, "CAP END").await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
"AUTHENTICATE" => {
|
||||
// Server sends "AUTHENTICATE +" to request credentials
|
||||
if sasl_pending && msg.params.first().is_some_and(|p| p == "+") {
|
||||
// sasl_password is loaded from runtime config, not hard-coded
|
||||
if let Some(password) = self.sasl_password.as_deref() {
|
||||
let encoded = encode_sasl_plain(¤t_nick, password);
|
||||
let mut guard = self.writer.lock().await;
|
||||
if let Some(ref mut w) = *guard {
|
||||
Self::send_raw(w, &format!("AUTHENTICATE {encoded}")).await?;
|
||||
}
|
||||
} else {
|
||||
// SASL was requested but no password is configured; abort SASL
|
||||
tracing::warn!(
|
||||
"SASL authentication requested but no SASL password is configured; aborting SASL"
|
||||
);
|
||||
sasl_pending = false;
|
||||
let mut guard = self.writer.lock().await;
|
||||
if let Some(ref mut w) = *guard {
|
||||
Self::send_raw(w, "CAP END").await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// RPL_SASLSUCCESS (903) — SASL done, end CAP
|
||||
"903" => {
|
||||
sasl_pending = false;
|
||||
let mut guard = self.writer.lock().await;
|
||||
if let Some(ref mut w) = *guard {
|
||||
Self::send_raw(w, "CAP END").await?;
|
||||
}
|
||||
}
|
||||
|
||||
// SASL failure (904, 905, 906, 907)
|
||||
"904" | "905" | "906" | "907" => {
|
||||
tracing::warn!("IRC SASL authentication failed ({})", msg.command);
|
||||
sasl_pending = false;
|
||||
let mut guard = self.writer.lock().await;
|
||||
if let Some(ref mut w) = *guard {
|
||||
Self::send_raw(w, "CAP END").await?;
|
||||
}
|
||||
}
|
||||
|
||||
// RPL_WELCOME — registration complete
|
||||
"001" => {
|
||||
registered = true;
|
||||
tracing::info!("IRC registered as {}", current_nick);
|
||||
|
||||
// NickServ authentication
|
||||
if let Some(ref pass) = self.nickserv_password {
|
||||
let mut guard = self.writer.lock().await;
|
||||
if let Some(ref mut w) = *guard {
|
||||
Self::send_raw(w, &format!("PRIVMSG NickServ :IDENTIFY {pass}"))
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
|
||||
// Join channels
|
||||
for chan in &self.channels {
|
||||
let mut guard = self.writer.lock().await;
|
||||
if let Some(ref mut w) = *guard {
|
||||
Self::send_raw(w, &format!("JOIN {chan}")).await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ERR_NICKNAMEINUSE (433)
|
||||
"433" => {
|
||||
let alt = format!("{current_nick}_");
|
||||
tracing::warn!("IRC nickname {current_nick} is in use, trying {alt}");
|
||||
let mut guard = self.writer.lock().await;
|
||||
if let Some(ref mut w) = *guard {
|
||||
Self::send_raw(w, &format!("NICK {alt}")).await?;
|
||||
}
|
||||
current_nick = alt;
|
||||
}
|
||||
|
||||
"PRIVMSG" => {
|
||||
if !registered {
|
||||
continue;
|
||||
}
|
||||
|
||||
let target = msg.params.first().map_or("", String::as_str);
|
||||
let text = msg.params.get(1).map_or("", String::as_str);
|
||||
let sender_nick = msg.nick().unwrap_or("unknown");
|
||||
|
||||
// Skip messages from NickServ/ChanServ
|
||||
if sender_nick.eq_ignore_ascii_case("NickServ")
|
||||
|| sender_nick.eq_ignore_ascii_case("ChanServ")
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
if !self.is_user_allowed(sender_nick) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Determine reply target: if sent to a channel, reply to channel;
|
||||
// if DM (target == our nick), reply to sender
|
||||
let is_channel = target.starts_with('#') || target.starts_with('&');
|
||||
let reply_target = if is_channel {
|
||||
target.to_string()
|
||||
} else {
|
||||
sender_nick.to_string()
|
||||
};
|
||||
let content = if is_channel {
|
||||
format!("{IRC_STYLE_PREFIX}<{sender_nick}> {text}")
|
||||
} else {
|
||||
format!("{IRC_STYLE_PREFIX}{text}")
|
||||
};
|
||||
|
||||
let seq = MSG_SEQ.fetch_add(1, Ordering::Relaxed);
|
||||
let channel_msg = ChannelMessage {
|
||||
id: format!("irc_{}_{seq}", chrono::Utc::now().timestamp_millis()),
|
||||
sender: sender_nick.to_string(),
|
||||
reply_target,
|
||||
content,
|
||||
channel: "irc".to_string(),
|
||||
timestamp: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
thread_ts: None,
|
||||
};
|
||||
|
||||
if tx.send(channel_msg).await.is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
// ERR_PASSWDMISMATCH (464) or other fatal errors
|
||||
"464" => {
|
||||
anyhow::bail!("IRC password mismatch");
|
||||
}
|
||||
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
// Lightweight connectivity check: TLS connect + QUIT
|
||||
match self.connect().await {
|
||||
Ok(tls) => {
|
||||
let (_, mut writer) = tokio::io::split(tls);
|
||||
let _ = Self::send_raw(&mut writer, "QUIT :health check").await;
|
||||
true
|
||||
}
|
||||
Err(_) => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "irc_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
#[cfg(any(test, debug_assertions))]
|
||||
pub mod test_support {
|
||||
//! Debug-build seams for raw integration tests. They cover IRC parsing and
|
||||
//! framing helpers without opening a TLS socket.
|
||||
|
||||
use super::*;
|
||||
|
||||
pub fn parse_line_for_test(
|
||||
line: &str,
|
||||
) -> Option<(Option<String>, String, Vec<String>, Option<String>)> {
|
||||
IrcMessage::parse(line).map(|msg| {
|
||||
let nick = msg.nick().map(str::to_string);
|
||||
(msg.prefix, msg.command, msg.params, nick)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn split_message_for_test(message: &str, max_bytes: usize) -> Vec<String> {
|
||||
split_message(message, max_bytes)
|
||||
}
|
||||
|
||||
pub fn encode_sasl_plain_for_test(nick: &str, password: &str) -> String {
|
||||
encode_sasl_plain(nick, password)
|
||||
}
|
||||
|
||||
pub fn is_user_allowed_for_test(allowed_users: Vec<String>, nick: &str) -> bool {
|
||||
IrcChannel::new(IrcChannelConfig {
|
||||
server: "irc.example.test".to_string(),
|
||||
port: 6697,
|
||||
nickname: "openhuman".to_string(),
|
||||
username: None,
|
||||
channels: vec!["#ops".to_string()],
|
||||
allowed_users,
|
||||
server_password: None,
|
||||
nickserv_password: None,
|
||||
sasl_password: None,
|
||||
verify_tls: false,
|
||||
})
|
||||
.is_user_allowed(nick)
|
||||
}
|
||||
}
|
||||
pub use tinychannels::providers::irc::*;
|
||||
|
||||
@@ -1,406 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
// ── IRC message parsing ──────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn parse_privmsg_with_prefix() {
|
||||
let msg = IrcMessage::parse(":nick!user@host PRIVMSG #channel :Hello world").unwrap();
|
||||
assert_eq!(msg.prefix.as_deref(), Some("nick!user@host"));
|
||||
assert_eq!(msg.command, "PRIVMSG");
|
||||
assert_eq!(msg.params, vec!["#channel", "Hello world"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_privmsg_dm() {
|
||||
let msg = IrcMessage::parse(":alice!a@host PRIVMSG botname :hi there").unwrap();
|
||||
assert_eq!(msg.command, "PRIVMSG");
|
||||
assert_eq!(msg.params, vec!["botname", "hi there"]);
|
||||
assert_eq!(msg.nick(), Some("alice"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_ping() {
|
||||
let msg = IrcMessage::parse("PING :server.example.com").unwrap();
|
||||
assert!(msg.prefix.is_none());
|
||||
assert_eq!(msg.command, "PING");
|
||||
assert_eq!(msg.params, vec!["server.example.com"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_numeric_reply() {
|
||||
let msg = IrcMessage::parse(":server 001 botname :Welcome to the IRC network").unwrap();
|
||||
assert_eq!(msg.prefix.as_deref(), Some("server"));
|
||||
assert_eq!(msg.command, "001");
|
||||
assert_eq!(msg.params, vec!["botname", "Welcome to the IRC network"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_no_trailing() {
|
||||
let msg = IrcMessage::parse(":server 433 * botname").unwrap();
|
||||
assert_eq!(msg.command, "433");
|
||||
assert_eq!(msg.params, vec!["*", "botname"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_cap_ack() {
|
||||
let msg = IrcMessage::parse(":server CAP * ACK :sasl").unwrap();
|
||||
assert_eq!(msg.command, "CAP");
|
||||
assert_eq!(msg.params, vec!["*", "ACK", "sasl"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_empty_line_returns_none() {
|
||||
assert!(IrcMessage::parse("").is_none());
|
||||
assert!(IrcMessage::parse("\r\n").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_strips_crlf() {
|
||||
let msg = IrcMessage::parse("PING :test\r\n").unwrap();
|
||||
assert_eq!(msg.params, vec!["test"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_command_uppercase() {
|
||||
let msg = IrcMessage::parse("ping :test").unwrap();
|
||||
assert_eq!(msg.command, "PING");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nick_extraction_full_prefix() {
|
||||
let msg = IrcMessage::parse(":nick!user@host PRIVMSG #ch :msg").unwrap();
|
||||
assert_eq!(msg.nick(), Some("nick"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nick_extraction_nick_only() {
|
||||
let msg = IrcMessage::parse(":server 001 bot :Welcome").unwrap();
|
||||
assert_eq!(msg.nick(), Some("server"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nick_extraction_no_prefix() {
|
||||
let msg = IrcMessage::parse("PING :token").unwrap();
|
||||
assert_eq!(msg.nick(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_authenticate_plus() {
|
||||
let msg = IrcMessage::parse("AUTHENTICATE +").unwrap();
|
||||
assert_eq!(msg.command, "AUTHENTICATE");
|
||||
assert_eq!(msg.params, vec!["+"]);
|
||||
}
|
||||
|
||||
// ── SASL PLAIN encoding ─────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn sasl_plain_encode() {
|
||||
let encoded = encode_sasl_plain("jilles", "sesame");
|
||||
// \0jilles\0sesame → base64
|
||||
assert_eq!(encoded, "AGppbGxlcwBzZXNhbWU=");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sasl_plain_empty_password() {
|
||||
let encoded = encode_sasl_plain("nick", "");
|
||||
// \0nick\0 → base64
|
||||
assert_eq!(encoded, "AG5pY2sA");
|
||||
}
|
||||
|
||||
// ── Message splitting ───────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn split_short_message() {
|
||||
let chunks = split_message("hello", 400);
|
||||
assert_eq!(chunks, vec!["hello"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_long_message() {
|
||||
let msg = "a".repeat(800);
|
||||
let chunks = split_message(&msg, 400);
|
||||
assert_eq!(chunks.len(), 2);
|
||||
assert_eq!(chunks[0].len(), 400);
|
||||
assert_eq!(chunks[1].len(), 400);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_exact_boundary() {
|
||||
let msg = "a".repeat(400);
|
||||
let chunks = split_message(&msg, 400);
|
||||
assert_eq!(chunks.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_unicode_safe() {
|
||||
// 'é' is 2 bytes in UTF-8; splitting at byte 3 would split mid-char
|
||||
let msg = "ééé"; // 6 bytes
|
||||
let chunks = split_message(msg, 3);
|
||||
// Should split at char boundary (2 bytes), not mid-char
|
||||
assert_eq!(chunks.len(), 3);
|
||||
assert_eq!(chunks[0], "é");
|
||||
assert_eq!(chunks[1], "é");
|
||||
assert_eq!(chunks[2], "é");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_empty_message() {
|
||||
let chunks = split_message("", 400);
|
||||
assert_eq!(chunks, vec![""]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_newlines_into_separate_lines() {
|
||||
let chunks = split_message("line one\nline two\nline three", 400);
|
||||
assert_eq!(chunks, vec!["line one", "line two", "line three"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_crlf_newlines() {
|
||||
let chunks = split_message("hello\r\nworld", 400);
|
||||
assert_eq!(chunks, vec!["hello", "world"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_skips_empty_lines() {
|
||||
let chunks = split_message("hello\n\n\nworld", 400);
|
||||
assert_eq!(chunks, vec!["hello", "world"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_trailing_newline() {
|
||||
let chunks = split_message("hello\n", 400);
|
||||
assert_eq!(chunks, vec!["hello"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_multiline_with_long_line() {
|
||||
let long = "a".repeat(800);
|
||||
let msg = format!("short\n{long}\nend");
|
||||
let chunks = split_message(&msg, 400);
|
||||
assert_eq!(chunks.len(), 4);
|
||||
assert_eq!(chunks[0], "short");
|
||||
assert_eq!(chunks[1].len(), 400);
|
||||
assert_eq!(chunks[2].len(), 400);
|
||||
assert_eq!(chunks[3], "end");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_only_newlines() {
|
||||
let chunks = split_message("\n\n\n", 400);
|
||||
assert_eq!(chunks, vec![""]);
|
||||
}
|
||||
|
||||
// ── Allowlist ───────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn wildcard_allows_anyone() {
|
||||
let ch = make_channel();
|
||||
// Default make_channel has wildcard
|
||||
assert!(ch.is_user_allowed("anyone"));
|
||||
assert!(ch.is_user_allowed("stranger"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn specific_user_allowed() {
|
||||
let ch = IrcChannel::new(IrcChannelConfig {
|
||||
server: "irc.test".into(),
|
||||
port: 6697,
|
||||
nickname: "bot".into(),
|
||||
username: None,
|
||||
channels: vec![],
|
||||
allowed_users: vec!["alice".into(), "bob".into()],
|
||||
server_password: None,
|
||||
nickserv_password: None,
|
||||
sasl_password: None,
|
||||
verify_tls: true,
|
||||
});
|
||||
assert!(ch.is_user_allowed("alice"));
|
||||
assert!(ch.is_user_allowed("bob"));
|
||||
assert!(!ch.is_user_allowed("eve"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allowlist_case_insensitive() {
|
||||
let ch = IrcChannel::new(IrcChannelConfig {
|
||||
server: "irc.test".into(),
|
||||
port: 6697,
|
||||
nickname: "bot".into(),
|
||||
username: None,
|
||||
channels: vec![],
|
||||
allowed_users: vec!["Alice".into()],
|
||||
server_password: None,
|
||||
nickserv_password: None,
|
||||
sasl_password: None,
|
||||
verify_tls: true,
|
||||
});
|
||||
assert!(ch.is_user_allowed("alice"));
|
||||
assert!(ch.is_user_allowed("ALICE"));
|
||||
assert!(ch.is_user_allowed("Alice"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_allowlist_denies_all() {
|
||||
let ch = IrcChannel::new(IrcChannelConfig {
|
||||
server: "irc.test".into(),
|
||||
port: 6697,
|
||||
nickname: "bot".into(),
|
||||
username: None,
|
||||
channels: vec![],
|
||||
allowed_users: vec![],
|
||||
server_password: None,
|
||||
nickserv_password: None,
|
||||
sasl_password: None,
|
||||
verify_tls: true,
|
||||
});
|
||||
assert!(!ch.is_user_allowed("anyone"));
|
||||
}
|
||||
|
||||
// ── Constructor ─────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn new_defaults_username_to_nickname() {
|
||||
let ch = IrcChannel::new(IrcChannelConfig {
|
||||
server: "irc.test".into(),
|
||||
port: 6697,
|
||||
nickname: "mybot".into(),
|
||||
username: None,
|
||||
channels: vec![],
|
||||
allowed_users: vec![],
|
||||
server_password: None,
|
||||
nickserv_password: None,
|
||||
sasl_password: None,
|
||||
verify_tls: true,
|
||||
});
|
||||
assert_eq!(ch.username, "mybot");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_uses_explicit_username() {
|
||||
let ch = IrcChannel::new(IrcChannelConfig {
|
||||
server: "irc.test".into(),
|
||||
port: 6697,
|
||||
nickname: "mybot".into(),
|
||||
username: Some("customuser".into()),
|
||||
channels: vec![],
|
||||
allowed_users: vec![],
|
||||
server_password: None,
|
||||
nickserv_password: None,
|
||||
sasl_password: None,
|
||||
verify_tls: true,
|
||||
});
|
||||
assert_eq!(ch.username, "customuser");
|
||||
assert_eq!(ch.nickname, "mybot");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn name_returns_irc() {
|
||||
let ch = make_channel();
|
||||
assert_eq!(ch.name(), "irc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_stores_all_fields() {
|
||||
let ch = IrcChannel::new(IrcChannelConfig {
|
||||
server: "irc.example.com".into(),
|
||||
port: 6697,
|
||||
nickname: "zcbot".into(),
|
||||
username: Some("openhuman".into()),
|
||||
channels: vec!["#test".into()],
|
||||
allowed_users: vec!["alice".into()],
|
||||
server_password: Some("serverpass".into()),
|
||||
nickserv_password: Some("nspass".into()),
|
||||
sasl_password: Some("saslpass".into()),
|
||||
verify_tls: false,
|
||||
});
|
||||
assert_eq!(ch.server, "irc.example.com");
|
||||
assert_eq!(ch.port, 6697);
|
||||
assert_eq!(ch.nickname, "zcbot");
|
||||
assert_eq!(ch.username, "openhuman");
|
||||
assert_eq!(ch.channels, vec!["#test"]);
|
||||
assert_eq!(ch.allowed_users, vec!["alice"]);
|
||||
assert_eq!(ch.server_password.as_deref(), Some("serverpass"));
|
||||
assert_eq!(ch.nickserv_password.as_deref(), Some("nspass"));
|
||||
assert_eq!(ch.sasl_password.as_deref(), Some("saslpass"));
|
||||
assert!(!ch.verify_tls);
|
||||
}
|
||||
|
||||
// ── Config serde ────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn irc_config_serde_roundtrip() {
|
||||
use crate::openhuman::config::schema::IrcConfig;
|
||||
|
||||
let config = IrcConfig {
|
||||
server: "irc.example.com".into(),
|
||||
port: 6697,
|
||||
nickname: "zcbot".into(),
|
||||
username: Some("openhuman".into()),
|
||||
channels: vec!["#test".into(), "#dev".into()],
|
||||
allowed_users: vec!["alice".into()],
|
||||
server_password: None,
|
||||
nickserv_password: Some("secret".into()),
|
||||
sasl_password: None,
|
||||
verify_tls: Some(true),
|
||||
};
|
||||
|
||||
let toml_str = toml::to_string(&config).unwrap();
|
||||
let parsed: IrcConfig = toml::from_str(&toml_str).unwrap();
|
||||
assert_eq!(parsed.server, "irc.example.com");
|
||||
assert_eq!(parsed.port, 6697);
|
||||
assert_eq!(parsed.nickname, "zcbot");
|
||||
assert_eq!(parsed.username.as_deref(), Some("openhuman"));
|
||||
assert_eq!(parsed.channels, vec!["#test", "#dev"]);
|
||||
assert_eq!(parsed.allowed_users, vec!["alice"]);
|
||||
assert!(parsed.server_password.is_none());
|
||||
assert_eq!(parsed.nickserv_password.as_deref(), Some("secret"));
|
||||
assert!(parsed.sasl_password.is_none());
|
||||
assert_eq!(parsed.verify_tls, Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn irc_config_minimal_toml() {
|
||||
use crate::openhuman::config::schema::IrcConfig;
|
||||
|
||||
let toml_str = r#"
|
||||
server = "irc.example.com"
|
||||
nickname = "bot"
|
||||
"#;
|
||||
let parsed: IrcConfig = toml::from_str(toml_str).unwrap();
|
||||
assert_eq!(parsed.server, "irc.example.com");
|
||||
assert_eq!(parsed.port, 6697); // default
|
||||
assert_eq!(parsed.nickname, "bot");
|
||||
assert!(parsed.username.is_none());
|
||||
assert!(parsed.channels.is_empty());
|
||||
assert!(parsed.allowed_users.is_empty());
|
||||
assert!(parsed.server_password.is_none());
|
||||
assert!(parsed.nickserv_password.is_none());
|
||||
assert!(parsed.sasl_password.is_none());
|
||||
assert!(parsed.verify_tls.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn irc_config_default_port() {
|
||||
use crate::openhuman::config::schema::IrcConfig;
|
||||
|
||||
let json = r#"{"server":"irc.test","nickname":"bot"}"#;
|
||||
let parsed: IrcConfig = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(parsed.port, 6697);
|
||||
}
|
||||
|
||||
// ── Helpers ─────────────────────────────────────────────
|
||||
|
||||
fn make_channel() -> IrcChannel {
|
||||
IrcChannel::new(IrcChannelConfig {
|
||||
server: "irc.example.com".into(),
|
||||
port: 6697,
|
||||
nickname: "zcbot".into(),
|
||||
username: None,
|
||||
channels: vec!["#openhuman".into()],
|
||||
allowed_users: vec!["*".into()],
|
||||
server_password: None,
|
||||
nickserv_password: None,
|
||||
sasl_password: None,
|
||||
verify_tls: true,
|
||||
})
|
||||
}
|
||||
@@ -1,992 +1 @@
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
use async_trait::async_trait;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use prost::Message as ProstMessage;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::sync::RwLock;
|
||||
use tokio_tungstenite::tungstenite::Message as WsMsg;
|
||||
use uuid::Uuid;
|
||||
|
||||
const FEISHU_BASE_URL: &str = "https://open.feishu.cn/open-apis";
|
||||
const FEISHU_WS_BASE_URL: &str = "https://open.feishu.cn";
|
||||
const LARK_BASE_URL: &str = "https://open.larksuite.com/open-apis";
|
||||
const LARK_WS_BASE_URL: &str = "https://open.larksuite.com";
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// Feishu WebSocket long-connection: pbbp2.proto frame codec
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[derive(Clone, PartialEq, prost::Message)]
|
||||
struct PbHeader {
|
||||
#[prost(string, tag = "1")]
|
||||
pub key: String,
|
||||
#[prost(string, tag = "2")]
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
/// Feishu WS frame (pbbp2.proto).
|
||||
/// method=0 → CONTROL (ping/pong) method=1 → DATA (events)
|
||||
#[derive(Clone, PartialEq, prost::Message)]
|
||||
struct PbFrame {
|
||||
#[prost(uint64, tag = "1")]
|
||||
pub seq_id: u64,
|
||||
#[prost(uint64, tag = "2")]
|
||||
pub log_id: u64,
|
||||
#[prost(int32, tag = "3")]
|
||||
pub service: i32,
|
||||
#[prost(int32, tag = "4")]
|
||||
pub method: i32,
|
||||
#[prost(message, repeated, tag = "5")]
|
||||
pub headers: Vec<PbHeader>,
|
||||
#[prost(bytes = "vec", optional, tag = "8")]
|
||||
pub payload: Option<Vec<u8>>,
|
||||
}
|
||||
|
||||
impl PbFrame {
|
||||
fn header_value<'a>(&'a self, key: &str) -> &'a str {
|
||||
self.headers
|
||||
.iter()
|
||||
.find(|h| h.key == key)
|
||||
.map(|h| h.value.as_str())
|
||||
.unwrap_or("")
|
||||
}
|
||||
}
|
||||
|
||||
/// Server-sent client config (parsed from pong payload)
|
||||
#[derive(Debug, serde::Deserialize, Default, Clone)]
|
||||
struct WsClientConfig {
|
||||
#[serde(rename = "PingInterval")]
|
||||
ping_interval: Option<u64>,
|
||||
}
|
||||
|
||||
/// POST /callback/ws/endpoint response
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct WsEndpointResp {
|
||||
code: i32,
|
||||
#[serde(default)]
|
||||
msg: Option<String>,
|
||||
#[serde(default)]
|
||||
data: Option<WsEndpoint>,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct WsEndpoint {
|
||||
#[serde(rename = "URL")]
|
||||
url: String,
|
||||
#[serde(rename = "ClientConfig")]
|
||||
client_config: Option<WsClientConfig>,
|
||||
}
|
||||
|
||||
/// LarkEvent envelope (method=1 / type=event payload)
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct LarkEvent {
|
||||
header: LarkEventHeader,
|
||||
event: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct LarkEventHeader {
|
||||
event_type: String,
|
||||
#[allow(dead_code)]
|
||||
event_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct MsgReceivePayload {
|
||||
sender: LarkSender,
|
||||
message: LarkMessage,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct LarkSender {
|
||||
sender_id: LarkSenderId,
|
||||
#[serde(default)]
|
||||
sender_type: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize, Default)]
|
||||
struct LarkSenderId {
|
||||
open_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct LarkMessage {
|
||||
message_id: String,
|
||||
chat_id: String,
|
||||
chat_type: String,
|
||||
message_type: String,
|
||||
#[serde(default)]
|
||||
content: String,
|
||||
#[serde(default)]
|
||||
mentions: Vec<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Heartbeat timeout for WS connection — must be larger than ping_interval (default 120 s).
|
||||
/// If no binary frame (pong or event) is received within this window, reconnect.
|
||||
const WS_HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(300);
|
||||
|
||||
/// Returns true when the WebSocket frame indicates live traffic that should
|
||||
/// refresh the heartbeat watchdog.
|
||||
fn should_refresh_last_recv(msg: &WsMsg) -> bool {
|
||||
matches!(msg, WsMsg::Binary(_) | WsMsg::Ping(_) | WsMsg::Pong(_))
|
||||
}
|
||||
|
||||
/// Lark/Feishu channel.
|
||||
///
|
||||
/// Supports two receive modes (configured via `receive_mode` in config):
|
||||
/// - **`websocket`** (default): persistent WSS long-connection; no public URL needed.
|
||||
/// - **`webhook`**: HTTP callback server; requires a public HTTPS endpoint.
|
||||
pub struct LarkChannel {
|
||||
app_id: String,
|
||||
app_secret: String,
|
||||
verification_token: String,
|
||||
port: Option<u16>,
|
||||
allowed_users: Vec<String>,
|
||||
/// When true, use Feishu (CN) endpoints; when false, use Lark (international).
|
||||
use_feishu: bool,
|
||||
/// How to receive events: WebSocket long-connection or HTTP webhook.
|
||||
receive_mode: crate::openhuman::config::schema::LarkReceiveMode,
|
||||
/// Cached tenant access token
|
||||
tenant_token: Arc<RwLock<Option<String>>>,
|
||||
/// Dedup set: WS message_ids seen in last ~30 min to prevent double-dispatch
|
||||
ws_seen_ids: Arc<RwLock<HashMap<String, Instant>>>,
|
||||
}
|
||||
|
||||
impl LarkChannel {
|
||||
pub fn new(
|
||||
app_id: String,
|
||||
app_secret: String,
|
||||
verification_token: String,
|
||||
port: Option<u16>,
|
||||
allowed_users: Vec<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
app_id,
|
||||
app_secret,
|
||||
verification_token,
|
||||
port,
|
||||
allowed_users,
|
||||
use_feishu: true,
|
||||
receive_mode: crate::openhuman::config::schema::LarkReceiveMode::default(),
|
||||
tenant_token: Arc::new(RwLock::new(None)),
|
||||
ws_seen_ids: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build from `LarkConfig` (preserves `use_feishu` and `receive_mode`).
|
||||
pub fn from_config(config: &crate::openhuman::config::schema::LarkConfig) -> Self {
|
||||
let mut ch = Self::new(
|
||||
config.app_id.clone(),
|
||||
config.app_secret.clone(),
|
||||
config.verification_token.clone().unwrap_or_default(),
|
||||
config.port,
|
||||
config.allowed_users.clone(),
|
||||
);
|
||||
ch.use_feishu = config.use_feishu;
|
||||
ch.receive_mode = config.receive_mode.clone();
|
||||
ch
|
||||
}
|
||||
|
||||
fn http_client(&self) -> reqwest::Client {
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.lark")
|
||||
}
|
||||
|
||||
fn api_base(&self) -> &'static str {
|
||||
if self.use_feishu {
|
||||
FEISHU_BASE_URL
|
||||
} else {
|
||||
LARK_BASE_URL
|
||||
}
|
||||
}
|
||||
|
||||
fn ws_base(&self) -> &'static str {
|
||||
if self.use_feishu {
|
||||
FEISHU_WS_BASE_URL
|
||||
} else {
|
||||
LARK_WS_BASE_URL
|
||||
}
|
||||
}
|
||||
|
||||
fn tenant_access_token_url(&self) -> String {
|
||||
format!("{}/auth/v3/tenant_access_token/internal", self.api_base())
|
||||
}
|
||||
|
||||
fn send_message_url(&self) -> String {
|
||||
format!("{}/im/v1/messages?receive_id_type=chat_id", self.api_base())
|
||||
}
|
||||
|
||||
/// POST /callback/ws/endpoint → (wss_url, client_config)
|
||||
async fn get_ws_endpoint(&self) -> anyhow::Result<(String, WsClientConfig)> {
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(format!("{}/callback/ws/endpoint", self.ws_base()))
|
||||
.header("locale", if self.use_feishu { "zh" } else { "en" })
|
||||
.json(&serde_json::json!({
|
||||
"AppID": self.app_id,
|
||||
"AppSecret": self.app_secret,
|
||||
}))
|
||||
.send()
|
||||
.await?
|
||||
.json::<WsEndpointResp>()
|
||||
.await?;
|
||||
if resp.code != 0 {
|
||||
anyhow::bail!(
|
||||
"Lark WS endpoint failed: code={} msg={}",
|
||||
resp.code,
|
||||
resp.msg.as_deref().unwrap_or("(none)")
|
||||
);
|
||||
}
|
||||
let ep = resp
|
||||
.data
|
||||
.ok_or_else(|| anyhow::anyhow!("Lark WS endpoint: empty data"))?;
|
||||
Ok((ep.url, ep.client_config.unwrap_or_default()))
|
||||
}
|
||||
|
||||
/// WS long-connection event loop. Returns Ok(()) when the connection closes
|
||||
/// (the caller reconnects).
|
||||
#[allow(clippy::too_many_lines)]
|
||||
async fn listen_ws(&self, tx: tokio::sync::mpsc::Sender<ChannelMessage>) -> anyhow::Result<()> {
|
||||
let (wss_url, client_config) = self.get_ws_endpoint().await?;
|
||||
let service_id = wss_url
|
||||
.split('?')
|
||||
.nth(1)
|
||||
.and_then(|qs| {
|
||||
qs.split('&')
|
||||
.find(|kv| kv.starts_with("service_id="))
|
||||
.and_then(|kv| kv.split('=').nth(1))
|
||||
.and_then(|v| v.parse::<i32>().ok())
|
||||
})
|
||||
.unwrap_or(0);
|
||||
tracing::info!("Lark: connecting to {wss_url}");
|
||||
|
||||
let (ws_stream, _) = tokio_tungstenite::connect_async(&wss_url).await?;
|
||||
let (mut write, mut read) = ws_stream.split();
|
||||
tracing::info!("Lark: WS connected (service_id={service_id})");
|
||||
|
||||
let mut ping_secs = client_config.ping_interval.unwrap_or(120).max(10);
|
||||
let mut hb_interval = tokio::time::interval(Duration::from_secs(ping_secs));
|
||||
let mut timeout_check = tokio::time::interval(Duration::from_secs(10));
|
||||
hb_interval.tick().await; // consume immediate tick
|
||||
|
||||
let mut seq: u64 = 0;
|
||||
let mut last_recv = Instant::now();
|
||||
|
||||
// Send initial ping immediately (like the official SDK) so the server
|
||||
// starts responding with pongs and we can calibrate the ping_interval.
|
||||
seq = seq.wrapping_add(1);
|
||||
let initial_ping = PbFrame {
|
||||
seq_id: seq,
|
||||
log_id: 0,
|
||||
service: service_id,
|
||||
method: 0,
|
||||
headers: vec![PbHeader {
|
||||
key: "type".into(),
|
||||
value: "ping".into(),
|
||||
}],
|
||||
payload: None,
|
||||
};
|
||||
if write
|
||||
.send(WsMsg::Binary(initial_ping.encode_to_vec()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
anyhow::bail!("Lark: initial ping failed");
|
||||
}
|
||||
// message_id → (fragment_slots, created_at) for multi-part reassembly
|
||||
type FragEntry = (Vec<Option<Vec<u8>>>, Instant);
|
||||
let mut frag_cache: HashMap<String, FragEntry> = HashMap::new();
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
|
||||
_ = hb_interval.tick() => {
|
||||
seq = seq.wrapping_add(1);
|
||||
let ping = PbFrame {
|
||||
seq_id: seq, log_id: 0, service: service_id, method: 0,
|
||||
headers: vec![PbHeader { key: "type".into(), value: "ping".into() }],
|
||||
payload: None,
|
||||
};
|
||||
if write.send(WsMsg::Binary(ping.encode_to_vec())).await.is_err() {
|
||||
tracing::warn!("Lark: ping failed, reconnecting");
|
||||
break;
|
||||
}
|
||||
// GC stale fragments > 5 min
|
||||
let cutoff = Instant::now().checked_sub(Duration::from_secs(300)).unwrap_or(Instant::now());
|
||||
frag_cache.retain(|_, (_, ts)| *ts > cutoff);
|
||||
}
|
||||
|
||||
_ = timeout_check.tick() => {
|
||||
if last_recv.elapsed() > WS_HEARTBEAT_TIMEOUT {
|
||||
tracing::warn!("Lark: heartbeat timeout, reconnecting");
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
msg = read.next() => {
|
||||
let raw = match msg {
|
||||
Some(Ok(ws_msg)) => {
|
||||
if should_refresh_last_recv(&ws_msg) {
|
||||
last_recv = Instant::now();
|
||||
}
|
||||
match ws_msg {
|
||||
WsMsg::Binary(b) => b,
|
||||
WsMsg::Ping(d) => { let _ = write.send(WsMsg::Pong(d)).await; continue; }
|
||||
WsMsg::Pong(_) => continue,
|
||||
WsMsg::Close(_) => { tracing::info!("Lark: WS closed — reconnecting"); break; }
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
None => { tracing::info!("Lark: WS closed — reconnecting"); break; }
|
||||
Some(Err(e)) => { tracing::error!("Lark: WS read error: {e}"); break; }
|
||||
};
|
||||
|
||||
let frame = match PbFrame::decode(&raw[..]) {
|
||||
Ok(f) => f,
|
||||
Err(e) => { tracing::error!("Lark: proto decode: {e}"); continue; }
|
||||
};
|
||||
|
||||
// CONTROL frame
|
||||
if frame.method == 0 {
|
||||
if frame.header_value("type") == "pong" {
|
||||
if let Some(p) = &frame.payload {
|
||||
if let Ok(cfg) = serde_json::from_slice::<WsClientConfig>(p) {
|
||||
if let Some(secs) = cfg.ping_interval {
|
||||
let secs = secs.max(10);
|
||||
if secs != ping_secs {
|
||||
ping_secs = secs;
|
||||
hb_interval = tokio::time::interval(Duration::from_secs(ping_secs));
|
||||
tracing::info!("Lark: ping_interval → {ping_secs}s");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
// DATA frame
|
||||
let msg_type = frame.header_value("type").to_string();
|
||||
let msg_id = frame.header_value("message_id").to_string();
|
||||
let sum = frame.header_value("sum").parse::<usize>().unwrap_or(1);
|
||||
let seq_num = frame.header_value("seq").parse::<usize>().unwrap_or(0);
|
||||
|
||||
// ACK immediately (Feishu requires within 3 s)
|
||||
{
|
||||
let mut ack = frame.clone();
|
||||
ack.payload = Some(br#"{"code":200,"headers":{},"data":[]}"#.to_vec());
|
||||
ack.headers.push(PbHeader { key: "biz_rt".into(), value: "0".into() });
|
||||
let _ = write.send(WsMsg::Binary(ack.encode_to_vec())).await;
|
||||
}
|
||||
|
||||
// Fragment reassembly
|
||||
let sum = if sum == 0 { 1 } else { sum };
|
||||
let payload: Vec<u8> = if sum == 1 || msg_id.is_empty() || seq_num >= sum {
|
||||
frame.payload.clone().unwrap_or_default()
|
||||
} else {
|
||||
let entry = frag_cache.entry(msg_id.clone())
|
||||
.or_insert_with(|| (vec![None; sum], Instant::now()));
|
||||
if entry.0.len() != sum { *entry = (vec![None; sum], Instant::now()); }
|
||||
entry.0[seq_num] = frame.payload.clone();
|
||||
if entry.0.iter().all(|s| s.is_some()) {
|
||||
let full: Vec<u8> = entry.0.iter()
|
||||
.flat_map(|s| s.as_deref().unwrap_or(&[]))
|
||||
.copied().collect();
|
||||
frag_cache.remove(&msg_id);
|
||||
full
|
||||
} else { continue; }
|
||||
};
|
||||
|
||||
if msg_type != "event" { continue; }
|
||||
|
||||
let event: LarkEvent = match serde_json::from_slice(&payload) {
|
||||
Ok(e) => e,
|
||||
Err(e) => { tracing::error!("Lark: event JSON: {e}"); continue; }
|
||||
};
|
||||
if event.header.event_type != "im.message.receive_v1" { continue; }
|
||||
|
||||
let recv: MsgReceivePayload = match serde_json::from_value(event.event) {
|
||||
Ok(r) => r,
|
||||
Err(e) => { tracing::error!("Lark: payload parse: {e}"); continue; }
|
||||
};
|
||||
|
||||
if recv.sender.sender_type == "app" || recv.sender.sender_type == "bot" { continue; }
|
||||
|
||||
let sender_open_id = recv.sender.sender_id.open_id.as_deref().unwrap_or("");
|
||||
if !self.is_user_allowed(sender_open_id) {
|
||||
tracing::warn!("Lark WS: ignoring {sender_open_id} (not in allowed_users)");
|
||||
continue;
|
||||
}
|
||||
|
||||
let lark_msg = &recv.message;
|
||||
|
||||
// Dedup
|
||||
{
|
||||
let now = Instant::now();
|
||||
let mut seen = self.ws_seen_ids.write().await;
|
||||
// GC
|
||||
seen.retain(|_, t| now.duration_since(*t) < Duration::from_secs(30 * 60));
|
||||
if seen.contains_key(&lark_msg.message_id) {
|
||||
tracing::debug!("Lark WS: dup {}", lark_msg.message_id);
|
||||
continue;
|
||||
}
|
||||
seen.insert(lark_msg.message_id.clone(), now);
|
||||
}
|
||||
|
||||
// Decode content by type (mirrors clawdbot-feishu parsing)
|
||||
let text = match lark_msg.message_type.as_str() {
|
||||
"text" => {
|
||||
let v: serde_json::Value = match serde_json::from_str(&lark_msg.content) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
match v.get("text").and_then(|t| t.as_str()).filter(|s| !s.is_empty()) {
|
||||
Some(t) => t.to_string(),
|
||||
None => continue,
|
||||
}
|
||||
}
|
||||
"post" => match parse_post_content(&lark_msg.content) {
|
||||
Some(t) => t,
|
||||
None => continue,
|
||||
},
|
||||
_ => { tracing::debug!("Lark WS: skipping unsupported type '{}'", lark_msg.message_type); continue; }
|
||||
};
|
||||
|
||||
// Strip @_user_N placeholders
|
||||
let text = strip_at_placeholders(&text);
|
||||
let text = text.trim().to_string();
|
||||
if text.is_empty() { continue; }
|
||||
|
||||
// Group-chat: only respond when explicitly @-mentioned
|
||||
if lark_msg.chat_type == "group" && !should_respond_in_group(&lark_msg.mentions) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let channel_msg = ChannelMessage {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
sender: lark_msg.chat_id.clone(),
|
||||
reply_target: lark_msg.chat_id.clone(),
|
||||
content: text,
|
||||
channel: "lark".to_string(),
|
||||
timestamp: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
thread_ts: None,
|
||||
};
|
||||
|
||||
tracing::debug!("Lark WS: message in {}", lark_msg.chat_id);
|
||||
if tx.send(channel_msg).await.is_err() { break; }
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check if a user open_id is allowed
|
||||
fn is_user_allowed(&self, open_id: &str) -> bool {
|
||||
self.allowed_users.iter().any(|u| u == "*" || u == open_id)
|
||||
}
|
||||
|
||||
/// Get or refresh tenant access token
|
||||
async fn get_tenant_access_token(&self) -> anyhow::Result<String> {
|
||||
// Check cache first
|
||||
{
|
||||
let cached = self.tenant_token.read().await;
|
||||
if let Some(ref token) = *cached {
|
||||
return Ok(token.clone());
|
||||
}
|
||||
}
|
||||
|
||||
let url = self.tenant_access_token_url();
|
||||
let body = serde_json::json!({
|
||||
"app_id": self.app_id,
|
||||
"app_secret": self.app_secret,
|
||||
});
|
||||
|
||||
let resp = self.http_client().post(&url).json(&body).send().await?;
|
||||
let data: serde_json::Value = resp.json().await?;
|
||||
|
||||
let code = data.get("code").and_then(|c| c.as_i64()).unwrap_or(-1);
|
||||
if code != 0 {
|
||||
let msg = data
|
||||
.get("msg")
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or("unknown error");
|
||||
anyhow::bail!("Lark tenant_access_token failed: {msg}");
|
||||
}
|
||||
|
||||
let token = data
|
||||
.get("tenant_access_token")
|
||||
.and_then(|t| t.as_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("missing tenant_access_token in response"))?
|
||||
.to_string();
|
||||
|
||||
// Cache it
|
||||
{
|
||||
let mut cached = self.tenant_token.write().await;
|
||||
*cached = Some(token.clone());
|
||||
}
|
||||
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
/// Invalidate cached token (called on 401)
|
||||
async fn invalidate_token(&self) {
|
||||
let mut cached = self.tenant_token.write().await;
|
||||
*cached = None;
|
||||
}
|
||||
|
||||
/// Parse an event callback payload and extract text messages
|
||||
pub fn parse_event_payload(&self, payload: &serde_json::Value) -> Vec<ChannelMessage> {
|
||||
let mut messages = Vec::new();
|
||||
|
||||
// Lark event v2 structure:
|
||||
// { "header": { "event_type": "im.message.receive_v1" }, "event": { "message": { ... }, "sender": { ... } } }
|
||||
let event_type = payload
|
||||
.pointer("/header/event_type")
|
||||
.and_then(|e| e.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
if event_type != "im.message.receive_v1" {
|
||||
return messages;
|
||||
}
|
||||
|
||||
let event = match payload.get("event") {
|
||||
Some(e) => e,
|
||||
None => return messages,
|
||||
};
|
||||
|
||||
// Extract sender open_id
|
||||
let open_id = event
|
||||
.pointer("/sender/sender_id/open_id")
|
||||
.and_then(|s| s.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
if open_id.is_empty() {
|
||||
return messages;
|
||||
}
|
||||
|
||||
// Check allowlist
|
||||
if !self.is_user_allowed(open_id) {
|
||||
tracing::warn!("Lark: ignoring message from unauthorized user: {open_id}");
|
||||
return messages;
|
||||
}
|
||||
|
||||
// Extract message content (text and post supported)
|
||||
let msg_type = event
|
||||
.pointer("/message/message_type")
|
||||
.and_then(|t| t.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
let content_str = event
|
||||
.pointer("/message/content")
|
||||
.and_then(|c| c.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
let text: String = match msg_type {
|
||||
"text" => {
|
||||
let extracted = serde_json::from_str::<serde_json::Value>(content_str)
|
||||
.ok()
|
||||
.and_then(|v| {
|
||||
v.get("text")
|
||||
.and_then(|t| t.as_str())
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(String::from)
|
||||
});
|
||||
match extracted {
|
||||
Some(t) => t,
|
||||
None => return messages,
|
||||
}
|
||||
}
|
||||
"post" => match parse_post_content(content_str) {
|
||||
Some(t) => t,
|
||||
None => return messages,
|
||||
},
|
||||
_ => {
|
||||
tracing::debug!("Lark: skipping unsupported message type: {msg_type}");
|
||||
return messages;
|
||||
}
|
||||
};
|
||||
|
||||
let timestamp = event
|
||||
.pointer("/message/create_time")
|
||||
.and_then(|t| t.as_str())
|
||||
.and_then(|t| t.parse::<u64>().ok())
|
||||
// Lark timestamps are in milliseconds
|
||||
.map(|ms| ms / 1000)
|
||||
.unwrap_or_else(|| {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
});
|
||||
|
||||
let chat_id = event
|
||||
.pointer("/message/chat_id")
|
||||
.and_then(|c| c.as_str())
|
||||
.unwrap_or(open_id);
|
||||
|
||||
messages.push(ChannelMessage {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
sender: chat_id.to_string(),
|
||||
reply_target: chat_id.to_string(),
|
||||
content: text,
|
||||
channel: "lark".to_string(),
|
||||
timestamp,
|
||||
thread_ts: None,
|
||||
});
|
||||
|
||||
messages
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for LarkChannel {
|
||||
fn name(&self) -> &str {
|
||||
"lark"
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> anyhow::Result<()> {
|
||||
let token = self.get_tenant_access_token().await?;
|
||||
let url = self.send_message_url();
|
||||
|
||||
let content = serde_json::json!({ "text": message.content }).to_string();
|
||||
let body = serde_json::json!({
|
||||
"receive_id": message.recipient,
|
||||
"msg_type": "text",
|
||||
"content": content,
|
||||
});
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {token}"))
|
||||
.header("Content-Type", "application/json; charset=utf-8")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if resp.status().as_u16() == 401 {
|
||||
// Token expired, invalidate and retry once
|
||||
self.invalidate_token().await;
|
||||
let new_token = self.get_tenant_access_token().await?;
|
||||
let retry_resp = self
|
||||
.http_client()
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {new_token}"))
|
||||
.header("Content-Type", "application/json; charset=utf-8")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !retry_resp.status().is_success() {
|
||||
let err = retry_resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!("Lark send failed after token refresh: {err}");
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!("Lark send failed: {err}");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn listen(&self, tx: tokio::sync::mpsc::Sender<ChannelMessage>) -> anyhow::Result<()> {
|
||||
use crate::openhuman::config::schema::LarkReceiveMode;
|
||||
match self.receive_mode {
|
||||
LarkReceiveMode::Websocket => self.listen_ws(tx).await,
|
||||
LarkReceiveMode::Webhook => self.listen_http(tx).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
self.get_tenant_access_token().await.is_ok()
|
||||
}
|
||||
}
|
||||
|
||||
impl LarkChannel {
|
||||
/// HTTP callback server (legacy — requires a public endpoint).
|
||||
/// Use `listen()` (WS long-connection) for new deployments.
|
||||
pub async fn listen_http(
|
||||
&self,
|
||||
tx: tokio::sync::mpsc::Sender<ChannelMessage>,
|
||||
) -> anyhow::Result<()> {
|
||||
use axum::{extract::State, routing::post, Json, Router};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct AppState {
|
||||
verification_token: String,
|
||||
channel: Arc<LarkChannel>,
|
||||
tx: tokio::sync::mpsc::Sender<ChannelMessage>,
|
||||
}
|
||||
|
||||
async fn handle_event(
|
||||
State(state): State<AppState>,
|
||||
Json(payload): Json<serde_json::Value>,
|
||||
) -> axum::response::Response {
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::IntoResponse;
|
||||
|
||||
// URL verification challenge
|
||||
if let Some(challenge) = payload.get("challenge").and_then(|c| c.as_str()) {
|
||||
// Verify token if present
|
||||
let token_ok = payload
|
||||
.get("token")
|
||||
.and_then(|t| t.as_str())
|
||||
.is_none_or(|t| t == state.verification_token);
|
||||
|
||||
if !token_ok {
|
||||
return (StatusCode::FORBIDDEN, "invalid token").into_response();
|
||||
}
|
||||
|
||||
let resp = serde_json::json!({ "challenge": challenge });
|
||||
return (StatusCode::OK, Json(resp)).into_response();
|
||||
}
|
||||
|
||||
// Parse event messages
|
||||
let messages = state.channel.parse_event_payload(&payload);
|
||||
for msg in messages {
|
||||
if state.tx.send(msg).await.is_err() {
|
||||
tracing::warn!("Lark: message channel closed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
(StatusCode::OK, "ok").into_response()
|
||||
}
|
||||
|
||||
let port = self.port.ok_or_else(|| {
|
||||
anyhow::anyhow!("Lark webhook mode requires `port` to be set in [channels_config.lark]")
|
||||
})?;
|
||||
|
||||
let state = AppState {
|
||||
verification_token: self.verification_token.clone(),
|
||||
channel: Arc::new(LarkChannel::new(
|
||||
self.app_id.clone(),
|
||||
self.app_secret.clone(),
|
||||
self.verification_token.clone(),
|
||||
None,
|
||||
self.allowed_users.clone(),
|
||||
)),
|
||||
tx,
|
||||
};
|
||||
|
||||
let app = Router::new()
|
||||
.route("/lark", post(handle_event))
|
||||
.with_state(state);
|
||||
|
||||
let addr = std::net::SocketAddr::from(([0, 0, 0, 0], port));
|
||||
tracing::info!("Lark event callback server listening on {addr}");
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(addr).await?;
|
||||
axum::serve(listener, app).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
// WS helper functions
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
/// Flatten a Feishu `post` rich-text message to plain text.
|
||||
///
|
||||
/// Returns `None` when the content cannot be parsed or yields no usable text,
|
||||
/// so callers can simply `continue` rather than forwarding a meaningless
|
||||
/// placeholder string to the agent.
|
||||
fn parse_post_content(content: &str) -> Option<String> {
|
||||
let parsed = serde_json::from_str::<serde_json::Value>(content).ok()?;
|
||||
let locale = parsed
|
||||
.get("zh_cn")
|
||||
.or_else(|| parsed.get("en_us"))
|
||||
.or_else(|| {
|
||||
parsed
|
||||
.as_object()
|
||||
.and_then(|m| m.values().find(|v| v.is_object()))
|
||||
})?;
|
||||
|
||||
let mut text = String::new();
|
||||
|
||||
if let Some(title) = locale
|
||||
.get("title")
|
||||
.and_then(|t| t.as_str())
|
||||
.filter(|s| !s.is_empty())
|
||||
{
|
||||
text.push_str(title);
|
||||
text.push_str("\n\n");
|
||||
}
|
||||
|
||||
if let Some(paragraphs) = locale.get("content").and_then(|c| c.as_array()) {
|
||||
for para in paragraphs {
|
||||
if let Some(elements) = para.as_array() {
|
||||
for el in elements {
|
||||
match el.get("tag").and_then(|t| t.as_str()).unwrap_or("") {
|
||||
"text" => {
|
||||
if let Some(t) = el.get("text").and_then(|t| t.as_str()) {
|
||||
text.push_str(t);
|
||||
}
|
||||
}
|
||||
"a" => {
|
||||
text.push_str(
|
||||
el.get("text")
|
||||
.and_then(|t| t.as_str())
|
||||
.filter(|s| !s.is_empty())
|
||||
.or_else(|| el.get("href").and_then(|h| h.as_str()))
|
||||
.unwrap_or(""),
|
||||
);
|
||||
}
|
||||
"at" => {
|
||||
let n = el
|
||||
.get("user_name")
|
||||
.and_then(|n| n.as_str())
|
||||
.or_else(|| el.get("user_id").and_then(|i| i.as_str()))
|
||||
.unwrap_or("user");
|
||||
text.push('@');
|
||||
text.push_str(n);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
text.push('\n');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let result = text.trim().to_string();
|
||||
if result.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(result)
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove `@_user_N` placeholder tokens injected by Feishu in group chats.
|
||||
fn strip_at_placeholders(text: &str) -> String {
|
||||
let mut result = String::with_capacity(text.len());
|
||||
let mut chars = text.char_indices().peekable();
|
||||
while let Some((_, ch)) = chars.next() {
|
||||
if ch == '@' {
|
||||
let rest: String = chars.clone().map(|(_, c)| c).collect();
|
||||
if let Some(after) = rest.strip_prefix("_user_") {
|
||||
let digit_count = after.chars().take_while(|c| c.is_ascii_digit()).count();
|
||||
if digit_count == 0 {
|
||||
result.push(ch);
|
||||
continue;
|
||||
}
|
||||
let skip = "_user_".len() + digit_count;
|
||||
for _ in 0..skip {
|
||||
chars.next();
|
||||
}
|
||||
if chars.peek().map(|(_, c)| *c == ' ').unwrap_or(false) {
|
||||
chars.next();
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
result.push(ch);
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// In group chats, only respond when the bot is explicitly @-mentioned.
|
||||
fn should_respond_in_group(mentions: &[serde_json::Value]) -> bool {
|
||||
!mentions.is_empty()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "lark_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
#[cfg(any(test, debug_assertions))]
|
||||
pub mod test_support {
|
||||
//! Debug-build helpers for raw integration tests. These expose pure parser
|
||||
//! seams without widening the production API surface.
|
||||
|
||||
use super::*;
|
||||
use prost::Message as ProstMessage;
|
||||
use tokio_tungstenite::tungstenite::Message as WsMsg;
|
||||
|
||||
pub fn parse_post_content_for_test(content: &str) -> Option<String> {
|
||||
parse_post_content(content)
|
||||
}
|
||||
|
||||
pub fn strip_at_placeholders_for_test(text: &str) -> String {
|
||||
strip_at_placeholders(text)
|
||||
}
|
||||
|
||||
pub fn should_respond_in_group_for_test(mentions: &[serde_json::Value]) -> bool {
|
||||
should_respond_in_group(mentions)
|
||||
}
|
||||
|
||||
pub fn should_refresh_last_recv_for_test(msg: &WsMsg) -> bool {
|
||||
should_refresh_last_recv(msg)
|
||||
}
|
||||
|
||||
pub fn endpoint_response_for_test(raw: &str) -> anyhow::Result<(String, Option<u64>)> {
|
||||
let resp = serde_json::from_str::<WsEndpointResp>(raw)?;
|
||||
if resp.code != 0 {
|
||||
anyhow::bail!(
|
||||
"Lark WS endpoint failed: code={} msg={}",
|
||||
resp.code,
|
||||
resp.msg.as_deref().unwrap_or("(none)")
|
||||
);
|
||||
}
|
||||
let ep = resp
|
||||
.data
|
||||
.ok_or_else(|| anyhow::anyhow!("Lark WS endpoint: empty data"))?;
|
||||
Ok((ep.url, ep.client_config.and_then(|cfg| cfg.ping_interval)))
|
||||
}
|
||||
|
||||
pub fn encode_frame_for_test(
|
||||
seq_id: u64,
|
||||
method: i32,
|
||||
frame_type: &str,
|
||||
payload: Option<Vec<u8>>,
|
||||
) -> Vec<u8> {
|
||||
PbFrame {
|
||||
seq_id,
|
||||
log_id: 0,
|
||||
service: 7,
|
||||
method,
|
||||
headers: vec![PbHeader {
|
||||
key: "type".to_string(),
|
||||
value: frame_type.to_string(),
|
||||
}],
|
||||
payload,
|
||||
}
|
||||
.encode_to_vec()
|
||||
}
|
||||
|
||||
pub fn decode_frame_for_test(
|
||||
raw: &[u8],
|
||||
) -> anyhow::Result<(u64, i32, String, Option<Vec<u8>>)> {
|
||||
let frame = PbFrame::decode(raw)?;
|
||||
Ok((
|
||||
frame.seq_id,
|
||||
frame.method,
|
||||
frame.header_value("type").to_string(),
|
||||
frame.payload,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn endpoint_urls_for_test(use_feishu: bool) -> (String, String) {
|
||||
let mut channel = LarkChannel::new(
|
||||
"app".to_string(),
|
||||
"secret".to_string(),
|
||||
String::new(),
|
||||
None,
|
||||
Vec::new(),
|
||||
);
|
||||
channel.use_feishu = use_feishu;
|
||||
(
|
||||
channel.tenant_access_token_url(),
|
||||
channel.send_message_url(),
|
||||
)
|
||||
}
|
||||
}
|
||||
pub use tinychannels::providers::lark::*;
|
||||
|
||||
@@ -1,613 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
fn make_channel() -> LarkChannel {
|
||||
LarkChannel::new(
|
||||
"cli_test_app_id".into(),
|
||||
"test_app_secret".into(),
|
||||
"test_verification_token".into(),
|
||||
None,
|
||||
vec!["ou_testuser123".into()],
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_channel_name() {
|
||||
let ch = make_channel();
|
||||
assert_eq!(ch.name(), "lark");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_ws_activity_refreshes_heartbeat_watchdog() {
|
||||
assert!(should_refresh_last_recv(&WsMsg::Binary(
|
||||
vec![1, 2, 3].into()
|
||||
)));
|
||||
assert!(should_refresh_last_recv(&WsMsg::Ping(vec![9, 9].into())));
|
||||
assert!(should_refresh_last_recv(&WsMsg::Pong(vec![8, 8].into())));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_ws_non_activity_frames_do_not_refresh_heartbeat_watchdog() {
|
||||
assert!(!should_refresh_last_recv(&WsMsg::Text("hello".into())));
|
||||
assert!(!should_refresh_last_recv(&WsMsg::Close(None)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_user_allowed_exact() {
|
||||
let ch = make_channel();
|
||||
assert!(ch.is_user_allowed("ou_testuser123"));
|
||||
assert!(!ch.is_user_allowed("ou_other"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_user_allowed_wildcard() {
|
||||
let ch = LarkChannel::new(
|
||||
"id".into(),
|
||||
"secret".into(),
|
||||
"token".into(),
|
||||
None,
|
||||
vec!["*".into()],
|
||||
);
|
||||
assert!(ch.is_user_allowed("ou_anyone"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_user_denied_empty() {
|
||||
let ch = LarkChannel::new("id".into(), "secret".into(), "token".into(), None, vec![]);
|
||||
assert!(!ch.is_user_allowed("ou_anyone"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_challenge() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"challenge": "abc123",
|
||||
"token": "test_verification_token",
|
||||
"type": "url_verification"
|
||||
});
|
||||
// Challenge payloads should not produce messages
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_valid_text_message() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"header": {
|
||||
"event_type": "im.message.receive_v1"
|
||||
},
|
||||
"event": {
|
||||
"sender": {
|
||||
"sender_id": {
|
||||
"open_id": "ou_testuser123"
|
||||
}
|
||||
},
|
||||
"message": {
|
||||
"message_type": "text",
|
||||
"content": "{\"text\":\"Hello OpenHuman!\"}",
|
||||
"chat_id": "oc_chat123",
|
||||
"create_time": "1699999999000"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(msgs[0].content, "Hello OpenHuman!");
|
||||
assert_eq!(msgs[0].sender, "oc_chat123");
|
||||
assert_eq!(msgs[0].channel, "lark");
|
||||
assert_eq!(msgs[0].timestamp, 1_699_999_999);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_unauthorized_user() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" },
|
||||
"event": {
|
||||
"sender": { "sender_id": { "open_id": "ou_unauthorized" } },
|
||||
"message": {
|
||||
"message_type": "text",
|
||||
"content": "{\"text\":\"spam\"}",
|
||||
"chat_id": "oc_chat",
|
||||
"create_time": "1000"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_non_text_message_skipped() {
|
||||
let ch = LarkChannel::new(
|
||||
"id".into(),
|
||||
"secret".into(),
|
||||
"token".into(),
|
||||
None,
|
||||
vec!["*".into()],
|
||||
);
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" },
|
||||
"event": {
|
||||
"sender": { "sender_id": { "open_id": "ou_user" } },
|
||||
"message": {
|
||||
"message_type": "image",
|
||||
"content": "{}",
|
||||
"chat_id": "oc_chat"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_empty_text_skipped() {
|
||||
let ch = LarkChannel::new(
|
||||
"id".into(),
|
||||
"secret".into(),
|
||||
"token".into(),
|
||||
None,
|
||||
vec!["*".into()],
|
||||
);
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" },
|
||||
"event": {
|
||||
"sender": { "sender_id": { "open_id": "ou_user" } },
|
||||
"message": {
|
||||
"message_type": "text",
|
||||
"content": "{\"text\":\"\"}",
|
||||
"chat_id": "oc_chat"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_wrong_event_type() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.chat.disbanded_v1" },
|
||||
"event": {}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_missing_sender() {
|
||||
let ch = LarkChannel::new(
|
||||
"id".into(),
|
||||
"secret".into(),
|
||||
"token".into(),
|
||||
None,
|
||||
vec!["*".into()],
|
||||
);
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" },
|
||||
"event": {
|
||||
"message": {
|
||||
"message_type": "text",
|
||||
"content": "{\"text\":\"hello\"}",
|
||||
"chat_id": "oc_chat"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_unicode_message() {
|
||||
let ch = LarkChannel::new(
|
||||
"id".into(),
|
||||
"secret".into(),
|
||||
"token".into(),
|
||||
None,
|
||||
vec!["*".into()],
|
||||
);
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" },
|
||||
"event": {
|
||||
"sender": { "sender_id": { "open_id": "ou_user" } },
|
||||
"message": {
|
||||
"message_type": "text",
|
||||
"content": "{\"text\":\"Hello world 🌍\"}",
|
||||
"chat_id": "oc_chat",
|
||||
"create_time": "1000"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(msgs[0].content, "Hello world 🌍");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_missing_event() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" }
|
||||
});
|
||||
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_invalid_content_json() {
|
||||
let ch = LarkChannel::new(
|
||||
"id".into(),
|
||||
"secret".into(),
|
||||
"token".into(),
|
||||
None,
|
||||
vec!["*".into()],
|
||||
);
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" },
|
||||
"event": {
|
||||
"sender": { "sender_id": { "open_id": "ou_user" } },
|
||||
"message": {
|
||||
"message_type": "text",
|
||||
"content": "not valid json",
|
||||
"chat_id": "oc_chat"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_config_serde() {
|
||||
use crate::openhuman::config::schema::{LarkConfig, LarkReceiveMode};
|
||||
let lc = LarkConfig {
|
||||
app_id: "cli_app123".into(),
|
||||
app_secret: "secret456".into(),
|
||||
encrypt_key: None,
|
||||
verification_token: Some("vtoken789".into()),
|
||||
allowed_users: vec!["ou_user1".into(), "ou_user2".into()],
|
||||
use_feishu: false,
|
||||
receive_mode: LarkReceiveMode::default(),
|
||||
port: None,
|
||||
};
|
||||
let json = serde_json::to_string(&lc).unwrap();
|
||||
let parsed: LarkConfig = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(parsed.app_id, "cli_app123");
|
||||
assert_eq!(parsed.app_secret, "secret456");
|
||||
assert_eq!(parsed.verification_token.as_deref(), Some("vtoken789"));
|
||||
assert_eq!(parsed.allowed_users.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_config_toml_roundtrip() {
|
||||
use crate::openhuman::config::schema::{LarkConfig, LarkReceiveMode};
|
||||
let lc = LarkConfig {
|
||||
app_id: "app".into(),
|
||||
app_secret: "secret".into(),
|
||||
encrypt_key: None,
|
||||
verification_token: Some("tok".into()),
|
||||
allowed_users: vec!["*".into()],
|
||||
use_feishu: false,
|
||||
receive_mode: LarkReceiveMode::Webhook,
|
||||
port: Some(9898),
|
||||
};
|
||||
let toml_str = toml::to_string(&lc).unwrap();
|
||||
let parsed: LarkConfig = toml::from_str(&toml_str).unwrap();
|
||||
assert_eq!(parsed.app_id, "app");
|
||||
assert_eq!(parsed.verification_token.as_deref(), Some("tok"));
|
||||
assert_eq!(parsed.allowed_users, vec!["*"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_config_defaults_optional_fields() {
|
||||
use crate::openhuman::config::schema::{LarkConfig, LarkReceiveMode};
|
||||
let json = r#"{"app_id":"a","app_secret":"s"}"#;
|
||||
let parsed: LarkConfig = serde_json::from_str(json).unwrap();
|
||||
assert!(parsed.verification_token.is_none());
|
||||
assert!(parsed.allowed_users.is_empty());
|
||||
assert_eq!(parsed.receive_mode, LarkReceiveMode::Websocket);
|
||||
assert!(parsed.port.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_from_config_preserves_mode_and_region() {
|
||||
use crate::openhuman::config::schema::{LarkConfig, LarkReceiveMode};
|
||||
|
||||
let cfg = LarkConfig {
|
||||
app_id: "cli_app123".into(),
|
||||
app_secret: "secret456".into(),
|
||||
encrypt_key: None,
|
||||
verification_token: Some("vtoken789".into()),
|
||||
allowed_users: vec!["*".into()],
|
||||
use_feishu: false,
|
||||
receive_mode: LarkReceiveMode::Webhook,
|
||||
port: Some(9898),
|
||||
};
|
||||
|
||||
let ch = LarkChannel::from_config(&cfg);
|
||||
|
||||
assert_eq!(ch.api_base(), LARK_BASE_URL);
|
||||
assert_eq!(ch.ws_base(), LARK_WS_BASE_URL);
|
||||
assert_eq!(ch.receive_mode, LarkReceiveMode::Webhook);
|
||||
assert_eq!(ch.port, Some(9898));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_fallback_sender_to_open_id() {
|
||||
// When chat_id is missing, sender should fall back to open_id
|
||||
let ch = LarkChannel::new(
|
||||
"id".into(),
|
||||
"secret".into(),
|
||||
"token".into(),
|
||||
None,
|
||||
vec!["*".into()],
|
||||
);
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" },
|
||||
"event": {
|
||||
"sender": { "sender_id": { "open_id": "ou_user" } },
|
||||
"message": {
|
||||
"message_type": "text",
|
||||
"content": "{\"text\":\"hello\"}",
|
||||
"create_time": "1000"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(msgs[0].sender, "ou_user");
|
||||
}
|
||||
|
||||
// ── parse_post_content ─────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn parse_post_content_returns_zh_cn_locale_content() {
|
||||
let post = serde_json::json!({
|
||||
"zh_cn": {
|
||||
"title": "标题",
|
||||
"content": [[{"tag": "text", "text": "你好"}]]
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
let out = parse_post_content(&post).expect("parsed");
|
||||
assert!(out.contains("标题"));
|
||||
assert!(out.contains("你好"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_post_content_falls_back_to_en_us_when_zh_cn_missing() {
|
||||
let post = serde_json::json!({
|
||||
"en_us": {
|
||||
"title": "Hello",
|
||||
"content": [[{"tag": "text", "text": "world"}]]
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
let out = parse_post_content(&post).expect("parsed");
|
||||
assert!(out.contains("Hello"));
|
||||
assert!(out.contains("world"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_post_content_returns_none_for_invalid_json() {
|
||||
assert!(parse_post_content("not json").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_post_content_handles_links_and_mentions() {
|
||||
let post = serde_json::json!({
|
||||
"zh_cn": {
|
||||
"title": "T",
|
||||
"content": [[
|
||||
{"tag": "text", "text": "pre "},
|
||||
{"tag": "a", "text": "link", "href": "https://x"},
|
||||
{"tag": "at", "user_name": "alice"}
|
||||
]]
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
let out = parse_post_content(&post).expect("parsed");
|
||||
assert!(out.contains("link"));
|
||||
assert!(out.contains("@alice"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_post_content_falls_back_to_href_when_anchor_text_missing() {
|
||||
// Anchor without `text` must surface the `href` — otherwise the
|
||||
// link is invisible in the rendered message.
|
||||
let post = serde_json::json!({
|
||||
"zh_cn": {
|
||||
"title": "T",
|
||||
"content": [[
|
||||
{"tag": "text", "text": "see "},
|
||||
{"tag": "a", "href": "https://example.com/no-text"}
|
||||
]]
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
let out = parse_post_content(&post).expect("parsed");
|
||||
assert!(
|
||||
out.contains("https://example.com/no-text"),
|
||||
"href fallback should surface when anchor has no text, got: {out}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_post_content_returns_none_when_all_sections_empty() {
|
||||
let post = serde_json::json!({ "zh_cn": { "title": "" } }).to_string();
|
||||
assert!(parse_post_content(&post).is_none());
|
||||
}
|
||||
|
||||
// ── strip_at_placeholders ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn strip_at_placeholders_removes_user_tokens() {
|
||||
assert_eq!(strip_at_placeholders("hello @_user_1 world"), "hello world");
|
||||
assert_eq!(
|
||||
strip_at_placeholders("@_user_42 message here"),
|
||||
"message here"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_at_placeholders_preserves_real_at_mentions() {
|
||||
assert_eq!(strip_at_placeholders("hello @alice"), "hello @alice");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_at_placeholders_handles_multiple_placeholders() {
|
||||
assert_eq!(strip_at_placeholders("@_user_1 hi @_user_2 bye"), "hi bye");
|
||||
}
|
||||
|
||||
// ── should_respond_in_group ────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn should_respond_in_group_requires_nonempty_mentions() {
|
||||
assert!(!should_respond_in_group(&[]));
|
||||
assert!(should_respond_in_group(&[
|
||||
serde_json::json!({"key": "val"})
|
||||
]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_refresh_last_recv_true_for_binary_ping_pong() {
|
||||
use tokio_tungstenite::tungstenite::Message as WsMsg;
|
||||
assert!(should_refresh_last_recv(&WsMsg::Binary(vec![1, 2, 3])));
|
||||
assert!(should_refresh_last_recv(&WsMsg::Ping(vec![])));
|
||||
assert!(should_refresh_last_recv(&WsMsg::Pong(vec![])));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_refresh_last_recv_false_for_text_and_close() {
|
||||
use tokio_tungstenite::tungstenite::Message as WsMsg;
|
||||
assert!(!should_refresh_last_recv(&WsMsg::Text("hello".into())));
|
||||
assert!(!should_refresh_last_recv(&WsMsg::Close(None)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_new_stores_fields_and_allowlist() {
|
||||
let ch = LarkChannel::new(
|
||||
"app_id".into(),
|
||||
"secret".into(),
|
||||
"verify".into(),
|
||||
Some(3001),
|
||||
vec!["u1".into(), "u2".into()],
|
||||
);
|
||||
assert_eq!(ch.app_id, "app_id");
|
||||
assert_eq!(ch.port, Some(3001));
|
||||
assert_eq!(ch.allowed_users.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_is_user_allowed_wildcard_allows_everyone() {
|
||||
let ch = LarkChannel::new("a".into(), "s".into(), "v".into(), None, vec!["*".into()]);
|
||||
assert!(ch.is_user_allowed("anyone"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_is_user_allowed_empty_allowlist_blocks_everyone() {
|
||||
// Empty allowlist matches nothing — explicit guard against the
|
||||
// "accidentally allowing all users" bug.
|
||||
let ch = LarkChannel::new("a".into(), "s".into(), "v".into(), None, vec![]);
|
||||
assert!(!ch.is_user_allowed("anyone"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_is_user_allowed_respects_allowlist() {
|
||||
let ch = LarkChannel::new("a".into(), "s".into(), "v".into(), None, vec!["u1".into()]);
|
||||
assert!(ch.is_user_allowed("u1"));
|
||||
assert!(!ch.is_user_allowed("u2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_event_payload_empty_object_returns_no_messages() {
|
||||
let ch = make_channel();
|
||||
let msgs = ch.parse_event_payload(&serde_json::json!({}));
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_event_payload_ignores_unsupported_message_type() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" },
|
||||
"event": {
|
||||
"sender": { "sender_id": { "open_id": "ou_testuser123" } },
|
||||
"message": {
|
||||
"message_type": "image",
|
||||
"content": r#"{"image_key":"abc"}"#,
|
||||
"create_time": "1700000000000",
|
||||
"chat_id": "chat_xyz"
|
||||
}
|
||||
}
|
||||
});
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_event_payload_empty_sender_returns_no_messages() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" },
|
||||
"event": {
|
||||
"sender": { "sender_id": { "open_id": "" } },
|
||||
"message": {
|
||||
"message_type": "text",
|
||||
"content": r#"{"text":"hi"}"#,
|
||||
"create_time": "1700000000000"
|
||||
}
|
||||
}
|
||||
});
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_event_payload_missing_event_returns_empty() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" }
|
||||
});
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lark_parse_event_payload_post_type_extracts_readable_text() {
|
||||
let ch = make_channel();
|
||||
let post_content = serde_json::json!({
|
||||
"zh_cn": {
|
||||
"title": "Title",
|
||||
"content": [[{"tag":"text","text":"Body"}]]
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
let payload = serde_json::json!({
|
||||
"header": { "event_type": "im.message.receive_v1" },
|
||||
"event": {
|
||||
"sender": { "sender_id": { "open_id": "ou_testuser123" } },
|
||||
"message": {
|
||||
"message_type": "post",
|
||||
"content": post_content,
|
||||
"create_time": "1700000000000",
|
||||
"chat_id": "chat_xyz"
|
||||
}
|
||||
}
|
||||
});
|
||||
let msgs = ch.parse_event_payload(&payload);
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert!(msgs[0].content.contains("Title"));
|
||||
}
|
||||
@@ -1,397 +1 @@
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
use async_trait::async_trait;
|
||||
use uuid::Uuid;
|
||||
|
||||
/// Linq channel — uses the Linq Partner V3 API for iMessage, RCS, and SMS.
|
||||
///
|
||||
/// This channel operates in webhook mode (push-based) rather than polling.
|
||||
/// The `listen` method here is a keepalive placeholder; inbound delivery depends on
|
||||
/// your deployment wiring Linq webhooks to the app.
|
||||
pub struct LinqChannel {
|
||||
api_token: String,
|
||||
from_phone: String,
|
||||
allowed_senders: Vec<String>,
|
||||
client: reqwest::Client,
|
||||
}
|
||||
|
||||
const LINQ_API_BASE: &str = "https://api.linqapp.com/api/partner/v3";
|
||||
|
||||
impl LinqChannel {
|
||||
pub fn new(api_token: String, from_phone: String, allowed_senders: Vec<String>) -> Self {
|
||||
Self {
|
||||
api_token,
|
||||
from_phone,
|
||||
allowed_senders,
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a sender phone number is allowed (E.164 format: +1234567890)
|
||||
fn is_sender_allowed(&self, phone: &str) -> bool {
|
||||
self.allowed_senders.iter().any(|n| n == "*" || n == phone)
|
||||
}
|
||||
|
||||
/// Get the bot's phone number
|
||||
pub fn phone_number(&self) -> &str {
|
||||
&self.from_phone
|
||||
}
|
||||
|
||||
fn media_part_to_image_marker(part: &serde_json::Value) -> Option<String> {
|
||||
let source = part
|
||||
.get("url")
|
||||
.or_else(|| part.get("value"))
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())?;
|
||||
|
||||
let mime_type = part
|
||||
.get("mime_type")
|
||||
.and_then(|value| value.as_str())
|
||||
.map(str::trim)
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase();
|
||||
|
||||
if !mime_type.starts_with("image/") {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(format!("[IMAGE:{source}]"))
|
||||
}
|
||||
|
||||
/// Parse an incoming webhook payload from Linq and extract messages.
|
||||
///
|
||||
/// Linq webhook envelope:
|
||||
/// ```json
|
||||
/// {
|
||||
/// "api_version": "v3",
|
||||
/// "event_type": "message.received",
|
||||
/// "event_id": "...",
|
||||
/// "created_at": "...",
|
||||
/// "trace_id": "...",
|
||||
/// "data": {
|
||||
/// "chat_id": "...",
|
||||
/// "from": "+1...",
|
||||
/// "recipient_phone": "+1...",
|
||||
/// "is_from_me": false,
|
||||
/// "service": "iMessage",
|
||||
/// "message": {
|
||||
/// "id": "...",
|
||||
/// "parts": [{ "type": "text", "value": "..." }]
|
||||
/// }
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
pub fn parse_webhook_payload(&self, payload: &serde_json::Value) -> Vec<ChannelMessage> {
|
||||
let mut messages = Vec::new();
|
||||
|
||||
// Only handle message.received events
|
||||
let event_type = payload
|
||||
.get("event_type")
|
||||
.and_then(|e| e.as_str())
|
||||
.unwrap_or("");
|
||||
if event_type != "message.received" {
|
||||
tracing::debug!("Linq: skipping non-message event: {event_type}");
|
||||
return messages;
|
||||
}
|
||||
|
||||
let Some(data) = payload.get("data") else {
|
||||
return messages;
|
||||
};
|
||||
|
||||
// Skip messages sent by the bot itself
|
||||
if data
|
||||
.get("is_from_me")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false)
|
||||
{
|
||||
tracing::debug!("Linq: skipping is_from_me message");
|
||||
return messages;
|
||||
}
|
||||
|
||||
// Get sender phone number
|
||||
let Some(from) = data.get("from").and_then(|f| f.as_str()) else {
|
||||
return messages;
|
||||
};
|
||||
|
||||
// Normalize to E.164 format
|
||||
let normalized_from = if from.starts_with('+') {
|
||||
from.to_string()
|
||||
} else {
|
||||
format!("+{from}")
|
||||
};
|
||||
|
||||
// Check allowlist
|
||||
if !self.is_sender_allowed(&normalized_from) {
|
||||
tracing::warn!(
|
||||
"Linq: ignoring message from unauthorized sender: {normalized_from}. \
|
||||
Add to allowed_senders in config.toml."
|
||||
);
|
||||
return messages;
|
||||
}
|
||||
|
||||
// Get chat_id for reply routing
|
||||
let chat_id = data
|
||||
.get("chat_id")
|
||||
.and_then(|c| c.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
// Extract text from message parts
|
||||
let Some(message) = data.get("message") else {
|
||||
return messages;
|
||||
};
|
||||
|
||||
let Some(parts) = message.get("parts").and_then(|p| p.as_array()) else {
|
||||
return messages;
|
||||
};
|
||||
|
||||
let content_parts: Vec<String> = parts
|
||||
.iter()
|
||||
.filter_map(|part| {
|
||||
let part_type = part.get("type").and_then(|t| t.as_str())?;
|
||||
match part_type {
|
||||
"text" => part
|
||||
.get("value")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(ToString::to_string),
|
||||
"media" | "image" => {
|
||||
if let Some(marker) = Self::media_part_to_image_marker(part) {
|
||||
Some(marker)
|
||||
} else {
|
||||
tracing::debug!("Linq: skipping unsupported {part_type} part");
|
||||
None
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
tracing::debug!("Linq: skipping {part_type} part");
|
||||
None
|
||||
}
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
if content_parts.is_empty() {
|
||||
return messages;
|
||||
}
|
||||
|
||||
let content = content_parts.join("\n").trim().to_string();
|
||||
|
||||
if content.is_empty() {
|
||||
return messages;
|
||||
}
|
||||
|
||||
// Get timestamp from created_at or use current time
|
||||
let timestamp = payload
|
||||
.get("created_at")
|
||||
.and_then(|t| t.as_str())
|
||||
.and_then(|t| {
|
||||
chrono::DateTime::parse_from_rfc3339(t)
|
||||
.ok()
|
||||
.map(|dt| dt.timestamp().cast_unsigned())
|
||||
})
|
||||
.unwrap_or_else(|| {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
});
|
||||
|
||||
// Use chat_id as reply_target so replies go to the right conversation
|
||||
let reply_target = if chat_id.is_empty() {
|
||||
normalized_from.clone()
|
||||
} else {
|
||||
chat_id
|
||||
};
|
||||
|
||||
messages.push(ChannelMessage {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
reply_target,
|
||||
sender: normalized_from,
|
||||
content,
|
||||
channel: "linq".to_string(),
|
||||
timestamp,
|
||||
thread_ts: None,
|
||||
});
|
||||
|
||||
messages
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for LinqChannel {
|
||||
fn name(&self) -> &str {
|
||||
"linq"
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> anyhow::Result<()> {
|
||||
// If reply_target looks like a chat_id, send to existing chat.
|
||||
// Otherwise create a new chat with the recipient phone number.
|
||||
let recipient = &message.recipient;
|
||||
|
||||
let body = serde_json::json!({
|
||||
"message": {
|
||||
"parts": [{
|
||||
"type": "text",
|
||||
"value": message.content
|
||||
}]
|
||||
}
|
||||
});
|
||||
|
||||
// Try sending to existing chat (recipient is chat_id)
|
||||
let url = format!("{LINQ_API_BASE}/chats/{recipient}/messages");
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(&url)
|
||||
.bearer_auth(&self.api_token)
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if resp.status().is_success() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// If the chat_id-based send failed with 404, try creating a new chat
|
||||
if resp.status() == reqwest::StatusCode::NOT_FOUND {
|
||||
let new_chat_body = serde_json::json!({
|
||||
"from": self.from_phone,
|
||||
"to": [recipient],
|
||||
"message": {
|
||||
"parts": [{
|
||||
"type": "text",
|
||||
"value": message.content
|
||||
}]
|
||||
}
|
||||
});
|
||||
|
||||
let create_resp = self
|
||||
.client
|
||||
.post(format!("{LINQ_API_BASE}/chats"))
|
||||
.bearer_auth(&self.api_token)
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&new_chat_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !create_resp.status().is_success() {
|
||||
let status = create_resp.status();
|
||||
let error_body = create_resp.text().await.unwrap_or_default();
|
||||
tracing::error!("Linq create chat failed: {status} — {error_body}");
|
||||
anyhow::bail!("Linq API error: {status}");
|
||||
}
|
||||
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let status = resp.status();
|
||||
let error_body = resp.text().await.unwrap_or_default();
|
||||
tracing::error!("Linq send failed: {status} — {error_body}");
|
||||
anyhow::bail!("Linq API error: {status}");
|
||||
}
|
||||
|
||||
async fn listen(&self, _tx: tokio::sync::mpsc::Sender<ChannelMessage>) -> anyhow::Result<()> {
|
||||
// Linq uses webhooks (push-based), not polling.
|
||||
tracing::info!(
|
||||
"Linq channel active (webhook mode). \
|
||||
Configure Linq to POST webhook events to your deployed HTTPS webhook URL."
|
||||
);
|
||||
|
||||
// Keep the task alive — it will be cancelled when the channel shuts down
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(3600)).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
// Check if we can reach the Linq API
|
||||
let url = format!("{LINQ_API_BASE}/phonenumbers");
|
||||
|
||||
self.client
|
||||
.get(&url)
|
||||
.bearer_auth(&self.api_token)
|
||||
.send()
|
||||
.await
|
||||
.map(|r| r.status().is_success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
async fn start_typing(&self, recipient: &str) -> anyhow::Result<()> {
|
||||
let url = format!("{LINQ_API_BASE}/chats/{recipient}/typing");
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(&url)
|
||||
.bearer_auth(&self.api_token)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
tracing::debug!("Linq start_typing failed: {}", resp.status());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stop_typing(&self, recipient: &str) -> anyhow::Result<()> {
|
||||
let url = format!("{LINQ_API_BASE}/chats/{recipient}/typing");
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.delete(&url)
|
||||
.bearer_auth(&self.api_token)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
tracing::debug!("Linq stop_typing failed: {}", resp.status());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Verify a Linq webhook signature.
|
||||
///
|
||||
/// Linq signs webhooks with HMAC-SHA256 over `"{timestamp}.{body}"`.
|
||||
/// The signature is sent in `X-Webhook-Signature` (hex-encoded) and the
|
||||
/// timestamp in `X-Webhook-Timestamp`. Reject timestamps older than 300s.
|
||||
pub fn verify_linq_signature(secret: &str, body: &str, timestamp: &str, signature: &str) -> bool {
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::Sha256;
|
||||
|
||||
// Reject stale timestamps (>300s old)
|
||||
if let Ok(ts) = timestamp.parse::<i64>() {
|
||||
let now = chrono::Utc::now().timestamp();
|
||||
if (now - ts).unsigned_abs() > 300 {
|
||||
tracing::warn!("Linq: rejecting stale webhook timestamp ({ts}, now={now})");
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
tracing::warn!("Linq: invalid webhook timestamp: {timestamp}");
|
||||
return false;
|
||||
}
|
||||
|
||||
// Compute HMAC-SHA256 over "{timestamp}.{body}"
|
||||
let message = format!("{timestamp}.{body}");
|
||||
let Ok(mut mac) = Hmac::<Sha256>::new_from_slice(secret.as_bytes()) else {
|
||||
return false;
|
||||
};
|
||||
mac.update(message.as_bytes());
|
||||
let signature_hex = signature
|
||||
.trim()
|
||||
.strip_prefix("sha256=")
|
||||
.unwrap_or(signature);
|
||||
let Ok(provided) = hex::decode(signature_hex.trim()) else {
|
||||
tracing::warn!("Linq: invalid webhook signature format");
|
||||
return false;
|
||||
};
|
||||
|
||||
// Constant-time comparison via HMAC verify.
|
||||
mac.verify_slice(&provided).is_ok()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "linq_tests.rs"]
|
||||
mod tests;
|
||||
pub use tinychannels::providers::linq::*;
|
||||
|
||||
@@ -1,390 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
fn make_channel() -> LinqChannel {
|
||||
LinqChannel::new(
|
||||
"test-token".into(),
|
||||
"+15551234567".into(),
|
||||
vec!["+1234567890".into()],
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_channel_name() {
|
||||
let ch = make_channel();
|
||||
assert_eq!(ch.name(), "linq");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_sender_allowed_exact() {
|
||||
let ch = make_channel();
|
||||
assert!(ch.is_sender_allowed("+1234567890"));
|
||||
assert!(!ch.is_sender_allowed("+9876543210"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_sender_allowed_wildcard() {
|
||||
let ch = LinqChannel::new("tok".into(), "+15551234567".into(), vec!["*".into()]);
|
||||
assert!(ch.is_sender_allowed("+1234567890"));
|
||||
assert!(ch.is_sender_allowed("+9999999999"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_sender_allowed_empty() {
|
||||
let ch = LinqChannel::new("tok".into(), "+15551234567".into(), vec![]);
|
||||
assert!(!ch.is_sender_allowed("+1234567890"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_valid_text_message() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"api_version": "v3",
|
||||
"event_type": "message.received",
|
||||
"event_id": "evt-123",
|
||||
"created_at": "2025-01-15T12:00:00Z",
|
||||
"trace_id": "trace-456",
|
||||
"data": {
|
||||
"chat_id": "chat-789",
|
||||
"from": "+1234567890",
|
||||
"recipient_phone": "+15551234567",
|
||||
"is_from_me": false,
|
||||
"service": "iMessage",
|
||||
"message": {
|
||||
"id": "msg-abc",
|
||||
"parts": [{
|
||||
"type": "text",
|
||||
"value": "Hello OpenHuman!"
|
||||
}]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(msgs[0].sender, "+1234567890");
|
||||
assert_eq!(msgs[0].content, "Hello OpenHuman!");
|
||||
assert_eq!(msgs[0].channel, "linq");
|
||||
assert_eq!(msgs[0].reply_target, "chat-789");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_skip_is_from_me() {
|
||||
let ch = LinqChannel::new("tok".into(), "+15551234567".into(), vec!["*".into()]);
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.received",
|
||||
"data": {
|
||||
"chat_id": "chat-789",
|
||||
"from": "+1234567890",
|
||||
"is_from_me": true,
|
||||
"message": {
|
||||
"id": "msg-abc",
|
||||
"parts": [{ "type": "text", "value": "My own message" }]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert!(msgs.is_empty(), "is_from_me messages should be skipped");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_skip_non_message_event() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.delivered",
|
||||
"data": {
|
||||
"chat_id": "chat-789",
|
||||
"message_id": "msg-abc"
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert!(msgs.is_empty(), "Non-message events should be skipped");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_unauthorized_sender() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.received",
|
||||
"data": {
|
||||
"chat_id": "chat-789",
|
||||
"from": "+9999999999",
|
||||
"is_from_me": false,
|
||||
"message": {
|
||||
"id": "msg-abc",
|
||||
"parts": [{ "type": "text", "value": "Spam" }]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert!(msgs.is_empty(), "Unauthorized senders should be filtered");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_empty_payload() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({});
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_media_only_translated_to_image_marker() {
|
||||
let ch = LinqChannel::new("tok".into(), "+15551234567".into(), vec!["*".into()]);
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.received",
|
||||
"data": {
|
||||
"chat_id": "chat-789",
|
||||
"from": "+1234567890",
|
||||
"is_from_me": false,
|
||||
"message": {
|
||||
"id": "msg-abc",
|
||||
"parts": [{
|
||||
"type": "media",
|
||||
"url": "https://example.com/image.jpg",
|
||||
"mime_type": "image/jpeg"
|
||||
}]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(msgs[0].content, "[IMAGE:https://example.com/image.jpg]");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_media_non_image_still_skipped() {
|
||||
let ch = LinqChannel::new("tok".into(), "+15551234567".into(), vec!["*".into()]);
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.received",
|
||||
"data": {
|
||||
"chat_id": "chat-789",
|
||||
"from": "+1234567890",
|
||||
"is_from_me": false,
|
||||
"message": {
|
||||
"id": "msg-abc",
|
||||
"parts": [{
|
||||
"type": "media",
|
||||
"url": "https://example.com/sound.mp3",
|
||||
"mime_type": "audio/mpeg"
|
||||
}]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert!(msgs.is_empty(), "Non-image media should still be skipped");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_multiple_text_parts() {
|
||||
let ch = LinqChannel::new("tok".into(), "+15551234567".into(), vec!["*".into()]);
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.received",
|
||||
"data": {
|
||||
"chat_id": "chat-789",
|
||||
"from": "+1234567890",
|
||||
"is_from_me": false,
|
||||
"message": {
|
||||
"id": "msg-abc",
|
||||
"parts": [
|
||||
{ "type": "text", "value": "First part" },
|
||||
{ "type": "text", "value": "Second part" }
|
||||
]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(msgs[0].content, "First part\nSecond part");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_signature_verification_valid() {
|
||||
let secret = "test_webhook_secret";
|
||||
let body = r#"{"event_type":"message.received"}"#;
|
||||
let now = chrono::Utc::now().timestamp().to_string();
|
||||
|
||||
// Compute expected signature
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::Sha256;
|
||||
let message = format!("{now}.{body}");
|
||||
let mut mac = Hmac::<Sha256>::new_from_slice(secret.as_bytes()).unwrap();
|
||||
mac.update(message.as_bytes());
|
||||
let signature = hex::encode(mac.finalize().into_bytes());
|
||||
|
||||
assert!(verify_linq_signature(secret, body, &now, &signature));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_signature_verification_invalid() {
|
||||
let secret = "test_webhook_secret";
|
||||
let body = r#"{"event_type":"message.received"}"#;
|
||||
let now = chrono::Utc::now().timestamp().to_string();
|
||||
|
||||
assert!(!verify_linq_signature(
|
||||
secret,
|
||||
body,
|
||||
&now,
|
||||
"deadbeefdeadbeefdeadbeef"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_signature_verification_stale_timestamp() {
|
||||
let secret = "test_webhook_secret";
|
||||
let body = r#"{"event_type":"message.received"}"#;
|
||||
// 10 minutes ago — stale
|
||||
let stale_ts = (chrono::Utc::now().timestamp() - 600).to_string();
|
||||
|
||||
// Even with correct signature, stale timestamp should fail
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::Sha256;
|
||||
let message = format!("{stale_ts}.{body}");
|
||||
let mut mac = Hmac::<Sha256>::new_from_slice(secret.as_bytes()).unwrap();
|
||||
mac.update(message.as_bytes());
|
||||
let signature = hex::encode(mac.finalize().into_bytes());
|
||||
|
||||
assert!(
|
||||
!verify_linq_signature(secret, body, &stale_ts, &signature),
|
||||
"Stale timestamps (>300s) should be rejected"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_signature_verification_accepts_sha256_prefix() {
|
||||
let secret = "test_webhook_secret";
|
||||
let body = r#"{"event_type":"message.received"}"#;
|
||||
let now = chrono::Utc::now().timestamp().to_string();
|
||||
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::Sha256;
|
||||
let message = format!("{now}.{body}");
|
||||
let mut mac = Hmac::<Sha256>::new_from_slice(secret.as_bytes()).unwrap();
|
||||
mac.update(message.as_bytes());
|
||||
let signature = format!("sha256={}", hex::encode(mac.finalize().into_bytes()));
|
||||
|
||||
assert!(verify_linq_signature(secret, body, &now, &signature));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_signature_verification_accepts_uppercase_hex() {
|
||||
let secret = "test_webhook_secret";
|
||||
let body = r#"{"event_type":"message.received"}"#;
|
||||
let now = chrono::Utc::now().timestamp().to_string();
|
||||
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::Sha256;
|
||||
let message = format!("{now}.{body}");
|
||||
let mut mac = Hmac::<Sha256>::new_from_slice(secret.as_bytes()).unwrap();
|
||||
mac.update(message.as_bytes());
|
||||
let signature = hex::encode(mac.finalize().into_bytes()).to_ascii_uppercase();
|
||||
|
||||
assert!(verify_linq_signature(secret, body, &now, &signature));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_normalizes_phone_with_plus() {
|
||||
let ch = LinqChannel::new(
|
||||
"tok".into(),
|
||||
"+15551234567".into(),
|
||||
vec!["+1234567890".into()],
|
||||
);
|
||||
// API sends without +, normalize to +
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.received",
|
||||
"data": {
|
||||
"chat_id": "chat-789",
|
||||
"from": "1234567890",
|
||||
"is_from_me": false,
|
||||
"message": {
|
||||
"id": "msg-abc",
|
||||
"parts": [{ "type": "text", "value": "Hi" }]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(msgs[0].sender, "+1234567890");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_missing_data() {
|
||||
let ch = make_channel();
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.received"
|
||||
});
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_missing_message_parts() {
|
||||
let ch = LinqChannel::new("tok".into(), "+15551234567".into(), vec!["*".into()]);
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.received",
|
||||
"data": {
|
||||
"chat_id": "chat-789",
|
||||
"from": "+1234567890",
|
||||
"is_from_me": false,
|
||||
"message": {
|
||||
"id": "msg-abc"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert!(msgs.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_empty_text_value() {
|
||||
let ch = LinqChannel::new("tok".into(), "+15551234567".into(), vec!["*".into()]);
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.received",
|
||||
"data": {
|
||||
"chat_id": "chat-789",
|
||||
"from": "+1234567890",
|
||||
"is_from_me": false,
|
||||
"message": {
|
||||
"id": "msg-abc",
|
||||
"parts": [{ "type": "text", "value": "" }]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert!(msgs.is_empty(), "Empty text should be skipped");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_parse_fallback_reply_target_when_no_chat_id() {
|
||||
let ch = LinqChannel::new("tok".into(), "+15551234567".into(), vec!["*".into()]);
|
||||
let payload = serde_json::json!({
|
||||
"event_type": "message.received",
|
||||
"data": {
|
||||
"from": "+1234567890",
|
||||
"is_from_me": false,
|
||||
"message": {
|
||||
"id": "msg-abc",
|
||||
"parts": [{ "type": "text", "value": "Hi" }]
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let msgs = ch.parse_webhook_payload(&payload);
|
||||
assert_eq!(msgs.len(), 1);
|
||||
// Falls back to sender phone number when no chat_id
|
||||
assert_eq!(msgs[0].reply_target, "+1234567890");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn linq_phone_number_accessor() {
|
||||
let ch = make_channel();
|
||||
assert_eq!(ch.phone_number(), "+15551234567");
|
||||
}
|
||||
@@ -1,490 +1 @@
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
use anyhow::{bail, Result};
|
||||
use async_trait::async_trait;
|
||||
use parking_lot::Mutex;
|
||||
|
||||
/// Mattermost channel — polls channel posts via REST API v4.
|
||||
/// Mattermost is API-compatible with many Slack patterns but uses a dedicated v4 structure.
|
||||
pub struct MattermostChannel {
|
||||
base_url: String, // e.g., https://mm.example.com
|
||||
bot_token: String,
|
||||
channel_id: Option<String>,
|
||||
allowed_users: Vec<String>,
|
||||
/// When true (default), replies thread on the original post's root_id.
|
||||
/// When false, replies go to the channel root.
|
||||
thread_replies: bool,
|
||||
/// When true, only respond to messages that @-mention the bot.
|
||||
mention_only: bool,
|
||||
/// Handle for the background typing-indicator loop (aborted on stop_typing).
|
||||
typing_handle: Mutex<Option<tokio::task::JoinHandle<()>>>,
|
||||
}
|
||||
|
||||
impl MattermostChannel {
|
||||
pub fn new(
|
||||
base_url: String,
|
||||
bot_token: String,
|
||||
channel_id: Option<String>,
|
||||
allowed_users: Vec<String>,
|
||||
thread_replies: bool,
|
||||
mention_only: bool,
|
||||
) -> Self {
|
||||
// Ensure base_url doesn't have a trailing slash for consistent path joining
|
||||
let base_url = base_url.trim_end_matches('/').to_string();
|
||||
Self {
|
||||
base_url,
|
||||
bot_token,
|
||||
channel_id,
|
||||
allowed_users,
|
||||
thread_replies,
|
||||
mention_only,
|
||||
typing_handle: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn http_client(&self) -> reqwest::Client {
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.mattermost")
|
||||
}
|
||||
|
||||
/// Check if a user ID is in the allowlist.
|
||||
/// Empty list means deny everyone. "*" means allow everyone.
|
||||
fn is_user_allowed(&self, user_id: &str) -> bool {
|
||||
self.allowed_users.iter().any(|u| u == "*" || u == user_id)
|
||||
}
|
||||
|
||||
/// Get the bot's own user ID and username so we can ignore our own messages
|
||||
/// and detect @-mentions by username.
|
||||
async fn get_bot_identity(&self) -> (String, String) {
|
||||
let resp: Option<serde_json::Value> = async {
|
||||
self.http_client()
|
||||
.get(format!("{}/api/v4/users/me", self.base_url))
|
||||
.bearer_auth(&self.bot_token)
|
||||
.send()
|
||||
.await
|
||||
.ok()?
|
||||
.json()
|
||||
.await
|
||||
.ok()
|
||||
}
|
||||
.await;
|
||||
|
||||
let id = resp
|
||||
.as_ref()
|
||||
.and_then(|v| v.get("id"))
|
||||
.and_then(|u| u.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let username = resp
|
||||
.as_ref()
|
||||
.and_then(|v| v.get("username"))
|
||||
.and_then(|u| u.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
(id, username)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for MattermostChannel {
|
||||
fn name(&self) -> &str {
|
||||
"mattermost"
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> Result<()> {
|
||||
// Mattermost supports threading via 'root_id'.
|
||||
// We pack 'channel_id:root_id' into recipient if it's a thread.
|
||||
let (channel_id, root_id) = if let Some((c, r)) = message.recipient.split_once(':') {
|
||||
(c, Some(r))
|
||||
} else {
|
||||
(message.recipient.as_str(), None)
|
||||
};
|
||||
|
||||
let mut body_map = serde_json::json!({
|
||||
"channel_id": channel_id,
|
||||
"message": message.content
|
||||
});
|
||||
|
||||
if let Some(root) = root_id {
|
||||
body_map.as_object_mut().unwrap().insert(
|
||||
"root_id".to_string(),
|
||||
serde_json::Value::String(root.to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(format!("{}/api/v4/posts", self.base_url))
|
||||
.bearer_auth(&self.bot_token)
|
||||
.json(&body_map)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
let body = resp
|
||||
.text()
|
||||
.await
|
||||
.unwrap_or_else(|e| format!("<failed to read response: {e}>"));
|
||||
bail!("Mattermost post failed ({status}): {body}");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn listen(&self, tx: tokio::sync::mpsc::Sender<ChannelMessage>) -> Result<()> {
|
||||
let channel_id = self
|
||||
.channel_id
|
||||
.clone()
|
||||
.ok_or_else(|| anyhow::anyhow!("Mattermost channel_id required for listening"))?;
|
||||
|
||||
let (bot_user_id, bot_username) = self.get_bot_identity().await;
|
||||
#[allow(clippy::cast_possible_truncation)]
|
||||
let mut last_create_at = (std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_millis()) as i64;
|
||||
|
||||
tracing::info!("Mattermost channel listening on {}...", channel_id);
|
||||
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_secs(3)).await;
|
||||
|
||||
let resp = match self
|
||||
.http_client()
|
||||
.get(format!(
|
||||
"{}/api/v4/channels/{}/posts",
|
||||
self.base_url, channel_id
|
||||
))
|
||||
.bearer_auth(&self.bot_token)
|
||||
.query(&[("since", last_create_at.to_string())])
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::warn!("Mattermost poll error: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let data: serde_json::Value = match resp.json().await {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
tracing::warn!("Mattermost parse error: {e}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(posts) = data.get("posts").and_then(|p| p.as_object()) {
|
||||
// Process in chronological order
|
||||
let mut post_list: Vec<_> = posts.values().collect();
|
||||
post_list.sort_by_key(|p| p.get("create_at").and_then(|c| c.as_i64()).unwrap_or(0));
|
||||
|
||||
for post in post_list {
|
||||
let msg = self.parse_mattermost_post(
|
||||
post,
|
||||
&bot_user_id,
|
||||
&bot_username,
|
||||
last_create_at,
|
||||
&channel_id,
|
||||
);
|
||||
let create_at = post
|
||||
.get("create_at")
|
||||
.and_then(|c| c.as_i64())
|
||||
.unwrap_or(last_create_at);
|
||||
last_create_at = last_create_at.max(create_at);
|
||||
|
||||
if let Some(channel_msg) = msg {
|
||||
if tx.send(channel_msg).await.is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
self.http_client()
|
||||
.get(format!("{}/api/v4/users/me", self.base_url))
|
||||
.bearer_auth(&self.bot_token)
|
||||
.send()
|
||||
.await
|
||||
.map(|r| r.status().is_success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
async fn start_typing(&self, recipient: &str) -> Result<()> {
|
||||
// Cancel any existing typing loop before starting a new one.
|
||||
self.stop_typing(recipient).await?;
|
||||
|
||||
let client = self.http_client();
|
||||
let token = self.bot_token.clone();
|
||||
let base_url = self.base_url.clone();
|
||||
|
||||
// recipient is "channel_id" or "channel_id:root_id"
|
||||
let (channel_id, parent_id) = match recipient.split_once(':') {
|
||||
Some((channel, parent)) => (channel.to_string(), Some(parent.to_string())),
|
||||
None => (recipient.to_string(), None),
|
||||
};
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
let url = format!("{base_url}/api/v4/users/me/typing");
|
||||
loop {
|
||||
let mut body = serde_json::json!({ "channel_id": channel_id });
|
||||
if let Some(ref pid) = parent_id {
|
||||
body.as_object_mut()
|
||||
.unwrap()
|
||||
.insert("parent_id".to_string(), serde_json::json!(pid));
|
||||
}
|
||||
|
||||
if let Ok(r) = client
|
||||
.post(&url)
|
||||
.bearer_auth(&token)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
if !r.status().is_success() {
|
||||
tracing::debug!(status = %r.status(), "Mattermost typing indicator failed");
|
||||
}
|
||||
}
|
||||
|
||||
// Mattermost typing events expire after ~6s; re-fire every 4s.
|
||||
tokio::time::sleep(std::time::Duration::from_secs(4)).await;
|
||||
}
|
||||
});
|
||||
|
||||
let mut guard = self.typing_handle.lock();
|
||||
*guard = Some(handle);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stop_typing(&self, _recipient: &str) -> Result<()> {
|
||||
let mut guard = self.typing_handle.lock();
|
||||
if let Some(handle) = guard.take() {
|
||||
handle.abort();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl MattermostChannel {
|
||||
fn parse_mattermost_post(
|
||||
&self,
|
||||
post: &serde_json::Value,
|
||||
bot_user_id: &str,
|
||||
bot_username: &str,
|
||||
last_create_at: i64,
|
||||
channel_id: &str,
|
||||
) -> Option<ChannelMessage> {
|
||||
let id = post.get("id").and_then(|i| i.as_str()).unwrap_or("");
|
||||
let user_id = post.get("user_id").and_then(|u| u.as_str()).unwrap_or("");
|
||||
let text = post.get("message").and_then(|m| m.as_str()).unwrap_or("");
|
||||
let create_at = post.get("create_at").and_then(|c| c.as_i64()).unwrap_or(0);
|
||||
let root_id = post.get("root_id").and_then(|r| r.as_str()).unwrap_or("");
|
||||
|
||||
if user_id == bot_user_id || create_at <= last_create_at || text.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
if !self.is_user_allowed(user_id) {
|
||||
tracing::warn!("Mattermost: ignoring message from unauthorized user: {user_id}");
|
||||
return None;
|
||||
}
|
||||
|
||||
// mention_only filtering: skip messages that don't @-mention the bot.
|
||||
let content = if self.mention_only {
|
||||
let normalized = normalize_mattermost_content(text, bot_user_id, bot_username, post);
|
||||
normalized?
|
||||
} else {
|
||||
text.to_string()
|
||||
};
|
||||
|
||||
// Reply routing depends on thread_replies config:
|
||||
// - Existing thread (root_id set): always stay in the thread.
|
||||
// - Top-level post + thread_replies=true: thread on the original post.
|
||||
// - Top-level post + thread_replies=false: reply at channel level.
|
||||
let reply_target = if !root_id.is_empty() {
|
||||
format!("{}:{}", channel_id, root_id)
|
||||
} else if self.thread_replies {
|
||||
format!("{}:{}", channel_id, id)
|
||||
} else {
|
||||
channel_id.to_string()
|
||||
};
|
||||
|
||||
Some(ChannelMessage {
|
||||
id: format!("mattermost_{id}"),
|
||||
sender: user_id.to_string(),
|
||||
reply_target,
|
||||
content,
|
||||
channel: "mattermost".to_string(),
|
||||
#[allow(clippy::cast_sign_loss)]
|
||||
timestamp: (create_at / 1000) as u64,
|
||||
thread_ts: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Check whether a Mattermost post contains an @-mention of the bot.
|
||||
///
|
||||
/// Checks two sources:
|
||||
/// 1. Text-based: looks for `@bot_username` in the message body (case-insensitive).
|
||||
/// 2. Metadata-based: checks the post's `metadata.mentions` array for the bot user ID.
|
||||
fn contains_bot_mention_mm(
|
||||
text: &str,
|
||||
bot_user_id: &str,
|
||||
bot_username: &str,
|
||||
post: &serde_json::Value,
|
||||
) -> bool {
|
||||
// 1. Text-based: @username (case-insensitive, word-boundary aware)
|
||||
if !find_bot_mention_spans(text, bot_username).is_empty() {
|
||||
return true;
|
||||
}
|
||||
|
||||
// 2. Metadata-based: Mattermost may include a "metadata.mentions" array of user IDs.
|
||||
if !bot_user_id.is_empty() {
|
||||
if let Some(mentions) = post
|
||||
.get("metadata")
|
||||
.and_then(|m| m.get("mentions"))
|
||||
.and_then(|m| m.as_array())
|
||||
{
|
||||
if mentions.iter().any(|m| m.as_str() == Some(bot_user_id)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
fn is_mattermost_username_char(c: char) -> bool {
|
||||
c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.'
|
||||
}
|
||||
|
||||
fn find_bot_mention_spans(text: &str, bot_username: &str) -> Vec<(usize, usize)> {
|
||||
if bot_username.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mention = format!("@{}", bot_username.to_ascii_lowercase());
|
||||
let mention_len = mention.len();
|
||||
if mention_len == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mention_bytes = mention.as_bytes();
|
||||
let text_bytes = text.as_bytes();
|
||||
let mut spans = Vec::new();
|
||||
let mut index = 0;
|
||||
|
||||
while index + mention_len <= text_bytes.len() {
|
||||
let is_match = text_bytes[index] == b'@'
|
||||
&& text_bytes[index..index + mention_len]
|
||||
.iter()
|
||||
.zip(mention_bytes.iter())
|
||||
.all(|(left, right)| left.eq_ignore_ascii_case(right));
|
||||
|
||||
if is_match {
|
||||
let end = index + mention_len;
|
||||
let at_boundary = text[end..]
|
||||
.chars()
|
||||
.next()
|
||||
.is_none_or(|next| !is_mattermost_username_char(next));
|
||||
if at_boundary {
|
||||
spans.push((index, end));
|
||||
index = end;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
let step = text[index..].chars().next().map_or(1, char::len_utf8);
|
||||
index += step;
|
||||
}
|
||||
|
||||
spans
|
||||
}
|
||||
|
||||
/// Normalize incoming Mattermost content when `mention_only` is enabled.
|
||||
///
|
||||
/// Returns `None` if the message doesn't mention the bot.
|
||||
/// Returns `Some(cleaned)` with the @-mention stripped and text trimmed.
|
||||
fn normalize_mattermost_content(
|
||||
text: &str,
|
||||
bot_user_id: &str,
|
||||
bot_username: &str,
|
||||
post: &serde_json::Value,
|
||||
) -> Option<String> {
|
||||
let mention_spans = find_bot_mention_spans(text, bot_username);
|
||||
let metadata_mentions_bot = !bot_user_id.is_empty()
|
||||
&& post
|
||||
.get("metadata")
|
||||
.and_then(|m| m.get("mentions"))
|
||||
.and_then(|m| m.as_array())
|
||||
.is_some_and(|mentions| mentions.iter().any(|m| m.as_str() == Some(bot_user_id)));
|
||||
|
||||
if mention_spans.is_empty() && !metadata_mentions_bot {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut cleaned = text.to_string();
|
||||
if !mention_spans.is_empty() {
|
||||
let mut result = String::with_capacity(text.len());
|
||||
let mut cursor = 0;
|
||||
for (start, end) in mention_spans {
|
||||
result.push_str(&text[cursor..start]);
|
||||
result.push(' ');
|
||||
cursor = end;
|
||||
}
|
||||
result.push_str(&text[cursor..]);
|
||||
cleaned = result;
|
||||
}
|
||||
|
||||
let cleaned = cleaned.trim().to_string();
|
||||
if cleaned.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(cleaned)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "mattermost_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
#[cfg(any(test, debug_assertions))]
|
||||
pub mod test_support {
|
||||
//! Debug-build seams for raw integration tests. These expose Mattermost's
|
||||
//! private parser helpers without widening the production API surface.
|
||||
|
||||
use super::*;
|
||||
|
||||
pub fn parse_mattermost_post_for_test(
|
||||
channel: &MattermostChannel,
|
||||
post: &serde_json::Value,
|
||||
bot_user_id: &str,
|
||||
bot_username: &str,
|
||||
last_create_at: i64,
|
||||
channel_id: &str,
|
||||
) -> Option<ChannelMessage> {
|
||||
channel.parse_mattermost_post(post, bot_user_id, bot_username, last_create_at, channel_id)
|
||||
}
|
||||
|
||||
pub fn contains_bot_mention_for_test(
|
||||
text: &str,
|
||||
bot_user_id: &str,
|
||||
bot_username: &str,
|
||||
post: &serde_json::Value,
|
||||
) -> bool {
|
||||
contains_bot_mention_mm(text, bot_user_id, bot_username, post)
|
||||
}
|
||||
|
||||
pub fn normalize_mattermost_content_for_test(
|
||||
text: &str,
|
||||
bot_user_id: &str,
|
||||
bot_username: &str,
|
||||
post: &serde_json::Value,
|
||||
) -> Option<String> {
|
||||
normalize_mattermost_content(text, bot_user_id, bot_username, post)
|
||||
}
|
||||
}
|
||||
pub use tinychannels::providers::mattermost::*;
|
||||
|
||||
@@ -1,462 +0,0 @@
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
// Helper: create a channel with mention_only=false (legacy behavior).
|
||||
fn make_channel(allowed: Vec<String>, thread_replies: bool) -> MattermostChannel {
|
||||
MattermostChannel::new(
|
||||
"url".into(),
|
||||
"token".into(),
|
||||
None,
|
||||
allowed,
|
||||
thread_replies,
|
||||
false,
|
||||
)
|
||||
}
|
||||
|
||||
// Helper: create a channel with mention_only=true.
|
||||
fn make_mention_only_channel() -> MattermostChannel {
|
||||
MattermostChannel::new(
|
||||
"url".into(),
|
||||
"token".into(),
|
||||
None,
|
||||
vec!["*".into()],
|
||||
true,
|
||||
true,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mattermost_url_trimming() {
|
||||
let ch = MattermostChannel::new(
|
||||
"https://mm.example.com/".into(),
|
||||
"token".into(),
|
||||
None,
|
||||
vec![],
|
||||
false,
|
||||
false,
|
||||
);
|
||||
assert_eq!(ch.base_url, "https://mm.example.com");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mattermost_allowlist_wildcard() {
|
||||
let ch = make_channel(vec!["*".into()], false);
|
||||
assert!(ch.is_user_allowed("any-id"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mattermost_parse_post_basic() {
|
||||
let ch = make_channel(vec!["*".into()], true);
|
||||
let post = json!({
|
||||
"id": "post123",
|
||||
"user_id": "user456",
|
||||
"message": "hello world",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "botname", 1_500_000_000_000_i64, "chan789")
|
||||
.unwrap();
|
||||
assert_eq!(msg.sender, "user456");
|
||||
assert_eq!(msg.content, "hello world");
|
||||
assert_eq!(msg.reply_target, "chan789:post123"); // Default threaded reply
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mattermost_parse_post_thread_replies_enabled() {
|
||||
let ch = make_channel(vec!["*".into()], true);
|
||||
let post = json!({
|
||||
"id": "post123",
|
||||
"user_id": "user456",
|
||||
"message": "hello world",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "botname", 1_500_000_000_000_i64, "chan789")
|
||||
.unwrap();
|
||||
assert_eq!(msg.reply_target, "chan789:post123"); // Threaded reply
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mattermost_parse_post_thread() {
|
||||
let ch = make_channel(vec!["*".into()], false);
|
||||
let post = json!({
|
||||
"id": "post123",
|
||||
"user_id": "user456",
|
||||
"message": "reply",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": "root789"
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "botname", 1_500_000_000_000_i64, "chan789")
|
||||
.unwrap();
|
||||
assert_eq!(msg.reply_target, "chan789:root789"); // Stays in the thread
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mattermost_parse_post_ignore_self() {
|
||||
let ch = make_channel(vec!["*".into()], false);
|
||||
let post = json!({
|
||||
"id": "post123",
|
||||
"user_id": "bot123",
|
||||
"message": "my own message",
|
||||
"create_at": 1_600_000_000_000_i64
|
||||
});
|
||||
|
||||
let msg =
|
||||
ch.parse_mattermost_post(&post, "bot123", "botname", 1_500_000_000_000_i64, "chan789");
|
||||
assert!(msg.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mattermost_parse_post_ignore_old() {
|
||||
let ch = make_channel(vec!["*".into()], false);
|
||||
let post = json!({
|
||||
"id": "post123",
|
||||
"user_id": "user456",
|
||||
"message": "old message",
|
||||
"create_at": 1_400_000_000_000_i64
|
||||
});
|
||||
|
||||
let msg =
|
||||
ch.parse_mattermost_post(&post, "bot123", "botname", 1_500_000_000_000_i64, "chan789");
|
||||
assert!(msg.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mattermost_parse_post_no_thread_when_disabled() {
|
||||
let ch = make_channel(vec!["*".into()], false);
|
||||
let post = json!({
|
||||
"id": "post123",
|
||||
"user_id": "user456",
|
||||
"message": "hello world",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "botname", 1_500_000_000_000_i64, "chan789")
|
||||
.unwrap();
|
||||
assert_eq!(msg.reply_target, "chan789"); // No thread suffix
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mattermost_existing_thread_always_threads() {
|
||||
// Even with thread_replies=false, replies to existing threads stay in the thread
|
||||
let ch = make_channel(vec!["*".into()], false);
|
||||
let post = json!({
|
||||
"id": "post123",
|
||||
"user_id": "user456",
|
||||
"message": "reply in thread",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": "root789"
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "botname", 1_500_000_000_000_i64, "chan789")
|
||||
.unwrap();
|
||||
assert_eq!(msg.reply_target, "chan789:root789"); // Stays in existing thread
|
||||
}
|
||||
|
||||
// ── mention_only tests ────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn mention_only_skips_message_without_mention() {
|
||||
let ch = make_mention_only_channel();
|
||||
let post = json!({
|
||||
"id": "post1",
|
||||
"user_id": "user1",
|
||||
"message": "hello everyone",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch.parse_mattermost_post(&post, "bot123", "mybot", 1_500_000_000_000_i64, "chan1");
|
||||
assert!(msg.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_only_accepts_message_with_at_mention() {
|
||||
let ch = make_mention_only_channel();
|
||||
let post = json!({
|
||||
"id": "post1",
|
||||
"user_id": "user1",
|
||||
"message": "@mybot what is the weather?",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "mybot", 1_500_000_000_000_i64, "chan1")
|
||||
.unwrap();
|
||||
assert_eq!(msg.content, "what is the weather?");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_only_strips_mention_and_trims() {
|
||||
let ch = make_mention_only_channel();
|
||||
let post = json!({
|
||||
"id": "post1",
|
||||
"user_id": "user1",
|
||||
"message": " @mybot run status ",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "mybot", 1_500_000_000_000_i64, "chan1")
|
||||
.unwrap();
|
||||
assert_eq!(msg.content, "run status");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_only_rejects_empty_after_stripping() {
|
||||
let ch = make_mention_only_channel();
|
||||
let post = json!({
|
||||
"id": "post1",
|
||||
"user_id": "user1",
|
||||
"message": "@mybot",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch.parse_mattermost_post(&post, "bot123", "mybot", 1_500_000_000_000_i64, "chan1");
|
||||
assert!(msg.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_only_case_insensitive() {
|
||||
let ch = make_mention_only_channel();
|
||||
let post = json!({
|
||||
"id": "post1",
|
||||
"user_id": "user1",
|
||||
"message": "@MyBot hello",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "mybot", 1_500_000_000_000_i64, "chan1")
|
||||
.unwrap();
|
||||
assert_eq!(msg.content, "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_only_detects_metadata_mentions() {
|
||||
// Even without @username in text, metadata.mentions should trigger.
|
||||
let ch = make_mention_only_channel();
|
||||
let post = json!({
|
||||
"id": "post1",
|
||||
"user_id": "user1",
|
||||
"message": "hey check this out",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": "",
|
||||
"metadata": {
|
||||
"mentions": ["bot123"]
|
||||
}
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "mybot", 1_500_000_000_000_i64, "chan1")
|
||||
.unwrap();
|
||||
// Content is preserved as-is since no @username was in the text to strip.
|
||||
assert_eq!(msg.content, "hey check this out");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_only_word_boundary_prevents_partial_match() {
|
||||
let ch = make_mention_only_channel();
|
||||
// "@mybotextended" should NOT match "@mybot" because it extends the username.
|
||||
let post = json!({
|
||||
"id": "post1",
|
||||
"user_id": "user1",
|
||||
"message": "@mybotextended hello",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch.parse_mattermost_post(&post, "bot123", "mybot", 1_500_000_000_000_i64, "chan1");
|
||||
assert!(msg.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_only_mention_in_middle_of_text() {
|
||||
let ch = make_mention_only_channel();
|
||||
let post = json!({
|
||||
"id": "post1",
|
||||
"user_id": "user1",
|
||||
"message": "hey @mybot how are you?",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "mybot", 1_500_000_000_000_i64, "chan1")
|
||||
.unwrap();
|
||||
assert_eq!(msg.content, "hey how are you?");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_only_disabled_passes_all_messages() {
|
||||
// With mention_only=false (default), messages pass through unfiltered.
|
||||
let ch = make_channel(vec!["*".into()], true);
|
||||
let post = json!({
|
||||
"id": "post1",
|
||||
"user_id": "user1",
|
||||
"message": "no mention here",
|
||||
"create_at": 1_600_000_000_000_i64,
|
||||
"root_id": ""
|
||||
});
|
||||
|
||||
let msg = ch
|
||||
.parse_mattermost_post(&post, "bot123", "mybot", 1_500_000_000_000_i64, "chan1")
|
||||
.unwrap();
|
||||
assert_eq!(msg.content, "no mention here");
|
||||
}
|
||||
|
||||
// ── contains_bot_mention_mm unit tests ────────────────────────
|
||||
|
||||
#[test]
|
||||
fn contains_mention_text_at_end() {
|
||||
let post = json!({});
|
||||
assert!(contains_bot_mention_mm(
|
||||
"hello @mybot",
|
||||
"bot123",
|
||||
"mybot",
|
||||
&post
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn contains_mention_text_at_start() {
|
||||
let post = json!({});
|
||||
assert!(contains_bot_mention_mm(
|
||||
"@mybot hello",
|
||||
"bot123",
|
||||
"mybot",
|
||||
&post
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn contains_mention_text_alone() {
|
||||
let post = json!({});
|
||||
assert!(contains_bot_mention_mm("@mybot", "bot123", "mybot", &post));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_mention_different_username() {
|
||||
let post = json!({});
|
||||
assert!(!contains_bot_mention_mm(
|
||||
"@otherbot hello",
|
||||
"bot123",
|
||||
"mybot",
|
||||
&post
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_mention_partial_username() {
|
||||
let post = json!({});
|
||||
// "mybot" is a prefix of "mybotx" — should NOT match
|
||||
assert!(!contains_bot_mention_mm(
|
||||
"@mybotx hello",
|
||||
"bot123",
|
||||
"mybot",
|
||||
&post
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_detects_later_valid_mention_after_partial_prefix() {
|
||||
let post = json!({});
|
||||
assert!(contains_bot_mention_mm(
|
||||
"@mybotx ignore this, but @mybot handle this",
|
||||
"bot123",
|
||||
"mybot",
|
||||
&post
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_followed_by_punctuation() {
|
||||
let post = json!({});
|
||||
// "@mybot," — comma is not alphanumeric/underscore/dash/dot, so it's a boundary
|
||||
assert!(contains_bot_mention_mm(
|
||||
"@mybot, hello",
|
||||
"bot123",
|
||||
"mybot",
|
||||
&post
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mention_via_metadata_only() {
|
||||
let post = json!({
|
||||
"metadata": { "mentions": ["bot123"] }
|
||||
});
|
||||
assert!(contains_bot_mention_mm(
|
||||
"no at mention",
|
||||
"bot123",
|
||||
"mybot",
|
||||
&post
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_mention_empty_username_no_metadata() {
|
||||
let post = json!({});
|
||||
assert!(!contains_bot_mention_mm("hello world", "bot123", "", &post));
|
||||
}
|
||||
|
||||
// ── normalize_mattermost_content unit tests ───────────────────
|
||||
|
||||
#[test]
|
||||
fn normalize_strips_and_trims() {
|
||||
let post = json!({});
|
||||
let result = normalize_mattermost_content(" @mybot do stuff ", "bot123", "mybot", &post);
|
||||
assert_eq!(result.as_deref(), Some("do stuff"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_returns_none_for_no_mention() {
|
||||
let post = json!({});
|
||||
let result = normalize_mattermost_content("hello world", "bot123", "mybot", &post);
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_returns_none_when_only_mention() {
|
||||
let post = json!({});
|
||||
let result = normalize_mattermost_content("@mybot", "bot123", "mybot", &post);
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_preserves_text_for_metadata_mention() {
|
||||
let post = json!({
|
||||
"metadata": { "mentions": ["bot123"] }
|
||||
});
|
||||
let result = normalize_mattermost_content("check this out", "bot123", "mybot", &post);
|
||||
assert_eq!(result.as_deref(), Some("check this out"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_strips_multiple_mentions() {
|
||||
let post = json!({});
|
||||
let result =
|
||||
normalize_mattermost_content("@mybot hello @mybot world", "bot123", "mybot", &post);
|
||||
assert_eq!(result.as_deref(), Some("hello world"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_keeps_partial_username_mentions() {
|
||||
let post = json!({});
|
||||
let result =
|
||||
normalize_mattermost_content("@mybot hello @mybotx world", "bot123", "mybot", &post);
|
||||
assert_eq!(result.as_deref(), Some("hello @mybotx world"));
|
||||
}
|
||||
@@ -8,12 +8,6 @@ pub mod irc;
|
||||
pub mod lark;
|
||||
pub mod linq;
|
||||
pub mod mattermost;
|
||||
// Public (like every sibling provider module) so cross-module callers reach it
|
||||
// in *all* profiles. It was previously `pub` only under test/debug and private
|
||||
// in release, which compiled in debug but broke the release build once a
|
||||
// cross-module caller appeared (`agent::task_dispatcher` →
|
||||
// `presentation::deliver_response`): release-only `E0603: module is private`.
|
||||
pub mod presentation;
|
||||
pub mod qq;
|
||||
pub mod signal;
|
||||
pub mod slack;
|
||||
|
||||
@@ -1,451 +1 @@
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
use async_trait::async_trait;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use serde_json::json;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use uuid::Uuid;
|
||||
|
||||
const QQ_API_BASE: &str = "https://api.sgroup.qq.com";
|
||||
const QQ_AUTH_URL: &str = "https://bots.qq.com/app/getAppAccessToken";
|
||||
|
||||
fn ensure_https(url: &str) -> anyhow::Result<()> {
|
||||
if !url.starts_with("https://") {
|
||||
anyhow::bail!(
|
||||
"Refusing to transmit sensitive data over non-HTTPS URL: URL scheme must be https"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Deduplication set capacity — evict half of entries when full.
|
||||
const DEDUP_CAPACITY: usize = 10_000;
|
||||
|
||||
/// QQ Official Bot channel — uses Tencent's official QQ Bot API with
|
||||
/// OAuth2 authentication and a Discord-like WebSocket gateway protocol.
|
||||
pub struct QQChannel {
|
||||
app_id: String,
|
||||
app_secret: String,
|
||||
allowed_users: Vec<String>,
|
||||
/// Cached access token + expiry timestamp.
|
||||
token_cache: Arc<RwLock<Option<(String, u64)>>>,
|
||||
/// Message deduplication set.
|
||||
dedup: Arc<RwLock<HashSet<String>>>,
|
||||
}
|
||||
|
||||
impl QQChannel {
|
||||
pub fn new(app_id: String, app_secret: String, allowed_users: Vec<String>) -> Self {
|
||||
Self {
|
||||
app_id,
|
||||
app_secret,
|
||||
allowed_users,
|
||||
token_cache: Arc::new(RwLock::new(None)),
|
||||
dedup: Arc::new(RwLock::new(HashSet::new())),
|
||||
}
|
||||
}
|
||||
|
||||
fn http_client(&self) -> reqwest::Client {
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.qq")
|
||||
}
|
||||
|
||||
fn is_user_allowed(&self, user_id: &str) -> bool {
|
||||
self.allowed_users.iter().any(|u| u == "*" || u == user_id)
|
||||
}
|
||||
|
||||
/// Fetch an access token from QQ's OAuth2 endpoint.
|
||||
async fn fetch_access_token(&self) -> anyhow::Result<(String, u64)> {
|
||||
let body = json!({
|
||||
"appId": self.app_id,
|
||||
"clientSecret": self.app_secret,
|
||||
});
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(QQ_AUTH_URL)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let err = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!("QQ token request failed ({status}): {err}");
|
||||
}
|
||||
|
||||
let data: serde_json::Value = resp.json().await?;
|
||||
let token = data
|
||||
.get("access_token")
|
||||
.and_then(|t| t.as_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("Missing access_token in QQ response"))?
|
||||
.to_string();
|
||||
|
||||
let expires_in = data
|
||||
.get("expires_in")
|
||||
.and_then(|e| e.as_str())
|
||||
.and_then(|e| e.parse::<u64>().ok())
|
||||
.unwrap_or(7200);
|
||||
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
// Expire 60 seconds early to avoid edge cases
|
||||
let expiry = now + expires_in.saturating_sub(60);
|
||||
|
||||
Ok((token, expiry))
|
||||
}
|
||||
|
||||
/// Get a valid access token, refreshing if expired.
|
||||
async fn get_token(&self) -> anyhow::Result<String> {
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
{
|
||||
let cache = self.token_cache.read().await;
|
||||
if let Some((ref token, expiry)) = *cache {
|
||||
if now < expiry {
|
||||
return Ok(token.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let (token, expiry) = self.fetch_access_token().await?;
|
||||
{
|
||||
let mut cache = self.token_cache.write().await;
|
||||
*cache = Some((token.clone(), expiry));
|
||||
}
|
||||
Ok(token)
|
||||
}
|
||||
|
||||
/// Get the WebSocket gateway URL.
|
||||
async fn get_gateway_url(&self, token: &str) -> anyhow::Result<String> {
|
||||
let resp = self
|
||||
.http_client()
|
||||
.get(format!("{QQ_API_BASE}/gateway"))
|
||||
.header("Authorization", format!("QQBot {token}"))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let err = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!("QQ gateway request failed ({status}): {err}");
|
||||
}
|
||||
|
||||
let data: serde_json::Value = resp.json().await?;
|
||||
let url = data
|
||||
.get("url")
|
||||
.and_then(|u| u.as_str())
|
||||
.ok_or_else(|| anyhow::anyhow!("Missing gateway URL in QQ response"))?
|
||||
.to_string();
|
||||
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
/// Check and insert message ID for deduplication.
|
||||
async fn is_duplicate(&self, msg_id: &str) -> bool {
|
||||
if msg_id.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let mut dedup = self.dedup.write().await;
|
||||
|
||||
if dedup.contains(msg_id) {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Evict oldest half when at capacity
|
||||
if dedup.len() >= DEDUP_CAPACITY {
|
||||
let to_remove: Vec<String> = dedup.iter().take(DEDUP_CAPACITY / 2).cloned().collect();
|
||||
for key in to_remove {
|
||||
dedup.remove(&key);
|
||||
}
|
||||
}
|
||||
|
||||
dedup.insert(msg_id.to_string());
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for QQChannel {
|
||||
fn name(&self) -> &str {
|
||||
"qq"
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> anyhow::Result<()> {
|
||||
let token = self.get_token().await?;
|
||||
|
||||
// Determine if this is a group or private message based on recipient format
|
||||
// Format: "user:{openid}" or "group:{group_openid}"
|
||||
let (url, body) = if let Some(group_id) = message.recipient.strip_prefix("group:") {
|
||||
(
|
||||
format!("{QQ_API_BASE}/v2/groups/{group_id}/messages"),
|
||||
json!({
|
||||
"content": &message.content,
|
||||
"msg_type": 0,
|
||||
}),
|
||||
)
|
||||
} else {
|
||||
let user_id = message
|
||||
.recipient
|
||||
.strip_prefix("user:")
|
||||
.unwrap_or(&message.recipient);
|
||||
(
|
||||
format!("{QQ_API_BASE}/v2/users/{user_id}/messages"),
|
||||
json!({
|
||||
"content": &message.content,
|
||||
"msg_type": 0,
|
||||
}),
|
||||
)
|
||||
};
|
||||
|
||||
ensure_https(&url)?;
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(&url)
|
||||
.header("Authorization", format!("QQBot {token}"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let err = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!("QQ send message failed ({status}): {err}");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_lines)]
|
||||
async fn listen(&self, tx: tokio::sync::mpsc::Sender<ChannelMessage>) -> anyhow::Result<()> {
|
||||
tracing::info!("QQ: authenticating...");
|
||||
let token = self.get_token().await?;
|
||||
|
||||
tracing::info!("QQ: fetching gateway URL...");
|
||||
let gw_url = self.get_gateway_url(&token).await?;
|
||||
|
||||
tracing::info!("QQ: connecting to gateway WebSocket...");
|
||||
let (ws_stream, _) = tokio_tungstenite::connect_async(&gw_url).await?;
|
||||
let (mut write, mut read) = ws_stream.split();
|
||||
|
||||
// Read Hello (opcode 10)
|
||||
let hello = read
|
||||
.next()
|
||||
.await
|
||||
.ok_or(anyhow::anyhow!("QQ: no hello frame"))??;
|
||||
let hello_data: serde_json::Value = serde_json::from_str(&hello.to_string())?;
|
||||
let heartbeat_interval = hello_data
|
||||
.get("d")
|
||||
.and_then(|d| d.get("heartbeat_interval"))
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.unwrap_or(41250);
|
||||
|
||||
// Send Identify (opcode 2)
|
||||
// Intents: PUBLIC_GUILD_MESSAGES (1<<30) | C2C_MESSAGE_CREATE & GROUP_AT_MESSAGE_CREATE (1<<25)
|
||||
let intents: u64 = (1 << 25) | (1 << 30);
|
||||
let identify = json!({
|
||||
"op": 2,
|
||||
"d": {
|
||||
"token": format!("QQBot {token}"),
|
||||
"intents": intents,
|
||||
"properties": {
|
||||
"os": "linux",
|
||||
"browser": "openhuman",
|
||||
"device": "openhuman",
|
||||
}
|
||||
}
|
||||
});
|
||||
write.send(Message::Text(identify.to_string())).await?;
|
||||
|
||||
tracing::info!("QQ: connected and identified");
|
||||
|
||||
let mut sequence: i64 = -1;
|
||||
|
||||
// Spawn heartbeat timer
|
||||
let (hb_tx, mut hb_rx) = tokio::sync::mpsc::channel::<()>(1);
|
||||
let hb_interval = heartbeat_interval;
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(std::time::Duration::from_millis(hb_interval));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
if hb_tx.send(()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = hb_rx.recv() => {
|
||||
let d = if sequence >= 0 { json!(sequence) } else { json!(null) };
|
||||
let hb = json!({"op": 1, "d": d});
|
||||
if write
|
||||
.send(Message::Text(hb.to_string()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
msg = read.next() => {
|
||||
let msg = match msg {
|
||||
Some(Ok(Message::Text(t))) => t,
|
||||
Some(Ok(Message::Close(_))) | None => break,
|
||||
_ => continue,
|
||||
};
|
||||
|
||||
let event: serde_json::Value = match serde_json::from_str(msg.as_ref()) {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
// Track sequence number
|
||||
if let Some(s) = event.get("s").and_then(serde_json::Value::as_i64) {
|
||||
sequence = s;
|
||||
}
|
||||
|
||||
let op = event.get("op").and_then(serde_json::Value::as_u64).unwrap_or(0);
|
||||
|
||||
match op {
|
||||
// Server requests immediate heartbeat
|
||||
1 => {
|
||||
let d = if sequence >= 0 { json!(sequence) } else { json!(null) };
|
||||
let hb = json!({"op": 1, "d": d});
|
||||
if write
|
||||
.send(Message::Text(hb.to_string()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
// Reconnect
|
||||
7 => {
|
||||
tracing::warn!("QQ: received Reconnect (op 7)");
|
||||
break;
|
||||
}
|
||||
// Invalid Session
|
||||
9 => {
|
||||
tracing::warn!("QQ: received Invalid Session (op 9)");
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
// Only process dispatch events (op 0)
|
||||
if op != 0 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let event_type = event.get("t").and_then(|t| t.as_str()).unwrap_or("");
|
||||
let d = match event.get("d") {
|
||||
Some(d) => d,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
match event_type {
|
||||
"C2C_MESSAGE_CREATE" => {
|
||||
let msg_id = d.get("id").and_then(|i| i.as_str()).unwrap_or("");
|
||||
if self.is_duplicate(msg_id).await {
|
||||
continue;
|
||||
}
|
||||
|
||||
let content = d.get("content").and_then(|c| c.as_str()).unwrap_or("").trim();
|
||||
if content.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let author_id = d.get("author").and_then(|a| a.get("id")).and_then(|i| i.as_str()).unwrap_or("unknown");
|
||||
// For QQ, user_openid is the identifier
|
||||
let user_openid = d.get("author").and_then(|a| a.get("user_openid")).and_then(|u| u.as_str()).unwrap_or(author_id);
|
||||
|
||||
if !self.is_user_allowed(user_openid) {
|
||||
tracing::warn!("QQ: ignoring C2C message from unauthorized user: {user_openid}");
|
||||
continue;
|
||||
}
|
||||
|
||||
let chat_id = format!("user:{user_openid}");
|
||||
|
||||
let channel_msg = ChannelMessage {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
sender: user_openid.to_string(),
|
||||
reply_target: chat_id,
|
||||
content: content.to_string(),
|
||||
channel: "qq".to_string(),
|
||||
timestamp: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
thread_ts: None,
|
||||
};
|
||||
|
||||
if tx.send(channel_msg).await.is_err() {
|
||||
tracing::warn!("QQ: message channel closed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
"GROUP_AT_MESSAGE_CREATE" => {
|
||||
let msg_id = d.get("id").and_then(|i| i.as_str()).unwrap_or("");
|
||||
if self.is_duplicate(msg_id).await {
|
||||
continue;
|
||||
}
|
||||
|
||||
let content = d.get("content").and_then(|c| c.as_str()).unwrap_or("").trim();
|
||||
if content.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let author_id = d.get("author").and_then(|a| a.get("member_openid")).and_then(|m| m.as_str()).unwrap_or("unknown");
|
||||
|
||||
if !self.is_user_allowed(author_id) {
|
||||
tracing::warn!("QQ: ignoring group message from unauthorized user: {author_id}");
|
||||
continue;
|
||||
}
|
||||
|
||||
let group_openid = d.get("group_openid").and_then(|g| g.as_str()).unwrap_or("unknown");
|
||||
let chat_id = format!("group:{group_openid}");
|
||||
|
||||
let channel_msg = ChannelMessage {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
sender: author_id.to_string(),
|
||||
reply_target: chat_id,
|
||||
content: content.to_string(),
|
||||
channel: "qq".to_string(),
|
||||
timestamp: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
thread_ts: None,
|
||||
};
|
||||
|
||||
if tx.send(channel_msg).await.is_err() {
|
||||
tracing::warn!("QQ: message channel closed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
anyhow::bail!("QQ WebSocket connection closed")
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
self.fetch_access_token().await.is_ok()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "qq_tests.rs"]
|
||||
mod tests;
|
||||
pub use tinychannels::providers::qq::*;
|
||||
|
||||
@@ -1,84 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_name() {
|
||||
let ch = QQChannel::new("id".into(), "secret".into(), vec![]);
|
||||
assert_eq!(ch.name(), "qq");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_allowed_wildcard() {
|
||||
let ch = QQChannel::new("id".into(), "secret".into(), vec!["*".into()]);
|
||||
assert!(ch.is_user_allowed("anyone"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_allowed_specific() {
|
||||
let ch = QQChannel::new("id".into(), "secret".into(), vec!["user123".into()]);
|
||||
assert!(ch.is_user_allowed("user123"));
|
||||
assert!(!ch.is_user_allowed("other"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_user_denied_empty() {
|
||||
let ch = QQChannel::new("id".into(), "secret".into(), vec![]);
|
||||
assert!(!ch.is_user_allowed("anyone"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_dedup() {
|
||||
let ch = QQChannel::new("id".into(), "secret".into(), vec![]);
|
||||
assert!(!ch.is_duplicate("msg1").await);
|
||||
assert!(ch.is_duplicate("msg1").await);
|
||||
assert!(!ch.is_duplicate("msg2").await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_dedup_empty_id() {
|
||||
let ch = QQChannel::new("id".into(), "secret".into(), vec![]);
|
||||
// Empty IDs should never be considered duplicates
|
||||
assert!(!ch.is_duplicate("").await);
|
||||
assert!(!ch.is_duplicate("").await);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_serde() {
|
||||
let toml_str = r#"
|
||||
app_id = "12345"
|
||||
app_secret = "secret_abc"
|
||||
allowed_users = ["user1"]
|
||||
"#;
|
||||
let config: crate::openhuman::config::schema::QQConfig = toml::from_str(toml_str).unwrap();
|
||||
assert_eq!(config.app_id, "12345");
|
||||
assert_eq!(config.app_secret, "secret_abc");
|
||||
assert_eq!(config.allowed_users, vec!["user1"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ensure_https_accepts_https_urls() {
|
||||
assert!(ensure_https("https://api.example.com").is_ok());
|
||||
assert!(ensure_https("https://api.sgroup.qq.com/v1").is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ensure_https_rejects_http_and_other_schemes() {
|
||||
assert!(ensure_https("http://example.com").is_err());
|
||||
assert!(ensure_https("ws://example.com").is_err());
|
||||
assert!(ensure_https("ftp://example.com").is_err());
|
||||
assert!(ensure_https("").is_err());
|
||||
assert!(ensure_https("example.com").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_base_and_auth_url_are_https_constants() {
|
||||
assert!(QQ_API_BASE.starts_with("https://"));
|
||||
assert!(QQ_AUTH_URL.starts_with("https://"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_constructor_stores_fields() {
|
||||
let ch = QQChannel::new("a".into(), "b".into(), vec!["u1".into()]);
|
||||
assert_eq!(ch.app_id, "a");
|
||||
assert_eq!(ch.app_secret, "b");
|
||||
assert_eq!(ch.allowed_users, vec!["u1".to_string()]);
|
||||
}
|
||||
@@ -1,130 +0,0 @@
|
||||
//! Attachment marker parsing and path/url detection.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum TelegramAttachmentKind {
|
||||
Image,
|
||||
Document,
|
||||
Video,
|
||||
Audio,
|
||||
Voice,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct TelegramAttachment {
|
||||
pub(crate) kind: TelegramAttachmentKind,
|
||||
pub(crate) target: String,
|
||||
}
|
||||
|
||||
impl TelegramAttachmentKind {
|
||||
fn from_marker(marker: &str) -> Option<Self> {
|
||||
match marker.trim().to_ascii_uppercase().as_str() {
|
||||
"IMAGE" | "PHOTO" => Some(Self::Image),
|
||||
"DOCUMENT" | "FILE" => Some(Self::Document),
|
||||
"VIDEO" => Some(Self::Video),
|
||||
"AUDIO" => Some(Self::Audio),
|
||||
"VOICE" => Some(Self::Voice),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_http_url(target: &str) -> bool {
|
||||
target.starts_with("http://") || target.starts_with("https://")
|
||||
}
|
||||
|
||||
pub(crate) fn infer_attachment_kind_from_target(target: &str) -> Option<TelegramAttachmentKind> {
|
||||
let normalized = target
|
||||
.split('?')
|
||||
.next()
|
||||
.unwrap_or(target)
|
||||
.split('#')
|
||||
.next()
|
||||
.unwrap_or(target);
|
||||
|
||||
let extension = Path::new(normalized)
|
||||
.extension()
|
||||
.and_then(|ext| ext.to_str())?
|
||||
.to_ascii_lowercase();
|
||||
|
||||
match extension.as_str() {
|
||||
"png" | "jpg" | "jpeg" | "gif" | "webp" | "bmp" => Some(TelegramAttachmentKind::Image),
|
||||
"mp4" | "mov" | "mkv" | "avi" | "webm" => Some(TelegramAttachmentKind::Video),
|
||||
"mp3" | "m4a" | "wav" | "flac" => Some(TelegramAttachmentKind::Audio),
|
||||
"ogg" | "oga" | "opus" => Some(TelegramAttachmentKind::Voice),
|
||||
"pdf" | "txt" | "md" | "csv" | "json" | "zip" | "tar" | "gz" | "doc" | "docx" | "xls"
|
||||
| "xlsx" | "ppt" | "pptx" => Some(TelegramAttachmentKind::Document),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn parse_path_only_attachment(message: &str) -> Option<TelegramAttachment> {
|
||||
let trimmed = message.trim();
|
||||
if trimmed.is_empty() || trimmed.contains('\n') {
|
||||
return None;
|
||||
}
|
||||
|
||||
let candidate = trimmed.trim_matches(|c| matches!(c, '`' | '"' | '\''));
|
||||
if candidate.chars().any(char::is_whitespace) {
|
||||
return None;
|
||||
}
|
||||
|
||||
let candidate = candidate.strip_prefix("file://").unwrap_or(candidate);
|
||||
let kind = infer_attachment_kind_from_target(candidate)?;
|
||||
|
||||
if !is_http_url(candidate) && !Path::new(candidate).exists() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(TelegramAttachment {
|
||||
kind,
|
||||
target: candidate.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn parse_attachment_markers(message: &str) -> (String, Vec<TelegramAttachment>) {
|
||||
let mut cleaned = String::with_capacity(message.len());
|
||||
let mut attachments = Vec::new();
|
||||
let mut cursor = 0;
|
||||
|
||||
while cursor < message.len() {
|
||||
let Some(open_rel) = message[cursor..].find('[') else {
|
||||
cleaned.push_str(&message[cursor..]);
|
||||
break;
|
||||
};
|
||||
|
||||
let open = cursor + open_rel;
|
||||
cleaned.push_str(&message[cursor..open]);
|
||||
|
||||
let Some(close_rel) = message[open..].find(']') else {
|
||||
cleaned.push_str(&message[open..]);
|
||||
break;
|
||||
};
|
||||
|
||||
let close = open + close_rel;
|
||||
let marker = &message[open + 1..close];
|
||||
|
||||
let parsed = marker.split_once(':').and_then(|(kind, target)| {
|
||||
let kind = TelegramAttachmentKind::from_marker(kind)?;
|
||||
let target = target.trim();
|
||||
if target.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(TelegramAttachment {
|
||||
kind,
|
||||
target: target.to_string(),
|
||||
})
|
||||
});
|
||||
|
||||
if let Some(attachment) = parsed {
|
||||
attachments.push(attachment);
|
||||
} else {
|
||||
cleaned.push_str(&message[open..=close]);
|
||||
}
|
||||
|
||||
cursor = close + 1;
|
||||
}
|
||||
|
||||
(cleaned.trim().to_string(), attachments)
|
||||
}
|
||||
@@ -1,21 +0,0 @@
|
||||
//! Telegram Bot API channel implementation.
|
||||
//!
|
||||
//! This module is the orchestration entry point for the Telegram channel.
|
||||
//! Implementation is split across sibling modules by concern:
|
||||
//!
|
||||
//! - [`super::channel_types`] — struct definition and private helper types
|
||||
//! - [`super::channel_core`] — constructor, config, pairing/auth, API plumbing
|
||||
//! - [`super::channel_recv`] — inbound parsing, allowlist checks, mention filtering
|
||||
//! - [`super::channel_send`] — outbound text, media, reactions, attachments
|
||||
//! - [`super::channel_ops`] — `Channel` trait impl (send/listen/draft/typing)
|
||||
|
||||
// Re-export so that the `#[path = "channel_tests.rs"]` test module can reach
|
||||
// `TelegramChannel` via `super::TelegramChannel`.
|
||||
#[allow(unused_imports)]
|
||||
pub use super::channel_types::TelegramChannel;
|
||||
#[cfg(test)]
|
||||
pub(super) use super::channel_types::TelegramTypingTask;
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "channel_tests.rs"]
|
||||
mod tests;
|
||||
@@ -1,350 +0,0 @@
|
||||
//! Telegram channel — constructor, configuration, auth/pairing, and API plumbing helpers.
|
||||
|
||||
use super::channel_types::{
|
||||
TelegramChannel, TelegramUpdateWindow, TELEGRAM_RECENT_UPDATE_CACHE_SIZE,
|
||||
};
|
||||
use super::text::{TELEGRAM_BIND_COMMAND, TELEGRAM_START_COMMAND};
|
||||
use crate::openhuman::config::{Config, StreamMode};
|
||||
use crate::openhuman::security::pairing::PairingGuard;
|
||||
use anyhow::Context;
|
||||
use directories::UserDirs;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use tokio::fs;
|
||||
|
||||
/// Resolve the Telegram API base URL from an optional env value. Pure function —
|
||||
/// callers in production pass `std::env::var("OPENHUMAN_TELEGRAM_BOT_API_BASE").ok()`
|
||||
/// (falling back to the legacy `OPENHUMAN_TELEGRAM_API_BASE`);
|
||||
/// tests can exercise this directly without mutating process env.
|
||||
pub(crate) fn resolve_api_base(raw: Option<String>) -> String {
|
||||
let base = raw
|
||||
.filter(|v| !v.trim().is_empty())
|
||||
.unwrap_or_else(|| "https://api.telegram.org".to_string());
|
||||
base.trim_end_matches('/').to_string()
|
||||
}
|
||||
|
||||
impl TelegramChannel {
|
||||
pub fn new(bot_token: String, allowed_users: Vec<String>, mention_only: bool) -> Self {
|
||||
let api_base = resolve_api_base(
|
||||
std::env::var("OPENHUMAN_TELEGRAM_BOT_API_BASE")
|
||||
.ok()
|
||||
.or_else(|| std::env::var("OPENHUMAN_TELEGRAM_API_BASE").ok()),
|
||||
);
|
||||
tracing::debug!(
|
||||
target: "telegram::api",
|
||||
api_base = %api_base,
|
||||
"Using Telegram API base URL"
|
||||
);
|
||||
|
||||
let normalized_allowed = Self::normalize_allowed_users(allowed_users);
|
||||
let pairing = if normalized_allowed.is_empty() {
|
||||
let (guard, code_opt) = PairingGuard::new(true, &[]);
|
||||
if let Some(code) = code_opt {
|
||||
println!(" 🔐 Telegram pairing required. One-time bind code: {code}");
|
||||
println!(" Send `{TELEGRAM_BIND_COMMAND} <code>` from your Telegram account.");
|
||||
}
|
||||
Some(guard)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Self {
|
||||
bot_token,
|
||||
chat_id: None,
|
||||
api_base,
|
||||
allowed_users: Arc::new(RwLock::new(normalized_allowed)),
|
||||
pairing,
|
||||
client: reqwest::Client::new(),
|
||||
stream_mode: StreamMode::Off,
|
||||
draft_update_interval_ms: 1000,
|
||||
silent_streaming: true,
|
||||
last_draft_edit: parking_lot::Mutex::new(std::collections::HashMap::new()),
|
||||
typing_handle: parking_lot::Mutex::new(None),
|
||||
mention_only,
|
||||
bot_username: parking_lot::Mutex::new(None),
|
||||
recent_updates: parking_lot::Mutex::new(TelegramUpdateWindow::default()),
|
||||
recent_approval_prompts: parking_lot::Mutex::new(std::collections::HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Configure streaming mode for progressive draft updates.
|
||||
/// Configure streaming mode for progressive draft updates.
|
||||
pub fn with_streaming(
|
||||
mut self,
|
||||
stream_mode: StreamMode,
|
||||
draft_update_interval_ms: u64,
|
||||
silent_streaming: bool,
|
||||
) -> Self {
|
||||
self.stream_mode = stream_mode;
|
||||
self.draft_update_interval_ms = draft_update_interval_ms;
|
||||
self.silent_streaming = silent_streaming;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the default chat for recipient-less proactive sends. A blank or
|
||||
/// whitespace-only value is treated as unset (`None`), so proactive routing
|
||||
/// skips Telegram rather than POSTing to an empty `chat_id`.
|
||||
pub fn with_chat_id(mut self, chat_id: Option<String>) -> Self {
|
||||
self.chat_id = chat_id
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty());
|
||||
self
|
||||
}
|
||||
|
||||
/// Parse reply_target into (chat_id, optional thread_id).
|
||||
pub(crate) fn parse_reply_target(reply_target: &str) -> (String, Option<String>) {
|
||||
if let Some((chat_id, thread_id)) = reply_target.split_once(':') {
|
||||
(chat_id.to_string(), Some(thread_id.to_string()))
|
||||
} else {
|
||||
(reply_target.to_string(), None)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn parse_message_id(value: Option<&str>) -> Option<i64> {
|
||||
value.and_then(|raw| raw.trim().parse::<i64>().ok())
|
||||
}
|
||||
|
||||
pub(crate) fn http_client(&self) -> reqwest::Client {
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.telegram")
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_identity(value: &str) -> String {
|
||||
value.trim().trim_start_matches('@').to_string()
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_allowed_users(allowed_users: Vec<String>) -> Vec<String> {
|
||||
allowed_users
|
||||
.into_iter()
|
||||
.map(|entry| Self::normalize_identity(&entry))
|
||||
.filter(|entry| !entry.is_empty())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn api_url(&self, method: &str) -> String {
|
||||
format!("{}/bot{}/{method}", self.api_base, self.bot_token)
|
||||
}
|
||||
|
||||
/// Point outbound Telegram API calls at `base` (test-only seam). Used to
|
||||
/// aim `send()` at a dead local port so onboarding tests exercise the
|
||||
/// decision logic without reaching api.telegram.org.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn set_api_base_for_tests(&mut self, base: impl Into<String>) {
|
||||
self.api_base = base.into();
|
||||
}
|
||||
|
||||
pub(crate) fn pairing_code_active(&self) -> bool {
|
||||
self.pairing
|
||||
.as_ref()
|
||||
.and_then(PairingGuard::pairing_code)
|
||||
.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn extract_bind_code(text: &str) -> Option<&str> {
|
||||
let mut parts = text.split_whitespace();
|
||||
let command = parts.next()?;
|
||||
let base_command = command.split('@').next().unwrap_or(command);
|
||||
if base_command != TELEGRAM_BIND_COMMAND {
|
||||
return None;
|
||||
}
|
||||
parts.next().map(str::trim).filter(|code| !code.is_empty())
|
||||
}
|
||||
|
||||
/// Whether `text` is the standard Telegram `/start` bot-onboarding command
|
||||
/// (optionally addressed as `/start@botname`, with or without a payload).
|
||||
///
|
||||
/// On the self-bot-token path this is the operator's explicit "I'm setting up
|
||||
/// my bot" signal: the first `/start` while pairing is still pending pairs the
|
||||
/// sender (see `handle_unauthorized_message`), matching the "first sender after
|
||||
/// /start" behaviour sanctioned by openhuman#4381.
|
||||
pub(crate) fn is_start_command(text: &str) -> bool {
|
||||
let Some(command) = text.split_whitespace().next() else {
|
||||
return false;
|
||||
};
|
||||
let base_command = command.split('@').next().unwrap_or(command);
|
||||
base_command == TELEGRAM_START_COMMAND
|
||||
}
|
||||
|
||||
pub(crate) fn track_update_id(&self, update_id: i64) -> bool {
|
||||
let mut window = self.recent_updates.lock();
|
||||
if window.recent_lookup.contains(&update_id) {
|
||||
tracing::debug!(
|
||||
update_id,
|
||||
"Telegram update dedupe hit: duplicate update skipped"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
|
||||
if update_id < window.max_seen_update_id {
|
||||
tracing::debug!(
|
||||
update_id,
|
||||
max_seen = window.max_seen_update_id,
|
||||
"Telegram update ordering safeguard: stale update skipped"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
|
||||
if update_id > window.max_seen_update_id {
|
||||
window.max_seen_update_id = update_id;
|
||||
}
|
||||
|
||||
window.recent_lookup.insert(update_id);
|
||||
window.recent_order.push_back(update_id);
|
||||
if window.recent_order.len() > TELEGRAM_RECENT_UPDATE_CACHE_SIZE {
|
||||
if let Some(evicted) = window.recent_order.pop_front() {
|
||||
window.recent_lookup.remove(&evicted);
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// Clears Bot API webhook mode so `getUpdates` long polling can run.
|
||||
pub(crate) async fn delete_webhook_for_long_polling(&self) -> bool {
|
||||
let url = self.api_url("deleteWebhook");
|
||||
let body = serde_json::json!({ "drop_pending_updates": false });
|
||||
tracing::info!(
|
||||
"[telegram] deleteWebhook: enabling getUpdates polling (drop_pending_updates=false)"
|
||||
);
|
||||
match self.http_client().post(&url).json(&body).send().await {
|
||||
Ok(resp) => Self::telegram_api_ok(resp).await,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "[telegram] deleteWebhook HTTP request failed");
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn telegram_api_ok(resp: reqwest::Response) -> bool {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
if !status.is_success() {
|
||||
tracing::warn!(status = ?status, body, "Telegram API request failed");
|
||||
return false;
|
||||
}
|
||||
|
||||
match serde_json::from_str::<serde_json::Value>(&body) {
|
||||
Ok(payload) => {
|
||||
if payload
|
||||
.get("ok")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
true
|
||||
} else {
|
||||
let error_code = payload
|
||||
.get("error_code")
|
||||
.and_then(serde_json::Value::as_i64)
|
||||
.unwrap_or_default();
|
||||
let description = payload
|
||||
.get("description")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or("unknown Telegram API error");
|
||||
tracing::warn!(
|
||||
status = ?status,
|
||||
error_code,
|
||||
description,
|
||||
body,
|
||||
"Telegram API responded with ok=false"
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
status = ?status,
|
||||
%error,
|
||||
body,
|
||||
"Telegram API returned non-JSON body"
|
||||
);
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn fetch_bot_username(&self) -> anyhow::Result<String> {
|
||||
let resp = self.http_client().get(self.api_url("getMe")).send().await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
anyhow::bail!("Failed to fetch bot info: {}", resp.status());
|
||||
}
|
||||
|
||||
let data: serde_json::Value = resp.json().await?;
|
||||
let username = data
|
||||
.get("result")
|
||||
.and_then(|r| r.get("username"))
|
||||
.and_then(|u| u.as_str())
|
||||
.context("Bot username not found in response")?;
|
||||
|
||||
Ok(username.to_string())
|
||||
}
|
||||
|
||||
pub(crate) async fn get_bot_username(&self) -> Option<String> {
|
||||
{
|
||||
let cache = self.bot_username.lock();
|
||||
if let Some(ref username) = *cache {
|
||||
return Some(username.clone());
|
||||
}
|
||||
}
|
||||
|
||||
match self.fetch_bot_username().await {
|
||||
Ok(username) => {
|
||||
let mut cache = self.bot_username.lock();
|
||||
*cache = Some(username.clone());
|
||||
Some(username)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to fetch bot username: {e}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_config_without_env() -> anyhow::Result<Config> {
|
||||
let home = UserDirs::new()
|
||||
.map(|u| u.home_dir().to_path_buf())
|
||||
.context("Could not find home directory")?;
|
||||
let openhuman_dir = home.join(".openhuman");
|
||||
let config_path = openhuman_dir.join("config.toml");
|
||||
|
||||
let contents = fs::read_to_string(&config_path)
|
||||
.await
|
||||
.with_context(|| format!("Failed to read config file: {}", config_path.display()))?;
|
||||
let mut config: Config = toml::from_str(&contents)
|
||||
.context("Failed to parse config file for Telegram binding")?;
|
||||
config.config_path = config_path;
|
||||
config.workspace_dir = openhuman_dir.join("workspace");
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
pub(crate) async fn persist_allowed_identity(&self, identity: &str) -> anyhow::Result<()> {
|
||||
let mut config = Self::load_config_without_env().await?;
|
||||
let Some(telegram) = config.channels_config.telegram.as_mut() else {
|
||||
anyhow::bail!("Telegram channel config is missing in config.toml");
|
||||
};
|
||||
|
||||
let normalized = Self::normalize_identity(identity);
|
||||
if normalized.is_empty() {
|
||||
anyhow::bail!("Cannot persist empty Telegram identity");
|
||||
}
|
||||
|
||||
if !telegram.allowed_users.iter().any(|u| u == &normalized) {
|
||||
telegram.allowed_users.push(normalized);
|
||||
config
|
||||
.save()
|
||||
.await
|
||||
.context("Failed to persist Telegram allowlist to config.toml")?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn add_allowed_identity_runtime(&self, identity: &str) {
|
||||
let normalized = Self::normalize_identity(identity);
|
||||
if normalized.is_empty() {
|
||||
return;
|
||||
}
|
||||
if let Ok(mut users) = self.allowed_users.write() {
|
||||
if !users.iter().any(|u| u == &normalized) {
|
||||
users.push(normalized);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,535 +0,0 @@
|
||||
//! Telegram channel — `Channel` trait implementation: send, listen, draft streaming, typing.
|
||||
|
||||
use super::attachments::{parse_attachment_markers, parse_path_only_attachment};
|
||||
use super::channel_types::{TelegramChannel, TelegramTypingTask};
|
||||
use super::text::{strip_tool_call_tags, TELEGRAM_MAX_MESSAGE_LENGTH};
|
||||
use crate::core::event_bus::{publish_global, DomainEvent};
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
use crate::openhuman::config::StreamMode;
|
||||
use async_trait::async_trait;
|
||||
use std::time::Duration;
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for TelegramChannel {
|
||||
fn name(&self) -> &str {
|
||||
"telegram"
|
||||
}
|
||||
|
||||
/// Recipient-less proactive sends (cron/heartbeat) deliver to the bot's
|
||||
/// configured default `chat_id`. `None` when unconfigured, so proactive
|
||||
/// routing skips Telegram rather than letting `send` POST to an empty
|
||||
/// `chat_id` (mirrors Discord — #3712 Telegram parity).
|
||||
fn proactive_target(&self) -> Option<String> {
|
||||
self.chat_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(str::to_string)
|
||||
}
|
||||
|
||||
fn supports_reactions(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn supports_draft_updates(&self) -> bool {
|
||||
self.stream_mode != StreamMode::Off
|
||||
}
|
||||
|
||||
async fn send_draft(&self, message: &SendMessage) -> anyhow::Result<Option<String>> {
|
||||
if self.stream_mode == StreamMode::Off {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let (chat_id, thread_id) = Self::parse_reply_target(&message.recipient);
|
||||
let parent_message_id = Self::parse_message_id(message.thread_ts.as_deref());
|
||||
let initial_text = if message.content.is_empty() {
|
||||
"...".to_string()
|
||||
} else {
|
||||
message.content.clone()
|
||||
};
|
||||
|
||||
let mut body = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"text": initial_text,
|
||||
});
|
||||
if let Some(tid) = thread_id {
|
||||
body["message_thread_id"] = serde_json::Value::String(tid.to_string());
|
||||
}
|
||||
if let Some(parent_id) = parent_message_id {
|
||||
body["reply_to_message_id"] = serde_json::Value::from(parent_id);
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(self.api_url("sendMessage"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!("Telegram sendMessage (draft) failed: {err}");
|
||||
}
|
||||
|
||||
let resp_json: serde_json::Value = resp.json().await?;
|
||||
let message_id = resp_json
|
||||
.get("result")
|
||||
.and_then(|r| r.get("message_id"))
|
||||
.and_then(|id| id.as_i64())
|
||||
.map(|id| id.to_string());
|
||||
|
||||
self.last_draft_edit
|
||||
.lock()
|
||||
.insert(chat_id.to_string(), std::time::Instant::now());
|
||||
|
||||
Ok(message_id)
|
||||
}
|
||||
|
||||
async fn update_draft(
|
||||
&self,
|
||||
recipient: &str,
|
||||
message_id: &str,
|
||||
text: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
let (chat_id, _) = Self::parse_reply_target(recipient);
|
||||
|
||||
// Rate-limit edits per chat
|
||||
{
|
||||
let last_edits = self.last_draft_edit.lock();
|
||||
if let Some(last_time) = last_edits.get(&chat_id) {
|
||||
let elapsed = u64::try_from(last_time.elapsed().as_millis()).unwrap_or(u64::MAX);
|
||||
if elapsed < self.draft_update_interval_ms {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Truncate to Telegram limit for mid-stream edits (UTF-8 safe)
|
||||
let display_text = if text.len() > TELEGRAM_MAX_MESSAGE_LENGTH {
|
||||
let mut end = 0;
|
||||
for (idx, ch) in text.char_indices() {
|
||||
let next = idx + ch.len_utf8();
|
||||
if next > TELEGRAM_MAX_MESSAGE_LENGTH {
|
||||
break;
|
||||
}
|
||||
end = next;
|
||||
}
|
||||
&text[..end]
|
||||
} else {
|
||||
text
|
||||
};
|
||||
|
||||
let message_id_parsed = match message_id.parse::<i64>() {
|
||||
Ok(id) => id,
|
||||
Err(e) => {
|
||||
tracing::warn!("Invalid Telegram message_id '{message_id}': {e}");
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
let body = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"message_id": message_id_parsed,
|
||||
"text": display_text,
|
||||
});
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(self.api_url("editMessageText"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if resp.status().is_success() {
|
||||
self.last_draft_edit
|
||||
.lock()
|
||||
.insert(chat_id.clone(), std::time::Instant::now());
|
||||
} else {
|
||||
let status = resp.status();
|
||||
let err = resp.text().await.unwrap_or_default();
|
||||
tracing::debug!("Telegram editMessageText failed ({status}): {err}");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn finalize_draft(
|
||||
&self,
|
||||
recipient: &str,
|
||||
message_id: &str,
|
||||
text: &str,
|
||||
thread_ts: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let text = &strip_tool_call_tags(text);
|
||||
let (chat_id, thread_id) = Self::parse_reply_target(recipient);
|
||||
let parent_message_id = Self::parse_message_id(thread_ts);
|
||||
|
||||
// Clean up rate-limit tracking for this chat
|
||||
self.last_draft_edit.lock().remove(&chat_id);
|
||||
|
||||
// If text exceeds limit, delete draft and send as chunked messages
|
||||
if text.len() > TELEGRAM_MAX_MESSAGE_LENGTH {
|
||||
let msg_id = match message_id.parse::<i64>() {
|
||||
Ok(id) => id,
|
||||
Err(e) => {
|
||||
tracing::warn!("Invalid Telegram message_id '{message_id}': {e}");
|
||||
return self
|
||||
.send_text_chunks(text, &chat_id, thread_id.as_deref(), parent_message_id)
|
||||
.await;
|
||||
}
|
||||
};
|
||||
|
||||
// Delete the draft
|
||||
let _ = self
|
||||
.client
|
||||
.post(self.api_url("deleteMessage"))
|
||||
.json(&serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"message_id": msg_id,
|
||||
}))
|
||||
.send()
|
||||
.await;
|
||||
|
||||
// Fall back to chunked send
|
||||
return self
|
||||
.send_text_chunks(text, &chat_id, thread_id.as_deref(), parent_message_id)
|
||||
.await;
|
||||
}
|
||||
|
||||
let msg_id = match message_id.parse::<i64>() {
|
||||
Ok(id) => id,
|
||||
Err(e) => {
|
||||
tracing::warn!("Invalid Telegram message_id '{message_id}': {e}");
|
||||
return self
|
||||
.send_text_chunks(text, &chat_id, thread_id.as_deref(), parent_message_id)
|
||||
.await;
|
||||
}
|
||||
};
|
||||
|
||||
// Try editing with Markdown formatting
|
||||
let body = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"message_id": msg_id,
|
||||
"text": text,
|
||||
"parse_mode": "Markdown",
|
||||
});
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(self.api_url("editMessageText"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if resp.status().is_success() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Markdown failed — retry without parse_mode
|
||||
let plain_body = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"message_id": msg_id,
|
||||
"text": text,
|
||||
});
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(self.api_url("editMessageText"))
|
||||
.json(&plain_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if resp.status().is_success() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Edit failed entirely — fall back to new message
|
||||
tracing::warn!("Telegram finalize_draft edit failed; falling back to sendMessage");
|
||||
self.send_text_chunks(text, &chat_id, thread_id.as_deref(), parent_message_id)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> anyhow::Result<()> {
|
||||
// Strip tool_call tags before processing to prevent Markdown parsing failures
|
||||
let content = strip_tool_call_tags(&message.content);
|
||||
let parent_message_id = Self::parse_message_id(message.thread_ts.as_deref());
|
||||
|
||||
// Parse recipient: "chat_id" or "chat_id:thread_id" format
|
||||
let (chat_id, thread_id) = match message.recipient.split_once(':') {
|
||||
Some((chat, thread)) => (chat, Some(thread)),
|
||||
None => (message.recipient.as_str(), None),
|
||||
};
|
||||
|
||||
let (reactionless_content, reaction_marker) = Self::parse_reaction_marker(&content);
|
||||
if let Some(reaction_marker) = reaction_marker.as_deref() {
|
||||
let (emoji, explicit_target_id) = match reaction_marker.split_once('|') {
|
||||
Some((emoji, target)) => (emoji.trim(), Self::parse_message_id(Some(target))),
|
||||
None => (reaction_marker.trim(), None),
|
||||
};
|
||||
let target_message_id = explicit_target_id.or(parent_message_id);
|
||||
if let Some(target_id) = target_message_id {
|
||||
let _ = self
|
||||
.send_message_reaction(chat_id, target_id, emoji)
|
||||
.await?;
|
||||
tracing::debug!(
|
||||
chat_id,
|
||||
target_id,
|
||||
emoji,
|
||||
has_reply = !reactionless_content.is_empty(),
|
||||
"[telegram] reaction sent; continuing to send reply text if present"
|
||||
);
|
||||
} else {
|
||||
tracing::warn!(
|
||||
recipient = message.recipient,
|
||||
marker = reaction_marker,
|
||||
"[telegram] reaction marker ignored: missing target message id"
|
||||
);
|
||||
}
|
||||
// If no text follows the reaction marker, we are done.
|
||||
if reactionless_content.trim().is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
let (text_without_markers, attachments) = parse_attachment_markers(&reactionless_content);
|
||||
|
||||
if !attachments.is_empty() {
|
||||
if !text_without_markers.is_empty() {
|
||||
self.send_text_chunks(&text_without_markers, chat_id, thread_id, parent_message_id)
|
||||
.await?;
|
||||
}
|
||||
|
||||
for attachment in &attachments {
|
||||
self.send_attachment(chat_id, thread_id, attachment).await?;
|
||||
}
|
||||
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if let Some(attachment) = parse_path_only_attachment(&reactionless_content) {
|
||||
self.send_attachment(chat_id, thread_id, &attachment)
|
||||
.await?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
self.send_text_chunks(&reactionless_content, chat_id, thread_id, parent_message_id)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn listen(&self, tx: tokio::sync::mpsc::Sender<ChannelMessage>) -> anyhow::Result<()> {
|
||||
let mut offset: i64 = 0;
|
||||
|
||||
if self.mention_only {
|
||||
let _ = self.get_bot_username().await;
|
||||
}
|
||||
|
||||
tracing::info!("Telegram channel listening for messages...");
|
||||
|
||||
loop {
|
||||
if self.mention_only {
|
||||
let missing_username = self.bot_username.lock().is_none();
|
||||
if missing_username {
|
||||
let _ = self.get_bot_username().await;
|
||||
}
|
||||
}
|
||||
|
||||
let url = self.api_url("getUpdates");
|
||||
let body = serde_json::json!({
|
||||
"offset": offset,
|
||||
"timeout": 30,
|
||||
"allowed_updates": ["message", "edited_message", "message_reaction"]
|
||||
});
|
||||
|
||||
let resp = match self.http_client().post(&url).json(&body).send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::warn!("Telegram poll error: {e}");
|
||||
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let data: serde_json::Value = match resp.json().await {
|
||||
Ok(d) => d,
|
||||
Err(e) => {
|
||||
tracing::warn!("Telegram parse error: {e}");
|
||||
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let ok = data
|
||||
.get("ok")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(true);
|
||||
if !ok {
|
||||
let error_code = data
|
||||
.get("error_code")
|
||||
.and_then(serde_json::Value::as_i64)
|
||||
.unwrap_or_default();
|
||||
let description = data
|
||||
.get("description")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or("unknown Telegram API error");
|
||||
|
||||
if error_code == 409 {
|
||||
let webhook_blocks_polling = description.to_lowercase().contains("webhook");
|
||||
if webhook_blocks_polling {
|
||||
tracing::warn!(
|
||||
"[telegram] getUpdates conflict (409): webhook is active; calling deleteWebhook"
|
||||
);
|
||||
if self.delete_webhook_for_long_polling().await {
|
||||
tracing::info!("[telegram] deleteWebhook ok; retrying getUpdates");
|
||||
continue;
|
||||
}
|
||||
tracing::warn!("[telegram] deleteWebhook did not succeed; backing off");
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"Telegram polling conflict (409): {description}. \
|
||||
Ensure only one `openhuman` process is using this bot token."
|
||||
);
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"Telegram getUpdates API error (code={}): {description}",
|
||||
error_code
|
||||
);
|
||||
tokio::time::sleep(std::time::Duration::from_secs(5)).await;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(results) = data.get("result").and_then(serde_json::Value::as_array) {
|
||||
for update in results {
|
||||
let update_id = update
|
||||
.get("update_id")
|
||||
.and_then(serde_json::Value::as_i64)
|
||||
.unwrap_or_default();
|
||||
if update_id > 0 && !self.track_update_id(update_id) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Advance offset past this update
|
||||
if let Some(uid) = update.get("update_id").and_then(serde_json::Value::as_i64) {
|
||||
offset = uid + 1;
|
||||
}
|
||||
|
||||
if let Some(reaction) = self.parse_update_reaction(update) {
|
||||
tracing::info!(
|
||||
sender = reaction.sender,
|
||||
reply_target = reaction.reply_target,
|
||||
target_message_id = reaction.target_message_id,
|
||||
emoji = reaction.emoji,
|
||||
"Telegram reaction received"
|
||||
);
|
||||
publish_global(DomainEvent::ChannelReactionReceived {
|
||||
channel: "telegram".to_string(),
|
||||
sender: reaction.sender,
|
||||
target_message_id: format!(
|
||||
"telegram_{}_{}",
|
||||
reaction.reply_target, reaction.target_message_id
|
||||
),
|
||||
emoji: reaction.emoji,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(msg) = self.parse_update_message_or_voice(update).await else {
|
||||
self.handle_unauthorized_message(update).await;
|
||||
continue;
|
||||
};
|
||||
|
||||
if tx.send(msg).await.is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
let timeout_duration = Duration::from_secs(5);
|
||||
|
||||
match tokio::time::timeout(
|
||||
timeout_duration,
|
||||
self.http_client().get(self.api_url("getMe")).send(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Ok(resp)) => resp.status().is_success(),
|
||||
Ok(Err(e)) => {
|
||||
tracing::debug!("Telegram health check failed: {e}");
|
||||
false
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::debug!("Telegram health check timed out after 5s");
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn start_typing(&self, recipient: &str) -> anyhow::Result<()> {
|
||||
tracing::info!(recipient, "Telegram start_typing invoked");
|
||||
// Emit immediately so short model turns still show "typing…"
|
||||
self.send_typing_action_once(recipient).await;
|
||||
|
||||
{
|
||||
let guard = self.typing_handle.lock();
|
||||
if guard
|
||||
.as_ref()
|
||||
.is_some_and(|task| task.recipient == recipient)
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
self.stop_typing(recipient).await?;
|
||||
|
||||
let client = self.http_client();
|
||||
let url = self.api_url("sendChatAction");
|
||||
let recipient_owned = recipient.to_string();
|
||||
let recipient_for_log = recipient_owned.clone();
|
||||
let body = Self::typing_body_for_recipient(recipient);
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
loop {
|
||||
match client.post(&url).json(&body).send().await {
|
||||
Ok(resp) => {
|
||||
if !Self::telegram_api_ok(resp).await {
|
||||
tracing::warn!(
|
||||
recipient = recipient_for_log,
|
||||
"Telegram typing refresh rejected"
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
recipient = recipient_for_log,
|
||||
%error,
|
||||
"Telegram typing refresh request failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
// Telegram typing indicator expires after 5s; refresh at 4s
|
||||
tokio::time::sleep(Duration::from_secs(4)).await;
|
||||
}
|
||||
});
|
||||
|
||||
let mut guard = self.typing_handle.lock();
|
||||
*guard = Some(TelegramTypingTask {
|
||||
recipient: recipient_owned,
|
||||
handle,
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stop_typing(&self, _recipient: &str) -> anyhow::Result<()> {
|
||||
tracing::info!("Telegram stop_typing invoked");
|
||||
let mut guard = self.typing_handle.lock();
|
||||
if let Some(task) = guard.take() {
|
||||
task.handle.abort();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,684 +0,0 @@
|
||||
//! Telegram channel — outbound message sending: text chunking, media uploads, reaction sending,
|
||||
//! and attachment dispatch.
|
||||
|
||||
use super::attachments::{is_http_url, TelegramAttachment, TelegramAttachmentKind};
|
||||
use super::channel_types::TelegramChannel;
|
||||
use super::text::split_message_for_telegram;
|
||||
use crate::core::event_bus::{publish_global, DomainEvent};
|
||||
use reqwest::multipart::{Form, Part};
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
impl TelegramChannel {
|
||||
pub(crate) fn parse_reaction_marker(content: &str) -> (String, Option<String>) {
|
||||
// Marker format at the start of the message: [REACTION:😀] or [REACTION:😀|12345]
|
||||
// The marker may be followed by a text reply: [REACTION:👍] Great point!
|
||||
// Returns (remaining_text, Some(marker_inner)) or (original, None).
|
||||
let trimmed = content.trim();
|
||||
let Some(rest) = trimmed.strip_prefix("[REACTION:") else {
|
||||
return (content.to_string(), None);
|
||||
};
|
||||
let Some(close_pos) = rest.find(']') else {
|
||||
return (content.to_string(), None);
|
||||
};
|
||||
let inner = rest[..close_pos].trim();
|
||||
if inner.is_empty() {
|
||||
return (String::new(), None);
|
||||
}
|
||||
let remaining = rest[close_pos + 1..].trim().to_string();
|
||||
(remaining, Some(inner.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn send_message_reaction(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
message_id: i64,
|
||||
emoji: &str,
|
||||
) -> anyhow::Result<bool> {
|
||||
let emoji = emoji.trim();
|
||||
if emoji.is_empty() {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let body = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"message_id": message_id,
|
||||
"reaction": [
|
||||
{
|
||||
"type": "emoji",
|
||||
"emoji": emoji
|
||||
}
|
||||
],
|
||||
"is_big": false
|
||||
});
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("setMessageReaction"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
if resp.status().is_success() {
|
||||
publish_global(DomainEvent::ChannelReactionSent {
|
||||
channel: "telegram".to_string(),
|
||||
target_message_id: format!("telegram_{chat_id}_{message_id}"),
|
||||
emoji: emoji.to_string(),
|
||||
success: true,
|
||||
});
|
||||
tracing::info!(chat_id, message_id, emoji, "Telegram reaction sent");
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let status = resp.status();
|
||||
let err = resp.text().await.unwrap_or_default();
|
||||
tracing::warn!(
|
||||
chat_id,
|
||||
message_id,
|
||||
emoji,
|
||||
status = ?status,
|
||||
error = err,
|
||||
"Telegram reaction not applied; continuing without failure"
|
||||
);
|
||||
publish_global(DomainEvent::ChannelReactionSent {
|
||||
channel: "telegram".to_string(),
|
||||
target_message_id: format!("telegram_{chat_id}_{message_id}"),
|
||||
emoji: emoji.to_string(),
|
||||
success: false,
|
||||
});
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
pub(crate) async fn send_text_chunks(
|
||||
&self,
|
||||
message: &str,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
reply_to_message_id: Option<i64>,
|
||||
) -> anyhow::Result<()> {
|
||||
let chunks = split_message_for_telegram(message);
|
||||
|
||||
for (index, chunk) in chunks.iter().enumerate() {
|
||||
let text = if chunks.len() > 1 {
|
||||
if index == 0 {
|
||||
format!("{chunk}\n\n(continues...)")
|
||||
} else if index == chunks.len() - 1 {
|
||||
format!("(continued)\n\n{chunk}")
|
||||
} else {
|
||||
format!("(continued)\n\n{chunk}\n\n(continues...)")
|
||||
}
|
||||
} else {
|
||||
chunk.to_string()
|
||||
};
|
||||
|
||||
let mut markdown_body = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"text": text,
|
||||
"parse_mode": "Markdown"
|
||||
});
|
||||
|
||||
// Add message_thread_id for forum topic support
|
||||
if let Some(tid) = thread_id {
|
||||
markdown_body["message_thread_id"] = serde_json::Value::String(tid.to_string());
|
||||
}
|
||||
if index == 0 {
|
||||
if let Some(parent_id) = reply_to_message_id {
|
||||
markdown_body["reply_to_message_id"] = serde_json::Value::from(parent_id);
|
||||
}
|
||||
}
|
||||
|
||||
let markdown_resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendMessage"))
|
||||
.json(&markdown_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if markdown_resp.status().is_success() {
|
||||
if index < chunks.len() - 1 {
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let markdown_status = markdown_resp.status();
|
||||
let markdown_err = markdown_resp.text().await.unwrap_or_default();
|
||||
tracing::warn!(
|
||||
status = ?markdown_status,
|
||||
"Telegram sendMessage with Markdown failed; retrying without parse_mode"
|
||||
);
|
||||
|
||||
let mut plain_body = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"text": text,
|
||||
});
|
||||
|
||||
// Add message_thread_id for forum topic support
|
||||
if let Some(tid) = thread_id {
|
||||
plain_body["message_thread_id"] = serde_json::Value::String(tid.to_string());
|
||||
}
|
||||
if index == 0 {
|
||||
if let Some(parent_id) = reply_to_message_id {
|
||||
plain_body["reply_to_message_id"] = serde_json::Value::from(parent_id);
|
||||
}
|
||||
}
|
||||
let plain_resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendMessage"))
|
||||
.json(&plain_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !plain_resp.status().is_success() {
|
||||
let plain_status = plain_resp.status();
|
||||
let plain_err = plain_resp.text().await.unwrap_or_default();
|
||||
anyhow::bail!(
|
||||
"Telegram sendMessage failed (markdown {}: {}; plain {}: {})",
|
||||
markdown_status,
|
||||
markdown_err,
|
||||
plain_status,
|
||||
plain_err
|
||||
);
|
||||
}
|
||||
|
||||
if index < chunks.len() - 1 {
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_media_by_url(
|
||||
&self,
|
||||
method: &str,
|
||||
media_field: &str,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
url: &str,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut body = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
});
|
||||
body[media_field] = serde_json::Value::String(url.to_string());
|
||||
|
||||
if let Some(tid) = thread_id {
|
||||
body["message_thread_id"] = serde_json::Value::String(tid.to_string());
|
||||
}
|
||||
|
||||
if let Some(cap) = caption {
|
||||
body["caption"] = serde_json::Value::String(cap.to_string());
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url(method))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await?;
|
||||
anyhow::bail!("Telegram {method} by URL failed: {err}");
|
||||
}
|
||||
|
||||
tracing::info!("Telegram {method} sent to {chat_id}: {url}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn send_attachment(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
attachment: &TelegramAttachment,
|
||||
) -> anyhow::Result<()> {
|
||||
let target = attachment.target.trim();
|
||||
|
||||
if is_http_url(target) {
|
||||
return match attachment.kind {
|
||||
TelegramAttachmentKind::Image => {
|
||||
self.send_photo_by_url(chat_id, thread_id, target, None)
|
||||
.await
|
||||
}
|
||||
TelegramAttachmentKind::Document => {
|
||||
self.send_document_by_url(chat_id, thread_id, target, None)
|
||||
.await
|
||||
}
|
||||
TelegramAttachmentKind::Video => {
|
||||
self.send_video_by_url(chat_id, thread_id, target, None)
|
||||
.await
|
||||
}
|
||||
TelegramAttachmentKind::Audio => {
|
||||
self.send_audio_by_url(chat_id, thread_id, target, None)
|
||||
.await
|
||||
}
|
||||
TelegramAttachmentKind::Voice => {
|
||||
self.send_voice_by_url(chat_id, thread_id, target, None)
|
||||
.await
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
let path = Path::new(target);
|
||||
if !path.exists() {
|
||||
anyhow::bail!("Telegram attachment path not found: {target}");
|
||||
}
|
||||
|
||||
match attachment.kind {
|
||||
TelegramAttachmentKind::Image => self.send_photo(chat_id, thread_id, path, None).await,
|
||||
TelegramAttachmentKind::Document => {
|
||||
self.send_document(chat_id, thread_id, path, None).await
|
||||
}
|
||||
TelegramAttachmentKind::Video => self.send_video(chat_id, thread_id, path, None).await,
|
||||
TelegramAttachmentKind::Audio => self.send_audio(chat_id, thread_id, path, None).await,
|
||||
TelegramAttachmentKind::Voice => self.send_voice(chat_id, thread_id, path, None).await,
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a document/file to a Telegram chat
|
||||
pub async fn send_document(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
file_path: &Path,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let file_name = file_path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or("file");
|
||||
|
||||
let file_bytes = tokio::fs::read(file_path).await?;
|
||||
let part = Part::bytes(file_bytes).file_name(file_name.to_string());
|
||||
|
||||
let mut form = Form::new()
|
||||
.text("chat_id", chat_id.to_string())
|
||||
.part("document", part);
|
||||
|
||||
if let Some(tid) = thread_id {
|
||||
form = form.text("message_thread_id", tid.to_string());
|
||||
}
|
||||
|
||||
if let Some(cap) = caption {
|
||||
form = form.text("caption", cap.to_string());
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendDocument"))
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await?;
|
||||
anyhow::bail!("Telegram sendDocument failed: {err}");
|
||||
}
|
||||
|
||||
tracing::info!("Telegram document sent to {chat_id}: {file_name}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send a document from bytes (in-memory) to a Telegram chat
|
||||
pub async fn send_document_bytes(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
file_bytes: Vec<u8>,
|
||||
file_name: &str,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let part = Part::bytes(file_bytes).file_name(file_name.to_string());
|
||||
|
||||
let mut form = Form::new()
|
||||
.text("chat_id", chat_id.to_string())
|
||||
.part("document", part);
|
||||
|
||||
if let Some(tid) = thread_id {
|
||||
form = form.text("message_thread_id", tid.to_string());
|
||||
}
|
||||
|
||||
if let Some(cap) = caption {
|
||||
form = form.text("caption", cap.to_string());
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendDocument"))
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await?;
|
||||
anyhow::bail!("Telegram sendDocument failed: {err}");
|
||||
}
|
||||
|
||||
tracing::info!("Telegram document sent to {chat_id}: {file_name}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send a photo to a Telegram chat
|
||||
pub async fn send_photo(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
file_path: &Path,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let file_name = file_path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or("photo.jpg");
|
||||
|
||||
let file_bytes = tokio::fs::read(file_path).await?;
|
||||
let part = Part::bytes(file_bytes).file_name(file_name.to_string());
|
||||
|
||||
let mut form = Form::new()
|
||||
.text("chat_id", chat_id.to_string())
|
||||
.part("photo", part);
|
||||
|
||||
if let Some(tid) = thread_id {
|
||||
form = form.text("message_thread_id", tid.to_string());
|
||||
}
|
||||
|
||||
if let Some(cap) = caption {
|
||||
form = form.text("caption", cap.to_string());
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendPhoto"))
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await?;
|
||||
anyhow::bail!("Telegram sendPhoto failed: {err}");
|
||||
}
|
||||
|
||||
tracing::info!("Telegram photo sent to {chat_id}: {file_name}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send a photo from bytes (in-memory) to a Telegram chat
|
||||
pub async fn send_photo_bytes(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
file_bytes: Vec<u8>,
|
||||
file_name: &str,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let part = Part::bytes(file_bytes).file_name(file_name.to_string());
|
||||
|
||||
let mut form = Form::new()
|
||||
.text("chat_id", chat_id.to_string())
|
||||
.part("photo", part);
|
||||
|
||||
if let Some(tid) = thread_id {
|
||||
form = form.text("message_thread_id", tid.to_string());
|
||||
}
|
||||
|
||||
if let Some(cap) = caption {
|
||||
form = form.text("caption", cap.to_string());
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendPhoto"))
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await?;
|
||||
anyhow::bail!("Telegram sendPhoto failed: {err}");
|
||||
}
|
||||
|
||||
tracing::info!("Telegram photo sent to {chat_id}: {file_name}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send a video to a Telegram chat
|
||||
pub async fn send_video(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
file_path: &Path,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let file_name = file_path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or("video.mp4");
|
||||
|
||||
let file_bytes = tokio::fs::read(file_path).await?;
|
||||
let part = Part::bytes(file_bytes).file_name(file_name.to_string());
|
||||
|
||||
let mut form = Form::new()
|
||||
.text("chat_id", chat_id.to_string())
|
||||
.part("video", part);
|
||||
|
||||
if let Some(tid) = thread_id {
|
||||
form = form.text("message_thread_id", tid.to_string());
|
||||
}
|
||||
|
||||
if let Some(cap) = caption {
|
||||
form = form.text("caption", cap.to_string());
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendVideo"))
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await?;
|
||||
anyhow::bail!("Telegram sendVideo failed: {err}");
|
||||
}
|
||||
|
||||
tracing::info!("Telegram video sent to {chat_id}: {file_name}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send an audio file to a Telegram chat
|
||||
pub async fn send_audio(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
file_path: &Path,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let file_name = file_path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or("audio.mp3");
|
||||
|
||||
let file_bytes = tokio::fs::read(file_path).await?;
|
||||
let part = Part::bytes(file_bytes).file_name(file_name.to_string());
|
||||
|
||||
let mut form = Form::new()
|
||||
.text("chat_id", chat_id.to_string())
|
||||
.part("audio", part);
|
||||
|
||||
if let Some(tid) = thread_id {
|
||||
form = form.text("message_thread_id", tid.to_string());
|
||||
}
|
||||
|
||||
if let Some(cap) = caption {
|
||||
form = form.text("caption", cap.to_string());
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendAudio"))
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await?;
|
||||
anyhow::bail!("Telegram sendAudio failed: {err}");
|
||||
}
|
||||
|
||||
tracing::info!("Telegram audio sent to {chat_id}: {file_name}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send a voice message to a Telegram chat
|
||||
pub async fn send_voice(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
file_path: &Path,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let file_name = file_path
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or("voice.ogg");
|
||||
|
||||
let file_bytes = tokio::fs::read(file_path).await?;
|
||||
let part = Part::bytes(file_bytes).file_name(file_name.to_string());
|
||||
|
||||
let mut form = Form::new()
|
||||
.text("chat_id", chat_id.to_string())
|
||||
.part("voice", part);
|
||||
|
||||
if let Some(tid) = thread_id {
|
||||
form = form.text("message_thread_id", tid.to_string());
|
||||
}
|
||||
|
||||
if let Some(cap) = caption {
|
||||
form = form.text("caption", cap.to_string());
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendVoice"))
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await?;
|
||||
anyhow::bail!("Telegram sendVoice failed: {err}");
|
||||
}
|
||||
|
||||
tracing::info!("Telegram voice sent to {chat_id}: {file_name}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send a file by URL (Telegram will download it)
|
||||
pub async fn send_document_by_url(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
url: &str,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut body = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"document": url
|
||||
});
|
||||
|
||||
if let Some(tid) = thread_id {
|
||||
body["message_thread_id"] = serde_json::Value::String(tid.to_string());
|
||||
}
|
||||
|
||||
if let Some(cap) = caption {
|
||||
body["caption"] = serde_json::Value::String(cap.to_string());
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendDocument"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await?;
|
||||
anyhow::bail!("Telegram sendDocument by URL failed: {err}");
|
||||
}
|
||||
|
||||
tracing::info!("Telegram document (URL) sent to {chat_id}: {url}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send a photo by URL (Telegram will download it)
|
||||
pub async fn send_photo_by_url(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
url: &str,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let mut body = serde_json::json!({
|
||||
"chat_id": chat_id,
|
||||
"photo": url
|
||||
});
|
||||
|
||||
if let Some(tid) = thread_id {
|
||||
body["message_thread_id"] = serde_json::Value::String(tid.to_string());
|
||||
}
|
||||
|
||||
if let Some(cap) = caption {
|
||||
body["caption"] = serde_json::Value::String(cap.to_string());
|
||||
}
|
||||
|
||||
let resp = self
|
||||
.http_client()
|
||||
.post(self.api_url("sendPhoto"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let err = resp.text().await?;
|
||||
anyhow::bail!("Telegram sendPhoto by URL failed: {err}");
|
||||
}
|
||||
|
||||
tracing::info!("Telegram photo (URL) sent to {chat_id}: {url}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Send a video by URL (Telegram will download it)
|
||||
pub async fn send_video_by_url(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
url: &str,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
self.send_media_by_url("sendVideo", "video", chat_id, thread_id, url, caption)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Send an audio file by URL (Telegram will download it)
|
||||
pub async fn send_audio_by_url(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
url: &str,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
self.send_media_by_url("sendAudio", "audio", chat_id, thread_id, url, caption)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Send a voice message by URL (Telegram will download it)
|
||||
pub async fn send_voice_by_url(
|
||||
&self,
|
||||
chat_id: &str,
|
||||
thread_id: Option<&str>,
|
||||
url: &str,
|
||||
caption: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
self.send_media_by_url("sendVoice", "voice", chat_id, thread_id, url, caption)
|
||||
.await
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,75 +0,0 @@
|
||||
//! Telegram channel — private types and the main struct definition.
|
||||
|
||||
use crate::openhuman::config::StreamMode;
|
||||
use crate::openhuman::security::pairing::PairingGuard;
|
||||
use parking_lot::Mutex;
|
||||
use std::collections::{HashSet, VecDeque};
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Instant;
|
||||
|
||||
pub(crate) const TELEGRAM_RECENT_UPDATE_CACHE_SIZE: usize = 4096;
|
||||
|
||||
/// Telegram Bot API caps downloaded files at 20MB for `getFile`.
|
||||
/// Keep our inbound voice path inside the same bound before base64 encoding
|
||||
/// and dispatching to STT.
|
||||
pub(crate) const TELEGRAM_MAX_VOICE_FILE_BYTES: u64 = 20 * 1024 * 1024;
|
||||
|
||||
/// De-bounce window for approval prompts: suppress duplicate prompts sent to the
|
||||
/// same chat+sender within this duration (prevents restart-race and rapid-fire spam).
|
||||
pub(crate) const APPROVAL_PROMPT_DEBOUNCE_SECS: u64 = 60;
|
||||
|
||||
pub(crate) struct TelegramTypingTask {
|
||||
pub(crate) recipient: String,
|
||||
pub(crate) handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct TelegramUpdateWindow {
|
||||
pub(crate) max_seen_update_id: i64,
|
||||
pub(crate) recent_order: VecDeque<i64>,
|
||||
pub(crate) recent_lookup: HashSet<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct TelegramReactionEvent {
|
||||
pub(crate) sender: String,
|
||||
pub(crate) reply_target: String,
|
||||
pub(crate) target_message_id: String,
|
||||
pub(crate) emoji: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct TelegramVoiceAttachment {
|
||||
pub(crate) file_id: String,
|
||||
pub(crate) file_unique_id: Option<String>,
|
||||
pub(crate) file_size: Option<u64>,
|
||||
pub(crate) mime_type: Option<String>,
|
||||
}
|
||||
|
||||
/// Telegram channel — long-polls the Bot API for updates
|
||||
pub struct TelegramChannel {
|
||||
pub(crate) bot_token: String,
|
||||
/// Default chat for recipient-less proactive sends. `None` ⇒ proactive
|
||||
/// routing skips Telegram (see `proactive_target`). Set from
|
||||
/// `TelegramConfig::chat_id` via [`TelegramChannel::with_chat_id`].
|
||||
pub(crate) chat_id: Option<String>,
|
||||
/// Base URL for the Telegram Bot API. Defaults to `https://api.telegram.org`.
|
||||
/// Override via `OPENHUMAN_TELEGRAM_BOT_API_BASE` for E2E testing against a
|
||||
/// mock server. The legacy `OPENHUMAN_TELEGRAM_API_BASE` alias is still accepted.
|
||||
pub(crate) api_base: String,
|
||||
pub(crate) allowed_users: Arc<RwLock<Vec<String>>>,
|
||||
pub(crate) pairing: Option<PairingGuard>,
|
||||
pub(crate) client: reqwest::Client,
|
||||
pub(crate) typing_handle: Mutex<Option<TelegramTypingTask>>,
|
||||
pub(crate) stream_mode: StreamMode,
|
||||
pub(crate) draft_update_interval_ms: u64,
|
||||
pub(crate) silent_streaming: bool,
|
||||
pub(crate) last_draft_edit: Mutex<std::collections::HashMap<String, Instant>>,
|
||||
pub(crate) mention_only: bool,
|
||||
pub(crate) bot_username: Mutex<Option<String>>,
|
||||
pub(crate) recent_updates: Mutex<TelegramUpdateWindow>,
|
||||
/// Tracks the last time an approval prompt was sent to a given "chat_id:sender" key.
|
||||
/// Prevents duplicate prompts during restart-overlap races and rapid re-sends.
|
||||
/// Mirrors the `last_draft_edit` pattern.
|
||||
pub(crate) recent_approval_prompts: Mutex<std::collections::HashMap<String, Instant>>,
|
||||
}
|
||||
@@ -1,32 +1,31 @@
|
||||
//! Telegram channel — long-polls the Bot API for updates.
|
||||
//! Telegram channel — host-side glue.
|
||||
//!
|
||||
//! The transport (Bot API driver, session store, pairing) now lives in
|
||||
//! `tinychannels::providers::telegram`; re-export it here so existing paths
|
||||
//! keep resolving. What stays are the host-coupled pieces that depend on the
|
||||
//! OpenHuman event bus, runtime context, and cross-channel calls:
|
||||
//! - [`remote_control`] — `/status /sessions /new` command handling (uses the
|
||||
//! agent runtime context + web session invalidation).
|
||||
//! - [`bus`] — the `TelegramRemoteSubscriber` busy-state event handler.
|
||||
//! - [`approval_surface`] — the `TelegramApprovalSurfaceSubscriber`.
|
||||
|
||||
mod approval_surface;
|
||||
mod attachments;
|
||||
mod bus;
|
||||
mod channel;
|
||||
mod channel_core;
|
||||
mod channel_ops;
|
||||
mod channel_recv;
|
||||
mod channel_send;
|
||||
mod channel_types;
|
||||
pub mod remote_control;
|
||||
mod session_store;
|
||||
mod text;
|
||||
|
||||
// Transport moved to tinychannels.
|
||||
pub use tinychannels::providers::telegram::{session_store, TelegramChannel};
|
||||
|
||||
pub use approval_surface::{TelegramApprovalSurfaceSubscriber, TELEGRAM_APPROVAL_CLIENT_ID};
|
||||
pub use bus::TelegramRemoteSubscriber;
|
||||
pub use channel_types::TelegramChannel;
|
||||
pub use remote_control::TelegramRemoteCommand;
|
||||
|
||||
#[cfg(any(test, debug_assertions))]
|
||||
pub mod test_support {
|
||||
//! Debug-build seams for raw integration coverage of Telegram send helpers.
|
||||
//! Delegates to the tinychannels transport crate where the logic now lives.
|
||||
|
||||
use super::TelegramChannel;
|
||||
|
||||
pub fn parse_reaction_marker_for_test(content: &str) -> (String, Option<String>) {
|
||||
TelegramChannel::parse_reaction_marker(content)
|
||||
}
|
||||
pub use tinychannels::providers::telegram::test_support::parse_reaction_marker_for_test;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1,236 +0,0 @@
|
||||
//! Workspace-backed Telegram chat → thread bindings for remote control.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
const STORE_FILE: &str = "state/telegram_remote_sessions.json";
|
||||
const LOG_PREFIX: &str = "[telegram-remote]";
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub(crate) struct TelegramChatBinding {
|
||||
pub(crate) thread_id: String,
|
||||
pub(crate) sender_key: String,
|
||||
pub(crate) updated_at: String,
|
||||
/// Human-readable title captured at `/new` time so `/status` can display it
|
||||
/// without listing all threads (O(1) instead of O(n) disk reads).
|
||||
#[serde(default)]
|
||||
pub(crate) title: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
struct TelegramSessionStoreFile {
|
||||
bindings: HashMap<String, TelegramChatBinding>,
|
||||
#[serde(default)]
|
||||
busy_reply_targets: HashMap<String, bool>,
|
||||
}
|
||||
|
||||
pub(crate) struct TelegramSessionStore {
|
||||
file: TelegramSessionStoreFile,
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
impl TelegramSessionStore {
|
||||
pub(crate) fn load(workspace_dir: &Path) -> anyhow::Result<Self> {
|
||||
let path = workspace_dir.join(STORE_FILE);
|
||||
let file = if path.exists() {
|
||||
let raw = std::fs::read_to_string(&path)?;
|
||||
serde_json::from_str(&raw).unwrap_or_else(|error| {
|
||||
tracing::warn!(
|
||||
"{LOG_PREFIX} corrupt session store at {}: {error}; resetting",
|
||||
path.display()
|
||||
);
|
||||
TelegramSessionStoreFile::default()
|
||||
})
|
||||
} else {
|
||||
TelegramSessionStoreFile::default()
|
||||
};
|
||||
tracing::debug!(
|
||||
"{LOG_PREFIX} loaded session store bindings={} busy={}",
|
||||
file.bindings.len(),
|
||||
file.busy_reply_targets.len()
|
||||
);
|
||||
Ok(Self { file, path })
|
||||
}
|
||||
|
||||
pub(crate) fn save(&self) -> anyhow::Result<()> {
|
||||
if let Some(parent) = self.path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
let raw = serde_json::to_string_pretty(&self.file)?;
|
||||
std::fs::write(&self.path, raw)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn binding(&self, reply_target: &str) -> Option<&TelegramChatBinding> {
|
||||
self.file.bindings.get(reply_target)
|
||||
}
|
||||
|
||||
pub(crate) fn set_binding(
|
||||
&mut self,
|
||||
reply_target: &str,
|
||||
thread_id: String,
|
||||
sender_key: String,
|
||||
title: Option<String>,
|
||||
) {
|
||||
let updated_at = chrono::Utc::now().to_rfc3339();
|
||||
self.file.bindings.insert(
|
||||
reply_target.to_string(),
|
||||
TelegramChatBinding {
|
||||
thread_id,
|
||||
sender_key,
|
||||
updated_at,
|
||||
title,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub(crate) fn set_busy(&mut self, reply_target: &str, busy: bool) {
|
||||
if busy {
|
||||
self.file
|
||||
.busy_reply_targets
|
||||
.insert(reply_target.to_string(), true);
|
||||
} else {
|
||||
self.file.busy_reply_targets.remove(reply_target);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn is_busy(&self, reply_target: &str) -> bool {
|
||||
self.file
|
||||
.busy_reply_targets
|
||||
.get(reply_target)
|
||||
.copied()
|
||||
.unwrap_or(false)
|
||||
}
|
||||
}
|
||||
|
||||
static STORE: std::sync::OnceLock<std::sync::Mutex<Option<TelegramSessionStore>>> =
|
||||
std::sync::OnceLock::new();
|
||||
|
||||
/// Read-write accessor: runs `f` against the cached store, then flushes to disk.
|
||||
/// Use for operations that mutate state (e.g. `set_binding`, `set_busy`).
|
||||
pub(crate) fn with_store<F, R>(workspace_dir: &Path, f: F) -> anyhow::Result<R>
|
||||
where
|
||||
F: FnOnce(&mut TelegramSessionStore) -> anyhow::Result<R>,
|
||||
{
|
||||
let lock = STORE.get_or_init(|| std::sync::Mutex::new(None));
|
||||
let mut guard = lock.lock().expect("telegram session store mutex poisoned");
|
||||
let expected_path = workspace_dir.join(STORE_FILE);
|
||||
let needs_load = guard
|
||||
.as_ref()
|
||||
.map(|store| store.path != expected_path)
|
||||
.unwrap_or(true);
|
||||
if needs_load {
|
||||
*guard = Some(TelegramSessionStore::load(workspace_dir)?);
|
||||
}
|
||||
let store = guard.as_mut().expect("store initialized");
|
||||
let result = f(store)?;
|
||||
store.save()?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Read-only accessor: runs `f` against the cached store but does **not** flush
|
||||
/// to disk. Use for operations that only read state (e.g. `binding`, `is_busy`)
|
||||
/// to avoid unnecessary serialization and disk I/O on every query.
|
||||
pub(crate) fn with_store_read<F, R>(workspace_dir: &Path, f: F) -> anyhow::Result<R>
|
||||
where
|
||||
F: FnOnce(&TelegramSessionStore) -> anyhow::Result<R>,
|
||||
{
|
||||
let lock = STORE.get_or_init(|| std::sync::Mutex::new(None));
|
||||
let mut guard = lock.lock().expect("telegram session store mutex poisoned");
|
||||
let expected_path = workspace_dir.join(STORE_FILE);
|
||||
let needs_load = guard
|
||||
.as_ref()
|
||||
.map(|store| store.path != expected_path)
|
||||
.unwrap_or(true);
|
||||
if needs_load {
|
||||
*guard = Some(TelegramSessionStore::load(workspace_dir)?);
|
||||
}
|
||||
let store = guard.as_ref().expect("store initialized");
|
||||
f(store)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[test]
|
||||
fn round_trip_binding_and_busy_flag() {
|
||||
let dir = tempdir().expect("tempdir");
|
||||
let mut store = TelegramSessionStore::load(dir.path()).expect("load");
|
||||
store.set_binding(
|
||||
"12345",
|
||||
"thread-abc".into(),
|
||||
"telegram_alice_12345".into(),
|
||||
Some("My Session".into()),
|
||||
);
|
||||
store.set_busy("12345", true);
|
||||
store.save().expect("save");
|
||||
|
||||
let reloaded = TelegramSessionStore::load(dir.path()).expect("reload");
|
||||
let binding = reloaded.binding("12345").expect("binding");
|
||||
assert_eq!(binding.thread_id, "thread-abc");
|
||||
assert_eq!(binding.title.as_deref(), Some("My Session"));
|
||||
assert!(reloaded.is_busy("12345"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn corrupt_store_resets_and_clearing_busy_removes_flag() {
|
||||
let dir = tempdir().expect("tempdir");
|
||||
let path = dir.path().join(STORE_FILE);
|
||||
std::fs::create_dir_all(path.parent().expect("state dir")).expect("state dir");
|
||||
std::fs::write(&path, "{ not valid json").expect("write corrupt store");
|
||||
|
||||
let mut store = TelegramSessionStore::load(dir.path()).expect("load corrupt store");
|
||||
assert!(store.binding("12345").is_none());
|
||||
|
||||
store.set_busy("12345", true);
|
||||
assert!(store.is_busy("12345"));
|
||||
store.set_busy("12345", false);
|
||||
assert!(!store.is_busy("12345"));
|
||||
store.save().expect("save reset store");
|
||||
|
||||
let raw = std::fs::read_to_string(path).expect("read saved store");
|
||||
assert!(raw.contains("\"bindings\""));
|
||||
assert!(!raw.contains("12345"));
|
||||
}
|
||||
|
||||
/// Tests `with_store` workspace-change detection by using `TelegramSessionStore`
|
||||
/// directly for cross-workspace assertions — avoids races with the process-global
|
||||
/// `STORE` singleton when tests run in parallel.
|
||||
#[test]
|
||||
fn store_isolates_bindings_across_workspaces() {
|
||||
let first = tempdir().expect("first tempdir");
|
||||
let second = tempdir().expect("second tempdir");
|
||||
|
||||
// Write binding into first workspace directly (no global singleton).
|
||||
let mut store_a = TelegramSessionStore::load(first.path()).expect("load first");
|
||||
store_a.set_binding("chat-a", "thread-a".into(), "telegram_a".into(), None);
|
||||
store_a.save().expect("save first");
|
||||
|
||||
// Write binding into second workspace directly.
|
||||
let mut store_b = TelegramSessionStore::load(second.path()).expect("load second");
|
||||
assert!(
|
||||
store_b.binding("chat-a").is_none(),
|
||||
"second workspace must not see first workspace's binding"
|
||||
);
|
||||
store_b.set_binding("chat-b", "thread-b".into(), "telegram_b".into(), None);
|
||||
store_b.save().expect("save second");
|
||||
|
||||
let first_store = TelegramSessionStore::load(first.path()).expect("reload first");
|
||||
let second_store = TelegramSessionStore::load(second.path()).expect("reload second");
|
||||
assert_eq!(
|
||||
first_store
|
||||
.binding("chat-a")
|
||||
.map(|binding| binding.thread_id.as_str()),
|
||||
Some("thread-a")
|
||||
);
|
||||
assert_eq!(
|
||||
second_store
|
||||
.binding("chat-b")
|
||||
.map(|binding| binding.thread_id.as_str()),
|
||||
Some("thread-b")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,131 +0,0 @@
|
||||
//! Text chunking and tool-call tag stripping for Telegram.
|
||||
|
||||
use tinychannels::channel::LengthUnit;
|
||||
use tinychannels::text::{chunk_text_with_options, ChunkMode, TextChunkOptions};
|
||||
|
||||
/// Telegram's maximum message length for text messages
|
||||
pub(crate) const TELEGRAM_MAX_MESSAGE_LENGTH: usize = 4096;
|
||||
pub(crate) const TELEGRAM_BIND_COMMAND: &str = "/bind";
|
||||
pub(crate) const TELEGRAM_START_COMMAND: &str = "/start";
|
||||
|
||||
pub(crate) fn split_message_for_telegram(message: &str) -> Vec<String> {
|
||||
if message.is_empty() {
|
||||
return vec![String::new()];
|
||||
}
|
||||
chunk_text_with_options(
|
||||
message,
|
||||
TextChunkOptions {
|
||||
limit: TELEGRAM_MAX_MESSAGE_LENGTH,
|
||||
length_unit: LengthUnit::Utf16Units,
|
||||
mode: ChunkMode::Length,
|
||||
markdown: true,
|
||||
indicators: false,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn strip_tool_call_tags(message: &str) -> String {
|
||||
const TOOL_CALL_OPEN_TAGS: [&str; 5] = [
|
||||
"<tool_call>",
|
||||
"<toolcall>",
|
||||
"<tool-call>",
|
||||
"<tool>",
|
||||
"<invoke>",
|
||||
];
|
||||
|
||||
fn find_first_tag<'a>(haystack: &str, tags: &'a [&'a str]) -> Option<(usize, &'a str)> {
|
||||
tags.iter()
|
||||
.filter_map(|tag| haystack.find(tag).map(|idx| (idx, *tag)))
|
||||
.min_by_key(|(idx, _)| *idx)
|
||||
}
|
||||
|
||||
fn matching_close_tag(open_tag: &str) -> Option<&'static str> {
|
||||
match open_tag {
|
||||
"<tool_call>" => Some("</tool_call>"),
|
||||
"<toolcall>" => Some("</toolcall>"),
|
||||
"<tool-call>" => Some("</tool-call>"),
|
||||
"<tool>" => Some("</tool>"),
|
||||
"<invoke>" => Some("</invoke>"),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_first_json_end(input: &str) -> Option<usize> {
|
||||
let trimmed = input.trim_start();
|
||||
let trim_offset = input.len().saturating_sub(trimmed.len());
|
||||
|
||||
for (byte_idx, ch) in trimmed.char_indices() {
|
||||
if ch != '{' && ch != '[' {
|
||||
continue;
|
||||
}
|
||||
|
||||
let slice = &trimmed[byte_idx..];
|
||||
let mut stream =
|
||||
serde_json::Deserializer::from_str(slice).into_iter::<serde_json::Value>();
|
||||
if let Some(Ok(_value)) = stream.next() {
|
||||
let consumed = stream.byte_offset();
|
||||
if consumed > 0 {
|
||||
return Some(trim_offset + byte_idx + consumed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn strip_leading_close_tags(mut input: &str) -> &str {
|
||||
loop {
|
||||
let trimmed = input.trim_start();
|
||||
if !trimmed.starts_with("</") {
|
||||
return trimmed;
|
||||
}
|
||||
|
||||
let Some(close_end) = trimmed.find('>') else {
|
||||
return "";
|
||||
};
|
||||
input = &trimmed[close_end + 1..];
|
||||
}
|
||||
}
|
||||
|
||||
let mut kept_segments = Vec::new();
|
||||
let mut remaining = message;
|
||||
|
||||
while let Some((start, open_tag)) = find_first_tag(remaining, &TOOL_CALL_OPEN_TAGS) {
|
||||
let before = &remaining[..start];
|
||||
if !before.is_empty() {
|
||||
kept_segments.push(before.to_string());
|
||||
}
|
||||
|
||||
let Some(close_tag) = matching_close_tag(open_tag) else {
|
||||
break;
|
||||
};
|
||||
let after_open = &remaining[start + open_tag.len()..];
|
||||
|
||||
if let Some(close_idx) = after_open.find(close_tag) {
|
||||
remaining = &after_open[close_idx + close_tag.len()..];
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Some(consumed_end) = extract_first_json_end(after_open) {
|
||||
remaining = strip_leading_close_tags(&after_open[consumed_end..]);
|
||||
continue;
|
||||
}
|
||||
|
||||
kept_segments.push(remaining[start..].to_string());
|
||||
remaining = "";
|
||||
break;
|
||||
}
|
||||
|
||||
if !remaining.is_empty() {
|
||||
kept_segments.push(remaining.to_string());
|
||||
}
|
||||
|
||||
let mut result = kept_segments.concat();
|
||||
|
||||
// Clean up any resulting blank lines (but preserve paragraphs)
|
||||
while result.contains("\n\n\n") {
|
||||
result = result.replace("\n\n\n", "\n\n");
|
||||
}
|
||||
|
||||
result.trim().to_string()
|
||||
}
|
||||
@@ -1,5 +1,8 @@
|
||||
mod event_bus;
|
||||
mod ops;
|
||||
// Response delivery/segmentation for the web surface (folded in from the former
|
||||
// standalone `presentation` provider — it is the web channel's delivery formatter).
|
||||
pub mod presentation;
|
||||
mod progress_bridge;
|
||||
mod run_task;
|
||||
mod schemas;
|
||||
|
||||
@@ -603,7 +603,7 @@ pub async fn start_chat(
|
||||
|
||||
match result {
|
||||
Ok(chat_result) => {
|
||||
crate::openhuman::channels::providers::presentation::deliver_response(
|
||||
crate::openhuman::channels::providers::web::presentation::deliver_response(
|
||||
&client_id_task,
|
||||
&thread_id_task,
|
||||
&request_id_task,
|
||||
@@ -809,7 +809,7 @@ async fn spawn_parallel_turn(
|
||||
|
||||
match result {
|
||||
Some(Ok(chat_result)) => {
|
||||
crate::openhuman::channels::providers::presentation::deliver_response(
|
||||
crate::openhuman::channels::providers::web::presentation::deliver_response(
|
||||
&client_id_task,
|
||||
&thread_id_task,
|
||||
&request_id_task,
|
||||
|
||||
+1
-1
@@ -13,7 +13,7 @@ use crate::core::socketio::{SubagentUsagePayload, TurnUsagePayload, WebChannelEv
|
||||
use crate::openhuman::agent::harness::turn_subagent_usage::LastTurnUsage;
|
||||
use crate::openhuman::config::rpc as config_rpc;
|
||||
|
||||
use super::web::publish_web_channel_event;
|
||||
use super::publish_web_channel_event;
|
||||
|
||||
const MIN_SEGMENT_CHARS: usize = 40;
|
||||
const MAX_SEGMENTS: usize = 5;
|
||||
@@ -1,590 +1 @@
|
||||
//! WhatsApp Web channel backed by upstream [`whatsapp-rust`] 0.5.
|
||||
//!
|
||||
//! # Why the upgrade
|
||||
//!
|
||||
//! The previous implementation used `wa-rs` 0.2 (a fork that pinned to stable
|
||||
//! Rust). That fork silently dropped `Event::Message` for LID-addressed
|
||||
//! contacts and group sender-key (`skmsg`) messages: the protocol layer
|
||||
//! decrypted the payload but never dispatched it to user code, breaking
|
||||
//! agent dispatch for the bulk of modern WhatsApp traffic (LID is the
|
||||
//! current default). Upstream `whatsapp-rust` 0.5 fixed this in PRs #170
|
||||
//! (SKDM tracking) + #181 (LID/PN mapping) + sender-key dispatch.
|
||||
//!
|
||||
//! # Feature Flag
|
||||
//!
|
||||
//! ```sh
|
||||
//! cargo build --features whatsapp-web
|
||||
//! ```
|
||||
//!
|
||||
//! # Configuration
|
||||
//!
|
||||
//! ```toml
|
||||
//! [channels.whatsapp]
|
||||
//! session_path = "~/.openhuman/whatsapp-session.db" # Reserved for durable Web mode
|
||||
//! pair_phone = "15551234567" # Optional: pair-code linking
|
||||
//! allowed_numbers = ["+1234567890", "*"] # Same shape as Cloud API
|
||||
//! ```
|
||||
//!
|
||||
//! # Runtime negotiation
|
||||
//!
|
||||
//! Selected automatically by [`crate::openhuman::channels::runtime::startup`]
|
||||
//! when `session_path` is set. The Cloud API channel ([`super::whatsapp`]) is
|
||||
//! used when `phone_number_id` is set instead.
|
||||
//!
|
||||
//! # Migration note
|
||||
//!
|
||||
//! The upstream 0.5 `sqlite-storage` feature currently uses Diesel, whose
|
||||
//! native sqlite binding conflicts with the TinyAgents 1.3 / rusqlite 0.40
|
||||
//! baseline. Until OpenHuman owns a rusqlite-backed durable store for the
|
||||
//! WhatsApp backend traits, this channel uses wacore's in-memory backend and
|
||||
//! requires re-linking after restart.
|
||||
//!
|
||||
//! [`whatsapp-rust`]: https://docs.rs/whatsapp-rust/0.5
|
||||
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
use anyhow::{anyhow, Result};
|
||||
use async_trait::async_trait;
|
||||
use parking_lot::Mutex;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
|
||||
/// WhatsApp Web channel.
|
||||
///
|
||||
/// Wraps a `whatsapp-rust` Bot with our `Channel` trait. The bot owns an
|
||||
/// `Arc<Client>` for outbound operations (`send`, typing) and a `BotHandle`
|
||||
/// for shutdown. Inbound messages are pushed onto an [`mpsc::Sender`] so
|
||||
/// the existing channel inbound subscriber pipeline can process them.
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
pub struct WhatsAppWebChannel {
|
||||
/// Path to the SQLite session database.
|
||||
session_path: String,
|
||||
/// Optional phone number for pair-code linking (E.164 digits, no leading `+`).
|
||||
pair_phone: Option<String>,
|
||||
/// Optional pre-allocated pair code paired with `pair_phone`.
|
||||
pair_code: Option<String>,
|
||||
/// E.164 numbers (with leading `+`) allowed to interact, or `["*"]` for any.
|
||||
/// Empty also means "allow all" — same convention as the Cloud API channel.
|
||||
allowed_numbers: Vec<String>,
|
||||
/// Bot run handle, retained for graceful shutdown.
|
||||
bot_handle: Arc<Mutex<Option<whatsapp_rust::bot::BotHandle>>>,
|
||||
/// Live client used for outbound calls; populated after `Bot::build` returns.
|
||||
client: Arc<Mutex<Option<Arc<whatsapp_rust::Client>>>>,
|
||||
/// Liveness signal driven by upstream `Event::Connected` / `LoggedOut` /
|
||||
/// `StreamError`. Used by `health_check` so a dropped session no longer
|
||||
/// reports healthy until process shutdown.
|
||||
connected: Arc<AtomicBool>,
|
||||
/// Group JIDs (`...@g.us`) we've already accepted an allowed inbound
|
||||
/// from. Acts as outbound provenance: replies into a group are only
|
||||
/// permitted after a participant on the per-number allowlist messaged
|
||||
/// in. Without this, any caller able to pass a `recipient` could post
|
||||
/// into arbitrary joined groups via the @g.us suffix.
|
||||
allowed_groups: Arc<Mutex<HashSet<String>>>,
|
||||
/// Sink for inbound `ChannelMessage`s. Populated when [`Channel::listen`]
|
||||
/// is called and shared with the event-handler closure.
|
||||
tx: Arc<Mutex<Option<tokio::sync::mpsc::Sender<ChannelMessage>>>>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
impl WhatsAppWebChannel {
|
||||
/// Construct a channel. The bot does not connect until [`Channel::listen`]
|
||||
/// is invoked.
|
||||
pub fn new(
|
||||
session_path: String,
|
||||
pair_phone: Option<String>,
|
||||
pair_code: Option<String>,
|
||||
allowed_numbers: Vec<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
session_path,
|
||||
pair_phone,
|
||||
pair_code,
|
||||
allowed_numbers,
|
||||
bot_handle: Arc::new(Mutex::new(None)),
|
||||
client: Arc::new(Mutex::new(None)),
|
||||
connected: Arc::new(AtomicBool::new(false)),
|
||||
allowed_groups: Arc::new(Mutex::new(HashSet::new())),
|
||||
tx: Arc::new(Mutex::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Allowlist check. Empty list ⇒ allow-all (matches Cloud API behaviour).
|
||||
fn is_number_allowed(&self, phone: &str) -> bool {
|
||||
self.allowed_numbers.is_empty()
|
||||
|| self.allowed_numbers.iter().any(|n| n == "*" || n == phone)
|
||||
}
|
||||
|
||||
/// Recognise WhatsApp group JIDs (`...@g.us`). Group recipients bypass
|
||||
/// the per-number outbound allowlist because group membership is
|
||||
/// governed by WhatsApp itself; the inbound side already gated on the
|
||||
/// participant's allowlist status before we ever decided to reply.
|
||||
fn is_group_jid(recipient: &str) -> bool {
|
||||
recipient.trim().ends_with("@g.us")
|
||||
}
|
||||
|
||||
/// Outbound gate combining group-provenance with the per-number allowlist.
|
||||
/// Group JIDs are only permitted when an allowed inbound has already
|
||||
/// been received from that exact group — populated in the inbound
|
||||
/// handler when an allow-listed participant posts. This narrows the
|
||||
/// previous "all `@g.us` is fine" path so an attacker that can supply
|
||||
/// a `recipient` cannot post into arbitrary groups the bot has joined.
|
||||
fn should_allow_outbound(&self, recipient: &str) -> bool {
|
||||
if Self::is_group_jid(recipient) {
|
||||
return self.allowed_groups.lock().contains(recipient.trim());
|
||||
}
|
||||
let normalized = self.normalize_phone(recipient);
|
||||
self.is_number_allowed(&normalized)
|
||||
}
|
||||
|
||||
/// Mask a recipient identifier for log emission. Handles bare phone
|
||||
/// numbers, `<digits>@s.whatsapp.net`/`@lid` DM JIDs, and `@g.us`
|
||||
/// group JIDs uniformly so warning paths never carry a full ID.
|
||||
fn redact_recipient(recipient: &str) -> String {
|
||||
let trimmed = recipient.trim();
|
||||
if let Some((user, server)) = trimmed.split_once('@') {
|
||||
format!("{}@{}", Self::redact_phone(user), server)
|
||||
} else {
|
||||
Self::redact_phone(trimmed)
|
||||
}
|
||||
}
|
||||
|
||||
/// Pick the address downstream replies should be sent back to.
|
||||
///
|
||||
/// Group chats are addressed by the group JID (`...@g.us`); a reply that
|
||||
/// targeted the participant's phone instead would leak the conversation
|
||||
/// into a private DM.
|
||||
fn compute_reply_target(chat_jid: &str, sender_normalized: &str) -> String {
|
||||
if chat_jid.ends_with("@g.us") {
|
||||
chat_jid.to_string()
|
||||
} else {
|
||||
sender_normalized.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
/// Mask the middle digits of an E.164 number so logs only carry a coarse
|
||||
/// fingerprint instead of the full identifier.
|
||||
fn redact_phone(phone: &str) -> String {
|
||||
let prefix = if phone.starts_with('+') { "+" } else { "" };
|
||||
if phone.len() <= prefix.len() + 4 {
|
||||
return format!("{prefix}****");
|
||||
}
|
||||
let tail = &phone[phone.len() - 4..];
|
||||
format!("{prefix}***{tail}")
|
||||
}
|
||||
|
||||
/// Pull the displayable text out of an inbound WhatsApp Message proto.
|
||||
/// Falls back from `conversation` to `extended_text_message.text`, then
|
||||
/// to an empty string for non-text payloads.
|
||||
fn extract_message_text(conversation: Option<&str>, extended_text: Option<&str>) -> String {
|
||||
conversation
|
||||
.or(extended_text)
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Render an arbitrary recipient string as E.164 with a leading `+`,
|
||||
/// stripping any `@server` JID suffix the caller passed in.
|
||||
fn normalize_phone(&self, phone: &str) -> String {
|
||||
let trimmed = phone.trim();
|
||||
let user_part = trimmed
|
||||
.split_once('@')
|
||||
.map(|(user, _)| user)
|
||||
.unwrap_or(trimmed);
|
||||
let normalized_user = user_part.trim_start_matches('+');
|
||||
format!("+{normalized_user}")
|
||||
}
|
||||
|
||||
/// Convert a recipient (full JID like `12345@s.whatsapp.net` or an E.164
|
||||
/// number like `+1234567890`) into a `whatsapp-rust` JID.
|
||||
fn recipient_to_jid(&self, recipient: &str) -> Result<whatsapp_rust::Jid> {
|
||||
let trimmed = recipient.trim();
|
||||
if trimmed.is_empty() {
|
||||
anyhow::bail!("Recipient cannot be empty");
|
||||
}
|
||||
|
||||
if trimmed.contains('@') {
|
||||
return trimmed
|
||||
.parse::<whatsapp_rust::Jid>()
|
||||
.map_err(|e| anyhow!("Invalid WhatsApp JID `{trimmed}`: {e}"));
|
||||
}
|
||||
|
||||
let digits: String = trimmed.chars().filter(|c| c.is_ascii_digit()).collect();
|
||||
if digits.is_empty() {
|
||||
anyhow::bail!("Recipient `{trimmed}` does not contain a valid phone number");
|
||||
}
|
||||
|
||||
Ok(whatsapp_rust::Jid::pn(digits))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
#[async_trait]
|
||||
impl Channel for WhatsAppWebChannel {
|
||||
fn name(&self) -> &str {
|
||||
"whatsapp"
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> Result<()> {
|
||||
let client = self.client.lock().clone();
|
||||
let Some(client) = client else {
|
||||
anyhow::bail!("WhatsApp Web client not connected. Initialize the bot first.");
|
||||
};
|
||||
|
||||
if !self.should_allow_outbound(&message.recipient) {
|
||||
tracing::warn!(
|
||||
"WhatsApp Web: recipient {} not in allowed list",
|
||||
Self::redact_recipient(&message.recipient)
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let to = self.recipient_to_jid(&message.recipient)?;
|
||||
let outgoing = whatsapp_rust::waproto::whatsapp::Message {
|
||||
conversation: Some(message.content.clone()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let message_id = client.send_message(to, outgoing).await?;
|
||||
tracing::debug!(
|
||||
"WhatsApp Web: sent message to {} (id: {})",
|
||||
message.recipient,
|
||||
message_id
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn listen(&self, tx: tokio::sync::mpsc::Sender<ChannelMessage>) -> Result<()> {
|
||||
*self.tx.lock() = Some(tx.clone());
|
||||
|
||||
use wacore::types::events::Event;
|
||||
use whatsapp_rust::bot::Bot;
|
||||
use whatsapp_rust::pair_code::PairCodeOptions;
|
||||
use whatsapp_rust::TokioRuntime;
|
||||
use whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory;
|
||||
use whatsapp_rust_ureq_http_client::UreqHttpClient;
|
||||
|
||||
tracing::info!(
|
||||
"WhatsApp Web channel starting (session: {})",
|
||||
self.session_path
|
||||
);
|
||||
|
||||
// Upstream's Diesel-backed SqliteStore pulls a separate sqlite3 native
|
||||
// link chain that conflicts with the TinyAgents 1.3 / rusqlite 0.40
|
||||
// baseline. Keep the feature buildable with wacore's full in-memory
|
||||
// Backend until a rusqlite-backed durable store is added here.
|
||||
tracing::warn!(
|
||||
"[whatsapp_web] using in-memory WhatsApp Web session backend; \
|
||||
session_path={} is reserved but not persisted in this build",
|
||||
self.session_path
|
||||
);
|
||||
let backend: Arc<dyn wacore::store::traits::Backend> =
|
||||
Arc::new(wacore::store::InMemoryBackend::new());
|
||||
|
||||
let mut transport_factory = TokioWebSocketTransportFactory::new();
|
||||
if let Ok(ws_url) = std::env::var("WHATSAPP_WS_URL") {
|
||||
transport_factory = transport_factory.with_url(ws_url);
|
||||
}
|
||||
|
||||
let http_client = UreqHttpClient::new();
|
||||
|
||||
let tx_for_handler = tx.clone();
|
||||
let allowed_numbers = self.allowed_numbers.clone();
|
||||
let connected_for_handler = Arc::clone(&self.connected);
|
||||
let allowed_groups_for_handler = Arc::clone(&self.allowed_groups);
|
||||
|
||||
let mut builder = Bot::builder()
|
||||
.with_backend(backend)
|
||||
.with_transport_factory(transport_factory)
|
||||
.with_http_client(http_client)
|
||||
.with_runtime(TokioRuntime)
|
||||
.on_event(move |event, _client| {
|
||||
let tx_inner = tx_for_handler.clone();
|
||||
let allowed_numbers = allowed_numbers.clone();
|
||||
let connected = Arc::clone(&connected_for_handler);
|
||||
let allowed_groups = Arc::clone(&allowed_groups_for_handler);
|
||||
async move {
|
||||
match event {
|
||||
Event::Message(msg, info) => {
|
||||
// Self-echoes (messages this user sent from another
|
||||
// linked device) are mirrored to all devices via
|
||||
// the WhatsApp protocol. Drop them so the agent
|
||||
// doesn't react to its own outgoing messages.
|
||||
if info.source.is_from_me {
|
||||
return;
|
||||
}
|
||||
|
||||
let text = Self::extract_message_text(
|
||||
msg.conversation.as_deref(),
|
||||
msg.extended_text_message
|
||||
.as_ref()
|
||||
.and_then(|e| e.text.as_deref()),
|
||||
);
|
||||
|
||||
// Sender JID can use either the legacy `s.whatsapp.net`
|
||||
// server (phone-number addressing) or the newer `lid`
|
||||
// server (privacy-preserving identifier). Render the
|
||||
// user portion in E.164 with a leading `+` for the
|
||||
// allowed-list check + downstream subscriber.
|
||||
let sender_user = info.source.sender.user.clone();
|
||||
let normalized = if sender_user.starts_with('+') {
|
||||
sender_user.clone()
|
||||
} else {
|
||||
format!("+{sender_user}")
|
||||
};
|
||||
let chat = info.source.chat.to_string();
|
||||
let reply_target = Self::compute_reply_target(&chat, &normalized);
|
||||
|
||||
// Routine logs only carry coarse metadata — no raw
|
||||
// sender identifier, no message body — so PII does
|
||||
// not leak into application logs at any level.
|
||||
// For DM chats `chat` is `<phone>@s.whatsapp.net`,
|
||||
// which still carries the participant's phone
|
||||
// number. Redact the user part so the routine
|
||||
// log keeps only the server suffix (DM vs group)
|
||||
// and a coarse identifier tail.
|
||||
tracing::info!(
|
||||
"📨 WhatsApp inbound: chat={} sender={} text_len={}",
|
||||
Self::redact_recipient(&chat),
|
||||
Self::redact_phone(&normalized),
|
||||
text.len()
|
||||
);
|
||||
|
||||
if allowed_numbers.is_empty()
|
||||
|| allowed_numbers.iter().any(|n| n == "*" || n == &normalized)
|
||||
{
|
||||
// Record group provenance: this group has had at
|
||||
// least one allow-listed participant message in,
|
||||
// so subsequent outbound replies into the same
|
||||
// group are legitimate. Outbound to groups
|
||||
// without provenance is rejected by
|
||||
// `should_allow_outbound`.
|
||||
if Self::is_group_jid(&chat) {
|
||||
allowed_groups.lock().insert(chat.clone());
|
||||
}
|
||||
if let Err(e) = tx_inner
|
||||
.send(ChannelMessage {
|
||||
id: uuid::Uuid::new_v4().to_string(),
|
||||
channel: "whatsapp".to_string(),
|
||||
sender: normalized.clone(),
|
||||
reply_target,
|
||||
content: text,
|
||||
timestamp: chrono::Utc::now().timestamp_millis() as u64,
|
||||
thread_ts: None,
|
||||
})
|
||||
.await
|
||||
{
|
||||
tracing::error!(
|
||||
"Failed to forward WhatsApp message to channel: {}",
|
||||
e
|
||||
);
|
||||
}
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"WhatsApp Web: message from {} not in allowed list",
|
||||
Self::redact_phone(&normalized)
|
||||
);
|
||||
}
|
||||
}
|
||||
Event::Connected(_) => {
|
||||
connected.store(true, Ordering::Release);
|
||||
tracing::info!("✅ WhatsApp Web connected successfully!");
|
||||
}
|
||||
Event::LoggedOut(_) => {
|
||||
connected.store(false, Ordering::Release);
|
||||
tracing::warn!("❌ WhatsApp Web was logged out!");
|
||||
}
|
||||
Event::StreamError(stream_error) => {
|
||||
connected.store(false, Ordering::Release);
|
||||
tracing::error!("❌ WhatsApp Web stream error: {:?}", stream_error);
|
||||
}
|
||||
// The pair code and QR payload are short-lived link
|
||||
// credentials — anyone reading the logs while they
|
||||
// are valid can hijack the session. Surface only a
|
||||
// non-sensitive notice; the raw payload is never
|
||||
// logged at any level. Surfacing the code to the
|
||||
// user is the responsibility of an upstream UX
|
||||
// path (e.g. a JSON-RPC event the frontend renders).
|
||||
Event::PairingCode { .. } => {
|
||||
tracing::info!(
|
||||
"🔑 WhatsApp Web pair code received. Enter the code shown on \
|
||||
your linking surface into WhatsApp > Linked Devices."
|
||||
);
|
||||
}
|
||||
Event::PairingQrCode { .. } => {
|
||||
tracing::info!(
|
||||
"📱 WhatsApp Web QR code received. Render via QR generator and \
|
||||
scan with WhatsApp > Linked Devices."
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if let Some(ref phone) = self.pair_phone {
|
||||
tracing::info!("WhatsApp Web: pair-code flow enabled for configured phone number");
|
||||
builder = builder.with_pair_code(PairCodeOptions {
|
||||
phone_number: phone.clone(),
|
||||
custom_code: self.pair_code.clone(),
|
||||
..Default::default()
|
||||
});
|
||||
} else if self.pair_code.is_some() {
|
||||
tracing::warn!(
|
||||
"WhatsApp Web: pair_code is set but pair_phone is missing; pair code config is ignored"
|
||||
);
|
||||
}
|
||||
|
||||
let mut bot = builder.build().await?;
|
||||
*self.client.lock() = Some(bot.client());
|
||||
|
||||
let bot_handle = bot.run().await?;
|
||||
*self.bot_handle.lock() = Some(bot_handle);
|
||||
|
||||
// Wire into the shared shutdown machinery in `core::shutdown` so
|
||||
// SIGTERM and SIGINT both trigger a coordinated tear-down. The
|
||||
// previous `tokio::signal::ctrl_c()` path silently ignored
|
||||
// SIGTERM and bypassed the registered cleanup hooks the rest of
|
||||
// the process uses.
|
||||
let shutdown_notify = Arc::new(tokio::sync::Notify::new());
|
||||
let bot_handle_for_hook = Arc::clone(&self.bot_handle);
|
||||
let connected_for_hook = Arc::clone(&self.connected);
|
||||
let client_for_hook = Arc::clone(&self.client);
|
||||
let notify_for_hook = Arc::clone(&shutdown_notify);
|
||||
crate::core::shutdown::register(move || {
|
||||
let bot_handle = Arc::clone(&bot_handle_for_hook);
|
||||
let connected = Arc::clone(&connected_for_hook);
|
||||
let client = Arc::clone(&client_for_hook);
|
||||
let notify = Arc::clone(¬ify_for_hook);
|
||||
async move {
|
||||
tracing::info!("[whatsapp_web] graceful shutdown hook firing — aborting bot");
|
||||
connected.store(false, Ordering::Release);
|
||||
*client.lock() = None;
|
||||
if let Some(handle) = bot_handle.lock().take() {
|
||||
handle.abort();
|
||||
}
|
||||
notify.notify_waiters();
|
||||
}
|
||||
});
|
||||
|
||||
shutdown_notify.notified().await;
|
||||
tracing::info!("WhatsApp Web channel exited via shared shutdown");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
self.connected.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
async fn start_typing(&self, recipient: &str) -> Result<()> {
|
||||
let client = self.client.lock().clone();
|
||||
let Some(client) = client else {
|
||||
anyhow::bail!("WhatsApp Web client not connected. Initialize the bot first.");
|
||||
};
|
||||
|
||||
if !self.should_allow_outbound(recipient) {
|
||||
tracing::warn!(
|
||||
"WhatsApp Web: typing target {} not in allowed list",
|
||||
Self::redact_recipient(recipient)
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let to = self.recipient_to_jid(recipient)?;
|
||||
client
|
||||
.chatstate()
|
||||
.send_composing(&to)
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to send typing state (composing): {e}"))?;
|
||||
|
||||
tracing::debug!("WhatsApp Web: start typing for {}", recipient);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stop_typing(&self, recipient: &str) -> Result<()> {
|
||||
let client = self.client.lock().clone();
|
||||
let Some(client) = client else {
|
||||
anyhow::bail!("WhatsApp Web client not connected. Initialize the bot first.");
|
||||
};
|
||||
|
||||
if !self.should_allow_outbound(recipient) {
|
||||
tracing::warn!(
|
||||
"WhatsApp Web: typing target {} not in allowed list",
|
||||
Self::redact_recipient(recipient)
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let to = self.recipient_to_jid(recipient)?;
|
||||
client
|
||||
.chatstate()
|
||||
.send_paused(&to)
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to send typing state (paused): {e}"))?;
|
||||
|
||||
tracing::debug!("WhatsApp Web: stop typing for {}", recipient);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// Stub implementation when the feature is not enabled. Keeps the public ctor
|
||||
// signature compatible so `runtime/startup.rs` compiles unchanged.
|
||||
#[cfg(not(feature = "whatsapp-web"))]
|
||||
pub struct WhatsAppWebChannel {
|
||||
_private: (),
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "whatsapp-web"))]
|
||||
impl WhatsAppWebChannel {
|
||||
pub fn new(
|
||||
_session_path: String,
|
||||
_pair_phone: Option<String>,
|
||||
_pair_code: Option<String>,
|
||||
_allowed_numbers: Vec<String>,
|
||||
) -> Self {
|
||||
Self { _private: () }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "whatsapp-web"))]
|
||||
#[async_trait]
|
||||
impl Channel for WhatsAppWebChannel {
|
||||
fn name(&self) -> &str {
|
||||
"whatsapp"
|
||||
}
|
||||
|
||||
async fn send(&self, _message: &SendMessage) -> Result<()> {
|
||||
anyhow::bail!(
|
||||
"WhatsApp Web channel requires the 'whatsapp-web' feature. \
|
||||
Enable with: cargo build --features whatsapp-web"
|
||||
);
|
||||
}
|
||||
|
||||
async fn listen(&self, _tx: tokio::sync::mpsc::Sender<ChannelMessage>) -> Result<()> {
|
||||
anyhow::bail!(
|
||||
"WhatsApp Web channel requires the 'whatsapp-web' feature. \
|
||||
Enable with: cargo build --features whatsapp-web"
|
||||
);
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
async fn start_typing(&self, _recipient: &str) -> Result<()> {
|
||||
anyhow::bail!(
|
||||
"WhatsApp Web channel requires the 'whatsapp-web' feature. \
|
||||
Enable with: cargo build --features whatsapp-web"
|
||||
);
|
||||
}
|
||||
|
||||
async fn stop_typing(&self, _recipient: &str) -> Result<()> {
|
||||
anyhow::bail!(
|
||||
"WhatsApp Web channel requires the 'whatsapp-web' feature. \
|
||||
Enable with: cargo build --features whatsapp-web"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[path = "whatsapp_web_tests.rs"]
|
||||
mod tests;
|
||||
pub use tinychannels::providers::whatsapp_web::WhatsAppWebChannel;
|
||||
|
||||
@@ -1,254 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn make_channel() -> WhatsAppWebChannel {
|
||||
WhatsAppWebChannel::new(
|
||||
"/tmp/test-whatsapp.db".into(),
|
||||
None,
|
||||
None,
|
||||
vec!["+1234567890".into()],
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_channel_name() {
|
||||
let ch = make_channel();
|
||||
assert_eq!(ch.name(), "whatsapp");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_number_allowed_exact() {
|
||||
let ch = make_channel();
|
||||
assert!(ch.is_number_allowed("+1234567890"));
|
||||
assert!(!ch.is_number_allowed("+9876543210"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_number_allowed_wildcard() {
|
||||
let ch = WhatsAppWebChannel::new("/tmp/test.db".into(), None, None, vec!["*".into()]);
|
||||
assert!(ch.is_number_allowed("+1234567890"));
|
||||
assert!(ch.is_number_allowed("+9999999999"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_number_denied_empty() {
|
||||
let ch = WhatsAppWebChannel::new("/tmp/test.db".into(), None, None, vec![]);
|
||||
// Empty allowed_numbers means "allow all" (same behavior as Cloud API)
|
||||
assert!(ch.is_number_allowed("+1234567890"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_normalize_phone_adds_plus() {
|
||||
let ch = make_channel();
|
||||
assert_eq!(ch.normalize_phone("1234567890"), "+1234567890");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_normalize_phone_preserves_plus() {
|
||||
let ch = make_channel();
|
||||
assert_eq!(ch.normalize_phone("+1234567890"), "+1234567890");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
async fn whatsapp_web_health_check_disconnected() {
|
||||
let ch = make_channel();
|
||||
assert!(!ch.health_check().await);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
async fn whatsapp_web_health_check_tracks_connected_flag() {
|
||||
let ch = make_channel();
|
||||
assert!(!ch.health_check().await);
|
||||
ch.connected.store(true, Ordering::Release);
|
||||
assert!(ch.health_check().await);
|
||||
ch.connected.store(false, Ordering::Release);
|
||||
assert!(!ch.health_check().await);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_compute_reply_target_dm_pn() {
|
||||
assert_eq!(
|
||||
WhatsAppWebChannel::compute_reply_target("123@s.whatsapp.net", "+1234567890"),
|
||||
"+1234567890"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_compute_reply_target_dm_lid() {
|
||||
assert_eq!(
|
||||
WhatsAppWebChannel::compute_reply_target("abc@lid", "+1234567890"),
|
||||
"+1234567890"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_compute_reply_target_group() {
|
||||
assert_eq!(
|
||||
WhatsAppWebChannel::compute_reply_target("987654@g.us", "+1234567890"),
|
||||
"987654@g.us"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_redact_phone_e164() {
|
||||
assert_eq!(WhatsAppWebChannel::redact_phone("+1234567890"), "+***7890");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_redact_phone_no_plus() {
|
||||
assert_eq!(WhatsAppWebChannel::redact_phone("1234567890"), "***7890");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_redact_phone_short_input() {
|
||||
// Pathological short inputs collapse to a generic mask rather than
|
||||
// exposing the entire identifier.
|
||||
assert_eq!(WhatsAppWebChannel::redact_phone("+12"), "+****");
|
||||
assert_eq!(WhatsAppWebChannel::redact_phone("12"), "****");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_extract_message_text_prefers_conversation() {
|
||||
assert_eq!(
|
||||
WhatsAppWebChannel::extract_message_text(Some("hello"), Some("ignored")),
|
||||
"hello"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_extract_message_text_falls_back_to_extended() {
|
||||
assert_eq!(
|
||||
WhatsAppWebChannel::extract_message_text(None, Some("from extended")),
|
||||
"from extended"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_extract_message_text_empty_when_missing() {
|
||||
assert_eq!(WhatsAppWebChannel::extract_message_text(None, None), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_is_group_jid_recognises_group() {
|
||||
assert!(WhatsAppWebChannel::is_group_jid("123456@g.us"));
|
||||
assert!(WhatsAppWebChannel::is_group_jid(" 4567@g.us "));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_is_group_jid_rejects_non_group() {
|
||||
assert!(!WhatsAppWebChannel::is_group_jid("+1234567890"));
|
||||
assert!(!WhatsAppWebChannel::is_group_jid("123@s.whatsapp.net"));
|
||||
assert!(!WhatsAppWebChannel::is_group_jid("abc@lid"));
|
||||
assert!(!WhatsAppWebChannel::is_group_jid(""));
|
||||
}
|
||||
|
||||
/// Regression for CodeRabbit finding: an `@g.us` reply target was being
|
||||
/// silently dropped because the outbound path normalised the JID to
|
||||
/// `+<group-id>` and missed the per-number allowlist. After provenance
|
||||
/// is recorded, an allowed user replying back into the group they came
|
||||
/// from must succeed.
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_should_allow_outbound_provenanced_group_allowed() {
|
||||
let ch = make_channel(); // allowed_numbers = ["+1234567890"]
|
||||
ch.allowed_groups
|
||||
.lock()
|
||||
.insert("987654321@g.us".to_string());
|
||||
assert!(ch.should_allow_outbound("987654321@g.us"));
|
||||
}
|
||||
|
||||
/// Regression for the follow-up CodeRabbit finding: a blanket `@g.us`
|
||||
/// bypass is itself a vulnerability — a caller able to set `recipient`
|
||||
/// could post into arbitrary joined groups. Groups without recorded
|
||||
/// provenance must stay blocked.
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_should_allow_outbound_unrelated_group_blocked() {
|
||||
let ch = make_channel();
|
||||
ch.allowed_groups
|
||||
.lock()
|
||||
.insert("987654321@g.us".to_string());
|
||||
assert!(!ch.should_allow_outbound("11111@g.us"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_should_allow_outbound_group_without_provenance_blocked() {
|
||||
let ch = make_channel();
|
||||
// empty allowed_groups
|
||||
assert!(!ch.should_allow_outbound("987654321@g.us"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_redact_recipient_pn_jid() {
|
||||
assert_eq!(
|
||||
WhatsAppWebChannel::redact_recipient("1234567890@s.whatsapp.net"),
|
||||
"***7890@s.whatsapp.net"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_redact_recipient_group_jid() {
|
||||
assert_eq!(
|
||||
WhatsAppWebChannel::redact_recipient("987654321@g.us"),
|
||||
"***4321@g.us"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_redact_recipient_bare_phone() {
|
||||
assert_eq!(
|
||||
WhatsAppWebChannel::redact_recipient("+1234567890"),
|
||||
"+***7890"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_should_allow_outbound_dm_blocks_unallowed() {
|
||||
let ch = make_channel();
|
||||
assert!(!ch.should_allow_outbound("+9999999999"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_should_allow_outbound_dm_allows_match() {
|
||||
let ch = make_channel();
|
||||
assert!(ch.should_allow_outbound("+1234567890"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_should_allow_outbound_wildcard_passes_dm() {
|
||||
let ch = WhatsAppWebChannel::new("/tmp/t.db".into(), None, None, vec!["*".into()]);
|
||||
assert!(ch.should_allow_outbound("+9999999999"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
fn whatsapp_web_should_allow_outbound_empty_allowlist_passes_dm() {
|
||||
let ch = WhatsAppWebChannel::new("/tmp/t.db".into(), None, None, vec![]);
|
||||
assert!(ch.should_allow_outbound("+9999999999"));
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub use tinychannels::providers::yuanbao::*;
|
||||
@@ -1,710 +0,0 @@
|
||||
//! Channel facade for the Yuanbao provider.
|
||||
//!
|
||||
//! This module owns the OpenHuman [`Channel`] implementation and keeps
|
||||
//! provider wiring out of `mod.rs`. Protocol decoding, transport, inbound
|
||||
//! filtering, and outbound sending remain delegated to sibling modules.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::{mpsc, watch, Mutex as TokioMutex};
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::openhuman::channels::traits::{Channel, ChannelMessage, SendMessage};
|
||||
|
||||
use super::config::YuanbaoConfig;
|
||||
use super::connection::{InboundEvent, YuanbaoConnection};
|
||||
use super::ids::{shorten_account_id, shorten_reply_target};
|
||||
use super::inbound::{InboundPipeline, PipelineOutcome, PipelineState};
|
||||
use super::outbound::OutboundSender;
|
||||
use super::proto::decode_push_msg;
|
||||
use super::sign::SignManager;
|
||||
use super::{splitter, types};
|
||||
|
||||
/// Reply Heartbeat keepalive interval. The yuanbao gateway expects the
|
||||
/// bot to ping (`SendPrivateHeartbeat RUNNING`) at this cadence so the
|
||||
/// "正在输入" indicator stays alive for long-running responses.
|
||||
const REPLY_HEARTBEAT_INTERVAL_SECS: u64 = 2;
|
||||
|
||||
/// Hard ceiling on the in-memory shortened-recipient → original-recipient
|
||||
/// map. Each entry is two short strings (~80 B), so 4096 distinct senders
|
||||
/// give ~320 KB — plenty for any realistic chat load and small enough
|
||||
/// that we can blow the whole map away when we hit the cap instead of
|
||||
/// dragging in an LRU dependency. See `register_recipient_alias`.
|
||||
const RECIPIENT_ALIAS_CAP: usize = 4096;
|
||||
|
||||
/// The yuanbao channel — owns one WebSocket and one inbound pipeline.
|
||||
pub struct YuanbaoChannel {
|
||||
config: YuanbaoConfig,
|
||||
connection: Arc<YuanbaoConnection>,
|
||||
outbound: Arc<OutboundSender>,
|
||||
pipeline: Arc<InboundPipeline>,
|
||||
shutdown_tx: watch::Sender<bool>,
|
||||
/// Holds the inbound receiver between `new()` and the first `listen()` call.
|
||||
///
|
||||
/// `Channel::listen` takes `&self`, so we can't move the receiver out of
|
||||
/// a field. Use a `Mutex<Option<…>>` so the first listener takes ownership
|
||||
/// and subsequent calls fail cleanly.
|
||||
inbound_rx: parking_lot::Mutex<Option<mpsc::UnboundedReceiver<InboundEvent>>>,
|
||||
/// Per-recipient Reply Heartbeat keepalive tasks (started on `start_typing`).
|
||||
heartbeat_tasks: TokioMutex<HashMap<String, JoinHandle<()>>>,
|
||||
/// Reverse lookup table from shortened recipient ids (the ones we
|
||||
/// emit on `ChannelMessage.sender` / `reply_target`) back to the
|
||||
/// original server-recognized ids that outbound `send_c2c_message`
|
||||
/// / `send_group_message` must use as `to_account` / `group_code`.
|
||||
///
|
||||
/// Why this exists: yuanbao uids are ~64-char hashes, and
|
||||
/// `super::ids::shorten_account_id` rewrites them as
|
||||
/// `<prefix>_<sha256-16hex>` so the conversation store's per-thread
|
||||
/// JSONL filenames stay under filesystem `NAME_MAX`. Without this
|
||||
/// table the agent loop sends replies addressed to the shortened
|
||||
/// hash, which the yuanbao gateway silently drops because no such
|
||||
/// user exists. See `register_recipient_alias` / `resolve_recipient`.
|
||||
recipient_aliases: TokioMutex<HashMap<String, String>>,
|
||||
}
|
||||
|
||||
impl YuanbaoChannel {
|
||||
/// Build a channel from a validated config. Returns an error if the
|
||||
/// config is missing required fields (so misconfiguration surfaces
|
||||
/// at startup, not on the first inbound message).
|
||||
pub fn new(mut config: YuanbaoConfig) -> anyhow::Result<Self> {
|
||||
config.apply_env_defaults();
|
||||
config.validate().map_err(anyhow::Error::msg)?;
|
||||
let (shutdown_tx, _shutdown_rx) = watch::channel(false);
|
||||
let (inbound_tx, inbound_rx) = mpsc::unbounded_channel::<InboundEvent>();
|
||||
|
||||
// SignManager is only useful when we have an app_secret — without
|
||||
// it we'd never call the sign endpoint anyway.
|
||||
let sign_manager: Option<Arc<SignManager>> = if !config.app_secret.is_empty() {
|
||||
Some(SignManager::new(reqwest::Client::new()))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let connection = YuanbaoConnection::new(config.clone(), inbound_tx, sign_manager.clone());
|
||||
let outbound = Arc::new(OutboundSender::new(
|
||||
Arc::clone(&connection),
|
||||
sign_manager.clone(),
|
||||
config.app_key.clone(),
|
||||
config.bot_id.clone(),
|
||||
));
|
||||
// PipelineState's `from_account` is used by the echo-guard stage to
|
||||
// drop self-sent messages. We feed it the static config value here
|
||||
// (which may be empty); the canonical server-issued bot_id only
|
||||
// becomes known after sign-token, so this is a known minor gap —
|
||||
// echo guard will simply not fire when bot_id isn't statically set.
|
||||
let pipeline_state = PipelineState::new(&config, config.bot_id.clone());
|
||||
let pipeline = Arc::new(InboundPipeline::new(pipeline_state));
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
connection,
|
||||
outbound,
|
||||
pipeline,
|
||||
shutdown_tx,
|
||||
inbound_rx: parking_lot::Mutex::new(Some(inbound_rx)),
|
||||
heartbeat_tasks: TokioMutex::new(HashMap::new()),
|
||||
recipient_aliases: TokioMutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Record a `shortened → original` recipient mapping so the outbound
|
||||
/// side can recover the server-recognized id when the agent loop
|
||||
/// addresses a reply with the shortened sender / reply_target it
|
||||
/// received on `ChannelMessage`.
|
||||
///
|
||||
/// No-op when the two are equal (uid is short enough to skip
|
||||
/// shortening, or this is the `g:` group-target case where the
|
||||
/// inner code is short). When the map crosses `RECIPIENT_ALIAS_CAP`
|
||||
/// we clear it — the next inbound message from each active sender
|
||||
/// re-populates the entry it needs, and stale entries from idle
|
||||
/// conversations are fine to lose.
|
||||
async fn register_recipient_alias(&self, shortened: &str, original: &str) {
|
||||
if shortened == original {
|
||||
return;
|
||||
}
|
||||
let mut m = self.recipient_aliases.lock().await;
|
||||
if m.len() >= RECIPIENT_ALIAS_CAP {
|
||||
warn!(
|
||||
"[yuanbao] recipient alias map hit cap ({}), clearing",
|
||||
RECIPIENT_ALIAS_CAP
|
||||
);
|
||||
m.clear();
|
||||
}
|
||||
m.insert(shortened.to_string(), original.to_string());
|
||||
}
|
||||
|
||||
/// Look up the server-recognized recipient for a (possibly
|
||||
/// shortened) inbound id. Falls back to the input unchanged when
|
||||
/// nothing is registered — which keeps the previous behavior for
|
||||
/// recipients that don't go through `shorten_account_id` (short
|
||||
/// uids, group codes, `imessage`-style ids).
|
||||
async fn resolve_recipient(&self, recipient: &str) -> String {
|
||||
let m = self.recipient_aliases.lock().await;
|
||||
m.get(recipient)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| recipient.to_string())
|
||||
}
|
||||
|
||||
fn split_message(&self, text: &str) -> Vec<String> {
|
||||
splitter::split_markdown(text, self.config.max_message_length)
|
||||
}
|
||||
|
||||
async fn start_heartbeat_task(&self, recipient: &str) {
|
||||
let mut tasks = self.heartbeat_tasks.lock().await;
|
||||
if tasks.contains_key(recipient) {
|
||||
return;
|
||||
}
|
||||
let outbound = Arc::clone(&self.outbound);
|
||||
let target = recipient.to_string();
|
||||
let handle = tokio::spawn(async move {
|
||||
let mut interval =
|
||||
tokio::time::interval(Duration::from_secs(REPLY_HEARTBEAT_INTERVAL_SECS));
|
||||
interval.tick().await; // skip first tick (start_typing already sent RUNNING)
|
||||
loop {
|
||||
interval.tick().await;
|
||||
if let Err(e) = outbound.start_heartbeat(&target).await {
|
||||
// Connection bouncing — bail out of this loop; the
|
||||
// next start_typing call will spawn a new one.
|
||||
warn!(
|
||||
"[yuanbao] reply heartbeat send failed: {} — stopping loop",
|
||||
e
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
tasks.insert(recipient.to_string(), handle);
|
||||
}
|
||||
|
||||
async fn stop_heartbeat_task(&self, recipient: &str) {
|
||||
let mut tasks = self.heartbeat_tasks.lock().await;
|
||||
if let Some(handle) = tasks.remove(recipient) {
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for YuanbaoChannel {
|
||||
fn name(&self) -> &str {
|
||||
"yuanbao"
|
||||
}
|
||||
|
||||
async fn send(&self, message: &SendMessage) -> anyhow::Result<()> {
|
||||
let chunks = self.split_message(&message.content);
|
||||
let ref_msg_id = message.thread_ts.as_deref();
|
||||
let recipient = self.resolve_recipient(&message.recipient).await;
|
||||
for chunk in &chunks {
|
||||
self.outbound
|
||||
.send_text(&recipient, chunk, ref_msg_id)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn supports_draft_updates(&self) -> bool {
|
||||
// Routes turns through the streaming code path even though Yuanbao
|
||||
// itself has no edit-message capability. We accept the UX cost (no
|
||||
// progressive rendering — the reply appears all at once in
|
||||
// `finalize_draft`) in exchange for streaming's tolerance of
|
||||
// malformed `usage` chunks; the non-streaming parser fails the
|
||||
// whole turn when an upstream LLM returns string-typed token counts.
|
||||
true
|
||||
}
|
||||
|
||||
async fn send_draft(&self, message: &SendMessage) -> anyhow::Result<Option<String>> {
|
||||
// Marker id so dispatch spins up the progress consumer task;
|
||||
// nothing is sent to the user here. Real content goes out in
|
||||
// `finalize_draft`. See `supports_draft_updates` for rationale.
|
||||
Ok(Some(format!("yb-draft:{}", message.recipient)))
|
||||
}
|
||||
|
||||
async fn update_draft(
|
||||
&self,
|
||||
_recipient: &str,
|
||||
_message_id: &str,
|
||||
_text: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn finalize_draft(
|
||||
&self,
|
||||
recipient: &str,
|
||||
_message_id: &str,
|
||||
text: &str,
|
||||
thread_ts: Option<&str>,
|
||||
) -> anyhow::Result<()> {
|
||||
let chunks = self.split_message(text);
|
||||
let recipient = self.resolve_recipient(recipient).await;
|
||||
for chunk in &chunks {
|
||||
self.outbound
|
||||
.send_text(&recipient, chunk, thread_ts)
|
||||
.await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn listen(&self, tx: mpsc::Sender<ChannelMessage>) -> anyhow::Result<()> {
|
||||
// Take the inbound receiver. A second listener would just exit early.
|
||||
let mut inbound_rx = match self.inbound_rx.lock().take() {
|
||||
Some(rx) => rx,
|
||||
None => {
|
||||
warn!("[yuanbao] listen() called twice — second call exits");
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
let conn = Arc::clone(&self.connection);
|
||||
let shutdown_rx = self.shutdown_tx.subscribe();
|
||||
let mut conn_task = tokio::spawn(async move {
|
||||
conn.run(shutdown_rx).await;
|
||||
});
|
||||
|
||||
info!("[yuanbao] channel listening — pipeline ready");
|
||||
let mut shutdown_rx2 = self.shutdown_tx.subscribe();
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx2.changed() => {
|
||||
info!("[yuanbao] listen loop received shutdown");
|
||||
break;
|
||||
}
|
||||
event = inbound_rx.recv() => {
|
||||
match event {
|
||||
Some(InboundEvent::Push(frame)) => {
|
||||
self.dispatch_push(frame, &tx).await;
|
||||
}
|
||||
Some(InboundEvent::Kickout(reason)) => {
|
||||
warn!("[yuanbao] kickout: {} — stopping listen loop", reason);
|
||||
break;
|
||||
}
|
||||
None => {
|
||||
warn!("[yuanbao] inbound channel closed");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let _ = self.shutdown_tx.send(true);
|
||||
// Give the connection task a brief window to run its own shutdown cleanup
|
||||
// (flush pending, update is_connected, etc.) before force-aborting.
|
||||
match tokio::time::timeout(std::time::Duration::from_secs(2), &mut conn_task).await {
|
||||
Ok(_) => {}
|
||||
Err(_) => conn_task.abort(),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> bool {
|
||||
self.connection.is_connected()
|
||||
}
|
||||
|
||||
async fn start_typing(&self, recipient: &str) -> anyhow::Result<()> {
|
||||
// Send RUNNING immediately, then spawn a 2s keepalive so the
|
||||
// indicator doesn't expire while we generate.
|
||||
let recipient = self.resolve_recipient(recipient).await;
|
||||
self.outbound.start_heartbeat(&recipient).await?;
|
||||
self.start_heartbeat_task(&recipient).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn stop_typing(&self, recipient: &str) -> anyhow::Result<()> {
|
||||
let recipient = self.resolve_recipient(recipient).await;
|
||||
self.stop_heartbeat_task(&recipient).await;
|
||||
self.outbound.stop_heartbeat(&recipient).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn supports_reactions(&self) -> bool {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
impl YuanbaoChannel {
|
||||
async fn dispatch_push(&self, frame: types::ConnFrame, tx: &mpsc::Sender<ChannelMessage>) {
|
||||
// The Yuanbao gateway pushes inbound messages with `cmd_type=Push`;
|
||||
// the actual `cmd` word is decided server-side and varies (mirrors
|
||||
// hermes-agent `yuanbao.py::_handle_received_frame` which routes
|
||||
// purely on cmd_type). The connection layer has already filtered
|
||||
// out non-push frames before we get here, so every frame we see
|
||||
// should be a candidate for the inbound pipeline.
|
||||
if frame.data.is_empty() {
|
||||
tracing::trace!("[yuanbao] empty push body cmd={} — skipping", frame.cmd);
|
||||
return;
|
||||
}
|
||||
// Some push frames wrap the biz body in an extra
|
||||
// `PushMsg { cmd, module, msg_id, data }` envelope; others (e.g.
|
||||
// cmd="inbound_message", module="yuanbao_openclaw_proxy") put the
|
||||
// InboundMessagePush bytes directly in `ConnMsg.data` with the
|
||||
// ConnMsg.head already carrying cmd/module. Mirrors plugin
|
||||
// client.ts::onPush (l. 813): try PushMsg first, but only accept
|
||||
// it when it has a non-empty cmd or module; otherwise treat the
|
||||
// raw frame.data as the biz body.
|
||||
let unwrapped: Option<Vec<u8>> = match decode_push_msg(&frame.data) {
|
||||
Ok(p) if (!p.cmd.is_empty() || !p.module.is_empty()) && !p.data.is_empty() => {
|
||||
info!(
|
||||
"[yuanbao] push envelope decoded: cmd={} module={} msg_id={} biz_len={}",
|
||||
p.cmd,
|
||||
p.module,
|
||||
p.msg_id,
|
||||
p.data.len()
|
||||
);
|
||||
Some(p.data)
|
||||
}
|
||||
_ => {
|
||||
info!(
|
||||
"[yuanbao] push has no PushMsg envelope — treating ConnMsg.data as biz body (conn_cmd={} module={} len={})",
|
||||
frame.cmd,
|
||||
frame.module,
|
||||
frame.data.len()
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
let biz_body: &[u8] = unwrapped.as_deref().unwrap_or(&frame.data);
|
||||
let outcome = self.pipeline.process(biz_body).await;
|
||||
match outcome {
|
||||
PipelineOutcome::Dispatch(ctx) => {
|
||||
// Shorten ids at the channel boundary so the composite thread_id
|
||||
// derived downstream (channel:yuanbao_<sender>_<reply_target>)
|
||||
// stays under filesystem NAME_MAX once hex-encoded for the
|
||||
// per-thread JSONL filename. Yuanbao internals (echo guard,
|
||||
// access control, owner-command check) keep the original
|
||||
// `from_account` — see `super::ids` for the format and rationale.
|
||||
let original_from = ctx.msg.from_account.clone();
|
||||
let original_reply_target = ctx.source.reply_target();
|
||||
let short_sender = shorten_account_id(&original_from);
|
||||
let short_reply_target = shorten_reply_target(&original_reply_target);
|
||||
// Remember the original ids so the outbound side can
|
||||
// recover them when the agent loop addresses a reply
|
||||
// with the shortened values it sees here.
|
||||
self.register_recipient_alias(&short_sender, &original_from)
|
||||
.await;
|
||||
self.register_recipient_alias(&short_reply_target, &original_reply_target)
|
||||
.await;
|
||||
let msg = ChannelMessage {
|
||||
id: ctx.msg.msg_id.clone(),
|
||||
sender: short_sender,
|
||||
reply_target: short_reply_target,
|
||||
content: if ctx.text.is_empty() && !ctx.image_urls.is_empty() {
|
||||
// Surface image URLs as content so downstream tools have something to work with.
|
||||
ctx.image_urls.join("\n")
|
||||
} else {
|
||||
ctx.text.clone()
|
||||
},
|
||||
channel: "yuanbao".into(),
|
||||
timestamp: ctx.msg.msg_time as u64,
|
||||
thread_ts: None,
|
||||
};
|
||||
if tx.send(msg).await.is_err() {
|
||||
warn!("[yuanbao] dispatch receiver gone — dropping message");
|
||||
}
|
||||
}
|
||||
PipelineOutcome::Filtered(reason) => {
|
||||
tracing::trace!("[yuanbao] filtered at {reason}");
|
||||
}
|
||||
PipelineOutcome::Failed(err) => {
|
||||
// Intentionally omit the raw biz payload — it can carry
|
||||
// user content / PII. The decoder error already encodes
|
||||
// the structural reason; only the length is safe to log.
|
||||
warn!(
|
||||
"[yuanbao] pipeline error: {err} | biz_len={}",
|
||||
biz_body.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for YuanbaoChannel {
|
||||
fn drop(&mut self) {
|
||||
let _ = self.shutdown_tx.send(true);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::openhuman::channels::traits::Channel;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn good_cfg() -> YuanbaoConfig {
|
||||
let mut c = YuanbaoConfig::default();
|
||||
c.app_key = "ak".into();
|
||||
c.ws_domain = "wss://example".into();
|
||||
c.token = "tok".into();
|
||||
c
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_construction_validates() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
assert_eq!(ch.name(), "yuanbao");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_config_rejected() {
|
||||
let mut c = YuanbaoConfig::default();
|
||||
c.app_key = "ak".into();
|
||||
// missing ws_domain
|
||||
assert!(YuanbaoChannel::new(c).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_short_message_returns_one() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
let chunks = ch.split_message("hello");
|
||||
assert_eq!(chunks, vec!["hello"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_respects_newlines() {
|
||||
let mut c = good_cfg();
|
||||
c.max_message_length = 12;
|
||||
let ch = YuanbaoChannel::new(c).unwrap();
|
||||
let chunks = ch.split_message("line one\nline two\nline three");
|
||||
assert!(chunks.len() >= 2);
|
||||
// No chunk exceeds the limit.
|
||||
for chunk in &chunks {
|
||||
assert!(chunk.len() <= 12, "chunk too long: {chunk:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_recipient_returns_input_when_no_alias_registered() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
assert_eq!(ch.resolve_recipient("short_uid").await, "short_uid");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn register_and_resolve_dm_alias_recovers_original_uid() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
let original = "x".repeat(64);
|
||||
let shortened = shorten_account_id(&original);
|
||||
assert_ne!(shortened, original, "test premise: should actually shorten");
|
||||
ch.register_recipient_alias(&shortened, &original).await;
|
||||
assert_eq!(ch.resolve_recipient(&shortened).await, original);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn register_recipient_alias_is_noop_for_equal_pair() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
// Short uid that wouldn't be shortened — caller still hands us
|
||||
// (s, s); we should silently skip and not eat a map slot.
|
||||
ch.register_recipient_alias("short", "short").await;
|
||||
let m = ch.recipient_aliases.lock().await;
|
||||
assert!(m.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_recipient_preserves_group_prefix_via_alias() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
let long_group_code = "g".repeat(64);
|
||||
let original = format!("g:{long_group_code}");
|
||||
let shortened = shorten_reply_target(&original);
|
||||
assert_ne!(shortened, original);
|
||||
assert!(shortened.starts_with("g:"));
|
||||
ch.register_recipient_alias(&shortened, &original).await;
|
||||
assert_eq!(ch.resolve_recipient(&shortened).await, original);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn alias_map_clears_when_cap_is_hit() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
// Pre-fill up to the cap with distinct entries.
|
||||
for i in 0..RECIPIENT_ALIAS_CAP {
|
||||
ch.register_recipient_alias(&format!("s{i}"), &format!("o{i}"))
|
||||
.await;
|
||||
}
|
||||
assert_eq!(ch.recipient_aliases.lock().await.len(), RECIPIENT_ALIAS_CAP);
|
||||
// One more entry must trigger a clear, then insert the new entry.
|
||||
ch.register_recipient_alias("new_short", "new_original")
|
||||
.await;
|
||||
let m = ch.recipient_aliases.lock().await;
|
||||
assert_eq!(m.len(), 1);
|
||||
assert_eq!(m.get("new_short").map(String::as_str), Some("new_original"));
|
||||
}
|
||||
|
||||
// ─── trivial trait methods ─────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn supports_draft_updates_is_true() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
assert!(ch.supports_draft_updates());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supports_reactions_is_false() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
assert!(!ch.supports_reactions());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_draft_returns_marker_id() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
let msg = SendMessage::new("ignored", "user-42");
|
||||
let id = ch.send_draft(&msg).await.unwrap();
|
||||
assert_eq!(id.as_deref(), Some("yb-draft:user-42"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_draft_is_a_noop_ok() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
assert!(ch.update_draft("user-42", "any-id", "text").await.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn health_check_is_false_when_socket_not_connected() {
|
||||
// Real connect requires a WebSocket; we only verify the
|
||||
// disconnected default here. The connected branch is exercised
|
||||
// by `connection::tests::set_state_connected_flips_is_connected_flag`.
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
assert!(!ch.health_check().await);
|
||||
}
|
||||
|
||||
// ─── dispatch_push branches ────────────────────────────────────
|
||||
|
||||
fn make_push_frame(cmd: &str, data: Vec<u8>) -> types::ConnFrame {
|
||||
types::ConnFrame {
|
||||
cmd_type: super::super::proto_constants::cmd_type::PUSH,
|
||||
cmd: cmd.into(),
|
||||
module: "yuanbao_openclaw_proxy".into(),
|
||||
seq_no: 0,
|
||||
msg_id: String::new(),
|
||||
need_ack: false,
|
||||
status: 0,
|
||||
data,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatch_push_empty_body_is_skipped() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
let (tx, mut rx) = mpsc::channel::<ChannelMessage>(4);
|
||||
let frame = make_push_frame("noop", Vec::new());
|
||||
ch.dispatch_push(frame, &tx).await;
|
||||
// No message should reach the sender.
|
||||
assert!(rx.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatch_push_garbage_body_does_not_dispatch() {
|
||||
// Body is not a valid protobuf push *and* not valid JSON → Failed.
|
||||
// dispatch_push should log + swallow, not propagate panic.
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
let (tx, mut rx) = mpsc::channel::<ChannelMessage>(4);
|
||||
let frame = make_push_frame("inbound_message", vec![0xFF, 0xFF, 0xFF, 0xFF]);
|
||||
ch.dispatch_push(frame, &tx).await;
|
||||
assert!(rx.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatch_push_dm_text_reaches_listener() {
|
||||
// Build a minimal `InboundMessagePush` directly in ConnFrame.data
|
||||
// (no PushMsg envelope), with a single TIMTextElem so the pipeline
|
||||
// dispatches.
|
||||
use super::super::proto::{encode_msg_body_element, encode_varint};
|
||||
let elem = types::MsgBodyElement {
|
||||
msg_type: "TIMTextElem".into(),
|
||||
msg_content: types::MsgContent {
|
||||
text: Some("hello".into()),
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
let elem_bytes = encode_msg_body_element(&elem);
|
||||
|
||||
// Hand-roll an InboundMessagePush so we don't depend on a helper:
|
||||
// field 2 = from_account, field 3 = to_account, field 12 = msg_id,
|
||||
// field 13 = repeated MsgBodyElement.
|
||||
let mut biz = Vec::new();
|
||||
let put_string = |fnum: u32, s: &str, b: &mut Vec<u8>| {
|
||||
encode_varint(((fnum as u64) << 3) | 2, b);
|
||||
encode_varint(s.len() as u64, b);
|
||||
b.extend_from_slice(s.as_bytes());
|
||||
};
|
||||
put_string(2, "alice", &mut biz);
|
||||
put_string(3, "bot1", &mut biz);
|
||||
put_string(12, "mid-x", &mut biz);
|
||||
encode_varint(((13u64) << 3) | 2, &mut biz);
|
||||
encode_varint(elem_bytes.len() as u64, &mut biz);
|
||||
biz.extend_from_slice(&elem_bytes);
|
||||
|
||||
// Disable group_at_required and use open dm_access so the
|
||||
// pipeline passes all stages for this DM.
|
||||
let mut cfg = good_cfg();
|
||||
cfg.dm_access = "open".into();
|
||||
cfg.bot_id = "bot1".into();
|
||||
let ch = YuanbaoChannel::new(cfg).unwrap();
|
||||
|
||||
let frame = make_push_frame("inbound_message", biz);
|
||||
let (tx, mut rx) = mpsc::channel::<ChannelMessage>(4);
|
||||
ch.dispatch_push(frame, &tx).await;
|
||||
let msg = rx.try_recv().expect("dispatch should produce one message");
|
||||
assert_eq!(msg.id, "mid-x");
|
||||
assert_eq!(msg.content, "hello");
|
||||
assert_eq!(msg.channel, "yuanbao");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dispatch_push_filtered_by_dedup_does_not_double_dispatch() {
|
||||
use super::super::proto::{encode_msg_body_element, encode_varint};
|
||||
let elem = types::MsgBodyElement {
|
||||
msg_type: "TIMTextElem".into(),
|
||||
msg_content: types::MsgContent {
|
||||
text: Some("dup".into()),
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
let elem_bytes = encode_msg_body_element(&elem);
|
||||
let mut biz = Vec::new();
|
||||
let put_string = |fnum: u32, s: &str, b: &mut Vec<u8>| {
|
||||
encode_varint(((fnum as u64) << 3) | 2, b);
|
||||
encode_varint(s.len() as u64, b);
|
||||
b.extend_from_slice(s.as_bytes());
|
||||
};
|
||||
put_string(2, "alice", &mut biz);
|
||||
put_string(3, "bot1", &mut biz);
|
||||
put_string(12, "dup-id", &mut biz);
|
||||
encode_varint(((13u64) << 3) | 2, &mut biz);
|
||||
encode_varint(elem_bytes.len() as u64, &mut biz);
|
||||
biz.extend_from_slice(&elem_bytes);
|
||||
|
||||
let mut cfg = good_cfg();
|
||||
cfg.dm_access = "open".into();
|
||||
cfg.bot_id = "bot1".into();
|
||||
let ch = YuanbaoChannel::new(cfg).unwrap();
|
||||
let (tx, mut rx) = mpsc::channel::<ChannelMessage>(4);
|
||||
ch.dispatch_push(make_push_frame("inbound_message", biz.clone()), &tx)
|
||||
.await;
|
||||
assert!(rx.try_recv().is_ok(), "first should dispatch");
|
||||
ch.dispatch_push(make_push_frame("inbound_message", biz), &tx)
|
||||
.await;
|
||||
assert!(rx.try_recv().is_err(), "second (same id) should dedup");
|
||||
}
|
||||
|
||||
// ─── heartbeat task lifecycle ──────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn start_heartbeat_task_inserts_and_stop_removes() {
|
||||
let ch = YuanbaoChannel::new(good_cfg()).unwrap();
|
||||
ch.start_heartbeat_task("recipient-1").await;
|
||||
assert!(
|
||||
ch.heartbeat_tasks.lock().await.contains_key("recipient-1"),
|
||||
"should have spawned a task for recipient-1"
|
||||
);
|
||||
// Second start for same recipient is a no-op (does not double-spawn).
|
||||
ch.start_heartbeat_task("recipient-1").await;
|
||||
assert_eq!(ch.heartbeat_tasks.lock().await.len(), 1);
|
||||
|
||||
ch.stop_heartbeat_task("recipient-1").await;
|
||||
assert!(ch.heartbeat_tasks.lock().await.is_empty());
|
||||
// Stopping a recipient with no task is also a no-op.
|
||||
ch.stop_heartbeat_task("never-started").await;
|
||||
}
|
||||
}
|
||||
@@ -1,79 +0,0 @@
|
||||
//! Yuanbao channel configuration re-exported from tinychannels.
|
||||
|
||||
pub use crate::openhuman::config::schema::YuanbaoConfig;
|
||||
|
||||
/// Default value for `DeviceInfo.app_version` (server-side `plugin_version`).
|
||||
pub(crate) const DEFAULT_PLUGIN_VERSION: &str = "0.1.0";
|
||||
|
||||
/// Strip legacy `openhuman/` prefix from version strings in config/TOML.
|
||||
pub(crate) fn strip_version_prefix(version: &str) -> &str {
|
||||
tinychannels::config::strip_yuanbao_version_prefix(version)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn default_config_is_invalid() {
|
||||
let cfg = YuanbaoConfig::default();
|
||||
assert!(cfg.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn env_defaults_fill_domains() {
|
||||
let mut cfg = YuanbaoConfig::default();
|
||||
cfg.apply_env_defaults();
|
||||
assert!(cfg.api_domain.contains("yuanbao.tencent.com"));
|
||||
assert!(cfg.ws_domain.contains("bot-wss.yuanbao.tencent.com"));
|
||||
assert!(cfg.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_only_config_can_skip_api_domain() {
|
||||
let mut cfg = YuanbaoConfig {
|
||||
app_key: "key".into(),
|
||||
token: "token".into(),
|
||||
..YuanbaoConfig::default()
|
||||
};
|
||||
cfg.apply_env_defaults();
|
||||
cfg.api_domain.clear();
|
||||
assert!(cfg.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pre_env_uses_pre_domains() {
|
||||
let mut cfg = YuanbaoConfig {
|
||||
env: "pre".into(),
|
||||
..YuanbaoConfig::default()
|
||||
};
|
||||
cfg.apply_env_defaults();
|
||||
assert!(cfg.api_domain.contains("bot-pre"));
|
||||
assert!(cfg.ws_domain.contains("bot-wss-pre"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_domains_are_preserved() {
|
||||
let mut cfg = YuanbaoConfig {
|
||||
api_domain: "https://api.example.test".into(),
|
||||
ws_domain: "wss://ws.example.test".into(),
|
||||
..YuanbaoConfig::default()
|
||||
};
|
||||
cfg.apply_env_defaults();
|
||||
assert_eq!(cfg.api_domain, "https://api.example.test");
|
||||
assert_eq!(cfg.ws_domain, "wss://ws.example.test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn plugin_version_alias_deserializes() {
|
||||
let cfg: YuanbaoConfig = toml::from_str(
|
||||
r#"
|
||||
app_key = "key"
|
||||
app_secret = "secret"
|
||||
plugin_version = "openhuman/9.9.9"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(strip_version_prefix(&cfg.bot_version), "9.9.9");
|
||||
}
|
||||
}
|
||||
@@ -1,903 +0,0 @@
|
||||
//! Yuanbao WebSocket connection manager.
|
||||
//!
|
||||
//! Owns one WebSocket to the gateway and runs:
|
||||
//! 1. token sign-fetch (via [`SignManager`]) → `auth-bind` handshake
|
||||
//! 2. periodic `ping` heartbeats
|
||||
//! 3. inbound frame fan-out (decoded `ConnFrame` → mpsc)
|
||||
//! 4. outbound request/response correlation via per-`msg_id` oneshot
|
||||
//! 5. exponential-backoff reconnect with a no-retry close-code allowlist
|
||||
//!
|
||||
//! All public APIs are `&self` so the connection can be wrapped in
|
||||
//! `Arc<…>` and shared between the listen loop and outbound senders.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use parking_lot::Mutex as ParkingMutex;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::{mpsc, oneshot, watch, Mutex};
|
||||
use tokio::time;
|
||||
use tokio_tungstenite::{
|
||||
connect_async, tungstenite::protocol::Message, MaybeTlsStream, WebSocketStream,
|
||||
};
|
||||
use tracing::{error, info, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::config::YuanbaoConfig;
|
||||
use super::errors::{YuanbaoError, NO_RECONNECT_CLOSE_CODES};
|
||||
use super::proto::{
|
||||
decode_auth_bind_rsp, decode_conn_msg, encode_auth_bind, encode_ping, encode_push_ack,
|
||||
};
|
||||
use super::proto_constants::*;
|
||||
use super::sign::SignManager;
|
||||
use super::types::{Account, ConnFrame, ConnectionState};
|
||||
|
||||
type WsSender =
|
||||
futures_util::stream::SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>;
|
||||
|
||||
/// One inbound event delivered to the listen loop.
|
||||
pub enum InboundEvent {
|
||||
/// A regular biz push.
|
||||
Push(ConnFrame),
|
||||
/// Server told us we were kicked off.
|
||||
Kickout(String),
|
||||
}
|
||||
|
||||
/// In-flight outbound request awaiting a matching `Response` frame.
|
||||
type PendingMap = HashMap<String, oneshot::Sender<ConnFrame>>;
|
||||
|
||||
/// Long-lived connection manager.
|
||||
pub struct YuanbaoConnection {
|
||||
config: YuanbaoConfig,
|
||||
state: ParkingMutex<ConnectionState>,
|
||||
is_connected: AtomicBool,
|
||||
msg_id_seq: AtomicU64,
|
||||
sender: Mutex<Option<WsSender>>,
|
||||
inbound_tx: mpsc::UnboundedSender<InboundEvent>,
|
||||
account: ParkingMutex<Account>,
|
||||
sign_manager: Option<Arc<SignManager>>,
|
||||
pending: ParkingMutex<PendingMap>,
|
||||
}
|
||||
|
||||
impl YuanbaoConnection {
|
||||
pub fn new(
|
||||
config: YuanbaoConfig,
|
||||
inbound_tx: mpsc::UnboundedSender<InboundEvent>,
|
||||
sign_manager: Option<Arc<SignManager>>,
|
||||
) -> Arc<Self> {
|
||||
let initial_account = Account {
|
||||
uid: config.bot_id.clone(),
|
||||
..Default::default()
|
||||
};
|
||||
Arc::new(Self {
|
||||
config,
|
||||
state: ParkingMutex::new(ConnectionState::Disconnected),
|
||||
is_connected: AtomicBool::new(false),
|
||||
msg_id_seq: AtomicU64::new(1),
|
||||
sender: Mutex::new(None),
|
||||
inbound_tx,
|
||||
account: ParkingMutex::new(initial_account),
|
||||
sign_manager,
|
||||
pending: ParkingMutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn is_connected(&self) -> bool {
|
||||
self.is_connected.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn state(&self) -> ConnectionState {
|
||||
*self.state.lock()
|
||||
}
|
||||
|
||||
fn set_state(&self, new: ConnectionState) {
|
||||
*self.state.lock() = new;
|
||||
self.is_connected
|
||||
.store(matches!(new, ConnectionState::Connected), Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Current account info (best-effort — empty fields until auth-bind succeeds).
|
||||
pub fn account(&self) -> Account {
|
||||
self.account.lock().clone()
|
||||
}
|
||||
|
||||
fn update_account(&self, f: impl FnOnce(&mut Account)) {
|
||||
let mut g = self.account.lock();
|
||||
f(&mut g);
|
||||
}
|
||||
|
||||
/// Per-process monotonic application msg_id.
|
||||
pub fn next_msg_id(&self, prefix: &str) -> String {
|
||||
let n = self.msg_id_seq.fetch_add(1, Ordering::Relaxed);
|
||||
format!("{prefix}_{n}")
|
||||
}
|
||||
|
||||
/// Send a raw binary frame. Returns `NotConnected` if the connection
|
||||
/// isn't currently up.
|
||||
pub async fn send_frame(&self, data: Vec<u8>) -> Result<(), YuanbaoError> {
|
||||
let mut guard = self.sender.lock().await;
|
||||
match guard.as_mut() {
|
||||
Some(s) => s
|
||||
.send(Message::Binary(data))
|
||||
.await
|
||||
.map_err(|e| YuanbaoError::WebSocket(e.to_string())),
|
||||
None => Err(YuanbaoError::NotConnected),
|
||||
}
|
||||
}
|
||||
|
||||
/// Send an already-encoded `ConnMsg` (alias of `send_frame`).
|
||||
pub async fn send_conn_msg(&self, frame_bytes: Vec<u8>) -> Result<(), YuanbaoError> {
|
||||
self.send_frame(frame_bytes).await
|
||||
}
|
||||
|
||||
/// Send a request and wait for the matching `Response` (correlated by
|
||||
/// `msg_id`). Times out after `timeout` and removes the pending entry.
|
||||
pub async fn send_and_wait(
|
||||
&self,
|
||||
msg_id: &str,
|
||||
frame_bytes: Vec<u8>,
|
||||
timeout: Duration,
|
||||
) -> Result<ConnFrame, YuanbaoError> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
{
|
||||
let mut p = self.pending.lock();
|
||||
p.insert(msg_id.to_string(), tx);
|
||||
}
|
||||
if let Err(e) = self.send_frame(frame_bytes).await {
|
||||
self.pending.lock().remove(msg_id);
|
||||
return Err(e);
|
||||
}
|
||||
match tokio::time::timeout(timeout, rx).await {
|
||||
Ok(Ok(frame)) => Ok(frame),
|
||||
Ok(Err(_)) => {
|
||||
self.pending.lock().remove(msg_id);
|
||||
Err(YuanbaoError::SendFailed(format!(
|
||||
"correlator channel closed for msg_id={msg_id}"
|
||||
)))
|
||||
}
|
||||
Err(_) => {
|
||||
self.pending.lock().remove(msg_id);
|
||||
Err(YuanbaoError::Timeout(format!("msg_id={msg_id}")))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trigger a graceful shutdown.
|
||||
pub async fn shutdown(&self) {
|
||||
let mut guard = self.sender.lock().await;
|
||||
if let Some(mut s) = guard.take() {
|
||||
let _ = s.send(Message::Close(None)).await;
|
||||
let _ = s.close().await;
|
||||
}
|
||||
// Drop all pending waiters so callers stop hanging.
|
||||
let mut pending = self.pending.lock();
|
||||
pending.clear();
|
||||
self.set_state(ConnectionState::Disconnected);
|
||||
}
|
||||
|
||||
/// Main reconnection loop. Returns when `shutdown` flips to `true`.
|
||||
pub async fn run(self: Arc<Self>, mut shutdown: watch::Receiver<bool>) {
|
||||
let max_attempts = if self.config.max_reconnect_attempts > 0 {
|
||||
self.config.max_reconnect_attempts
|
||||
} else {
|
||||
MAX_RECONNECT_ATTEMPTS
|
||||
};
|
||||
let mut attempt: u32 = 0;
|
||||
|
||||
loop {
|
||||
if *shutdown.borrow() {
|
||||
info!("[yuanbao] shutdown signaled, stopping connection loop");
|
||||
self.shutdown().await;
|
||||
return;
|
||||
}
|
||||
if attempt >= max_attempts {
|
||||
error!("[yuanbao] giving up after {} reconnect attempts", attempt);
|
||||
return;
|
||||
}
|
||||
|
||||
self.set_state(if attempt == 0 {
|
||||
ConnectionState::Connecting
|
||||
} else {
|
||||
ConnectionState::Reconnecting
|
||||
});
|
||||
|
||||
let outcome = self.connect_once(&mut shutdown).await;
|
||||
match outcome {
|
||||
Ok(Some(code)) if NO_RECONNECT_CLOSE_CODES.contains(&code) => {
|
||||
error!("[yuanbao] no-reconnect close code {} — stopping", code);
|
||||
return;
|
||||
}
|
||||
Ok(close_code) => {
|
||||
// Successful connection: reset the attempt counter so
|
||||
// intermittent disconnects don't permanently exhaust the
|
||||
// reconnect budget.
|
||||
attempt = 0;
|
||||
info!("[yuanbao] connection closed (code={:?})", close_code);
|
||||
}
|
||||
Err(e) => warn!("[yuanbao] connection error: {}", e),
|
||||
}
|
||||
|
||||
self.set_state(ConnectionState::Disconnected);
|
||||
*self.sender.lock().await = None;
|
||||
self.pending.lock().clear();
|
||||
|
||||
// `connect_once` may have returned because shutdown fired inside
|
||||
// its read loop. In that case we must not sleep through the
|
||||
// reconnect backoff — exit immediately so stop is responsive.
|
||||
if *shutdown.borrow() {
|
||||
info!("[yuanbao] shutdown signaled, stopping connection loop");
|
||||
self.shutdown().await;
|
||||
return;
|
||||
}
|
||||
|
||||
attempt += 1;
|
||||
let delay = backoff_seconds(attempt);
|
||||
info!(
|
||||
"[yuanbao] reconnecting in {}s (attempt {}/{})",
|
||||
delay, attempt, max_attempts
|
||||
);
|
||||
tokio::select! {
|
||||
_ = time::sleep(Duration::from_secs(delay)) => {}
|
||||
_ = shutdown.changed() => {
|
||||
info!("[yuanbao] shutdown received during backoff");
|
||||
self.shutdown().await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn connect_once(
|
||||
&self,
|
||||
shutdown: &mut watch::Receiver<bool>,
|
||||
) -> Result<Option<u16>, YuanbaoError> {
|
||||
// Resolve token (may hit the sign endpoint).
|
||||
let (token, bot_id, source) = self.resolve_token().await?;
|
||||
if !bot_id.is_empty() {
|
||||
self.update_account(|a| {
|
||||
if a.uid.is_empty() {
|
||||
a.uid = bot_id.clone();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
let url = &self.config.ws_domain;
|
||||
info!("[yuanbao] connecting to {}", url);
|
||||
let (ws_stream, _resp) = connect_async(url)
|
||||
.await
|
||||
.map_err(|e| YuanbaoError::WebSocket(e.to_string()))?;
|
||||
|
||||
let (sender, mut receiver) = ws_stream.split();
|
||||
*self.sender.lock().await = Some(sender);
|
||||
info!("[yuanbao] WebSocket connected — sending auth-bind");
|
||||
|
||||
self.set_state(ConnectionState::Authenticating);
|
||||
self.send_auth_bind(&token, &bot_id, &source).await?;
|
||||
|
||||
// Wait for auth-bind response.
|
||||
let auth_timeout = Duration::from_secs(AUTH_TIMEOUT_SECS);
|
||||
let auth_msg = tokio::time::timeout(auth_timeout, receiver.next())
|
||||
.await
|
||||
.map_err(|_| YuanbaoError::AuthTimeout)?
|
||||
.ok_or_else(|| YuanbaoError::WebSocket("closed during auth-bind".into()))?
|
||||
.map_err(|e| YuanbaoError::WebSocket(e.to_string()))?;
|
||||
|
||||
self.handle_auth_response(&auth_msg)?;
|
||||
self.set_state(ConnectionState::Connected);
|
||||
info!("[yuanbao] auth-bind successful, entering read loop");
|
||||
|
||||
let ping_secs = if self.config.heartbeat_interval_secs > 0 {
|
||||
self.config.heartbeat_interval_secs
|
||||
} else {
|
||||
PING_INTERVAL_SECS
|
||||
};
|
||||
let mut ping_interval = time::interval(Duration::from_secs(ping_secs));
|
||||
ping_interval.tick().await; // skip first tick
|
||||
|
||||
let mut close_code: Option<u16> = None;
|
||||
let mut consecutive_ping_failures: u32 = 0;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown.changed() => {
|
||||
info!("[yuanbao] shutdown received in read loop");
|
||||
return Ok(None);
|
||||
}
|
||||
_ = ping_interval.tick() => {
|
||||
let msg_id = self.next_msg_id("ping");
|
||||
let frame = encode_ping(&msg_id);
|
||||
if let Err(e) = self.send_frame(frame).await {
|
||||
warn!("[yuanbao] ping send failed: {}", e);
|
||||
consecutive_ping_failures += 1;
|
||||
if consecutive_ping_failures >= HEARTBEAT_TIMEOUT_THRESHOLD {
|
||||
warn!(
|
||||
"[yuanbao] {} consecutive ping failures — dropping",
|
||||
consecutive_ping_failures
|
||||
);
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
consecutive_ping_failures = 0;
|
||||
}
|
||||
}
|
||||
msg = receiver.next() => {
|
||||
match msg {
|
||||
Some(Ok(Message::Binary(data))) => self.handle_binary(data).await,
|
||||
Some(Ok(Message::Close(frame))) => {
|
||||
close_code = frame.map(|f| u16::from(f.code));
|
||||
info!("[yuanbao] received close frame: {:?}", close_code);
|
||||
break;
|
||||
}
|
||||
Some(Ok(Message::Ping(payload))) => {
|
||||
let mut guard = self.sender.lock().await;
|
||||
if let Some(s) = guard.as_mut() {
|
||||
let _ = s.send(Message::Pong(payload)).await;
|
||||
}
|
||||
}
|
||||
Some(Ok(_)) => {}
|
||||
Some(Err(e)) => {
|
||||
warn!("[yuanbao] websocket read error: {}", e);
|
||||
break;
|
||||
}
|
||||
None => {
|
||||
info!("[yuanbao] websocket stream ended");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(close_code)
|
||||
}
|
||||
|
||||
async fn resolve_token(&self) -> Result<(String, String, String), YuanbaoError> {
|
||||
let cfg = &self.config;
|
||||
if !cfg.token.is_empty() {
|
||||
// Pre-signed token: no source returned by the sign endpoint.
|
||||
// Mirrors yuanbao-openclaw-plugin's static-token branch, which
|
||||
// returns source="bot".
|
||||
return Ok((cfg.token.clone(), cfg.bot_id.clone(), String::new()));
|
||||
}
|
||||
let mgr = self
|
||||
.sign_manager
|
||||
.as_ref()
|
||||
.ok_or_else(|| YuanbaoError::AuthFailed("no token and no SignManager".into()))?;
|
||||
if cfg.app_secret.is_empty() {
|
||||
return Err(YuanbaoError::AuthFailed(
|
||||
"app_secret required to sign".into(),
|
||||
));
|
||||
}
|
||||
let entry = mgr
|
||||
.get_token(
|
||||
&cfg.app_key,
|
||||
&cfg.app_secret,
|
||||
&cfg.api_domain,
|
||||
&cfg.route_env,
|
||||
)
|
||||
.await?;
|
||||
Ok((entry.token, entry.bot_id, entry.source))
|
||||
}
|
||||
|
||||
async fn send_auth_bind(
|
||||
&self,
|
||||
token: &str,
|
||||
bot_id: &str,
|
||||
source: &str,
|
||||
) -> Result<(), YuanbaoError> {
|
||||
let cfg = &self.config;
|
||||
let uid = if bot_id.is_empty() {
|
||||
self.account.lock().uid.clone()
|
||||
} else {
|
||||
bot_id.to_string()
|
||||
};
|
||||
let msg_id = format!("auth_{}", Uuid::new_v4());
|
||||
// Auth-bind payload aligned with yuanbao-openclaw-plugin:
|
||||
// biz_id = "ybBot" (server rejects raw app_key with 40011).
|
||||
// source comes from the sign endpoint response; fall back to
|
||||
// "bot" when missing (matches the plugin's static-token branch
|
||||
// and `data.source || "bot"` resolution).
|
||||
let resolved_source = if source.is_empty() { "bot" } else { source };
|
||||
// Align with yuanbao-openclaw-plugin: app_version → plugin_version,
|
||||
// DeviceInfo field 24 → bot_version (OpenHuman framework / CARGO_PKG_VERSION).
|
||||
let plugin_version = super::config::strip_version_prefix(&cfg.bot_version);
|
||||
let framework_version = env!("CARGO_PKG_VERSION");
|
||||
let frame = encode_auth_bind(
|
||||
"ybBot",
|
||||
&uid,
|
||||
resolved_source,
|
||||
token,
|
||||
&msg_id,
|
||||
plugin_version,
|
||||
std::env::consts::OS,
|
||||
framework_version,
|
||||
&cfg.route_env,
|
||||
);
|
||||
self.send_frame(frame).await
|
||||
}
|
||||
|
||||
fn handle_auth_response(&self, msg: &Message) -> Result<(), YuanbaoError> {
|
||||
let data = match msg {
|
||||
Message::Binary(b) => b,
|
||||
_ => {
|
||||
return Err(YuanbaoError::AuthFailed(
|
||||
"expected binary auth-bind response".into(),
|
||||
))
|
||||
}
|
||||
};
|
||||
let frame = decode_conn_msg(data)?;
|
||||
if frame.cmd != cmd::AUTH_BIND {
|
||||
return Err(YuanbaoError::AuthFailed(format!(
|
||||
"unexpected cmd in auth response: {:?}",
|
||||
frame.cmd
|
||||
)));
|
||||
}
|
||||
if frame.status != 0 {
|
||||
return Err(YuanbaoError::AuthFailed(format!(
|
||||
"auth rejected: status={}",
|
||||
frame.status
|
||||
)));
|
||||
}
|
||||
// Body carries code/message/connect_id — back-fill the account.
|
||||
if !frame.data.is_empty() {
|
||||
let rsp = decode_auth_bind_rsp(&frame.data)?;
|
||||
if rsp.code != 0 {
|
||||
return Err(YuanbaoError::AuthFailed(format!(
|
||||
"auth-bind code={} message={}",
|
||||
rsp.code, rsp.message
|
||||
)));
|
||||
}
|
||||
if !rsp.connect_id.is_empty() {
|
||||
self.update_account(|a| a.connect_id = rsp.connect_id.clone());
|
||||
info!("[yuanbao] auth-bind connect_id={}", rsp.connect_id);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_binary(&self, data: Vec<u8>) {
|
||||
let frame = match decode_conn_msg(&data) {
|
||||
Ok(f) => f,
|
||||
Err(e) => {
|
||||
warn!("[yuanbao] failed to decode binary frame: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
info!(
|
||||
"[yuanbao] rx cmd={} module={} cmd_type={} seq={} msg_id={} data_len={}",
|
||||
frame.cmd,
|
||||
frame.module,
|
||||
frame.cmd_type,
|
||||
frame.seq_no,
|
||||
frame.msg_id,
|
||||
frame.data.len()
|
||||
);
|
||||
|
||||
// Responses → match against pending requests via msg_id.
|
||||
if frame.cmd_type == cmd_type::RESPONSE {
|
||||
if !frame.msg_id.is_empty() {
|
||||
if let Some(tx) = self.pending.lock().remove(&frame.msg_id) {
|
||||
let _ = tx.send(frame);
|
||||
return;
|
||||
}
|
||||
}
|
||||
info!(
|
||||
"[yuanbao] response with no waiter cmd={} msg_id={}",
|
||||
frame.cmd, frame.msg_id
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// For server-driven pushes, ACK first when the head asks for it.
|
||||
if frame.cmd_type == cmd_type::PUSH && frame.need_ack {
|
||||
let ack = encode_push_ack(&frame);
|
||||
if let Err(e) = self.send_frame(ack).await {
|
||||
warn!("[yuanbao] failed to send PushAck: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
// Handle conn-level builtin pushes inline.
|
||||
if frame.cmd == cmd::KICKOUT {
|
||||
let reason = String::from_utf8_lossy(&frame.data).into_owned();
|
||||
warn!("[yuanbao] kickout received: {}", reason);
|
||||
let _ = self.inbound_tx.send(InboundEvent::Kickout(reason));
|
||||
return;
|
||||
}
|
||||
if frame.cmd == cmd::UPDATE_META {
|
||||
return;
|
||||
}
|
||||
|
||||
if frame.cmd_type != cmd_type::PUSH {
|
||||
info!(
|
||||
"[yuanbao] dropping non-push frame cmd_type={} cmd={}",
|
||||
frame.cmd_type, frame.cmd
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
info!(
|
||||
"[yuanbao] push forwarded to listener cmd={} module={} seq={}",
|
||||
frame.cmd, frame.module, frame.seq_no
|
||||
);
|
||||
if self.inbound_tx.send(InboundEvent::Push(frame)).is_err() {
|
||||
error!("[yuanbao] inbound channel closed — listener gone");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Backoff schedule used by `run()`. After the configured table is
|
||||
/// exhausted we cap at the last entry forever (until the attempt budget
|
||||
/// trips). Indexing is 1-based so attempt=1 → table[0].
|
||||
fn backoff_seconds(attempt: u32) -> u64 {
|
||||
let idx = attempt.saturating_sub(1) as usize;
|
||||
if idx < RECONNECT_DELAYS.len() {
|
||||
RECONNECT_DELAYS[idx]
|
||||
} else {
|
||||
*RECONNECT_DELAYS.last().unwrap_or(&60)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(test, debug_assertions))]
|
||||
pub mod test_support {
|
||||
use super::*;
|
||||
use crate::openhuman::channels::providers::yuanbao::proto::encode_conn_msg;
|
||||
use crate::openhuman::channels::providers::yuanbao::wire::{
|
||||
encode_field_bytes, encode_field_string, encode_field_varint,
|
||||
};
|
||||
|
||||
fn cfg() -> YuanbaoConfig {
|
||||
let mut c = YuanbaoConfig::default();
|
||||
c.app_key = "ak".into();
|
||||
c.ws_domain = "wss://example".into();
|
||||
c.token = "tok".into();
|
||||
c.bot_id = "bot1".into();
|
||||
c
|
||||
}
|
||||
|
||||
pub fn auth_response_success_connect_id_for_test() -> Result<String, YuanbaoError> {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
let mut body = Vec::new();
|
||||
encode_field_varint(1, 0, &mut body);
|
||||
encode_field_string(2, "ok", &mut body);
|
||||
encode_field_string(3, "connect-123", &mut body);
|
||||
let msg = Message::Binary(encode_conn_msg(
|
||||
cmd_type::RESPONSE,
|
||||
cmd::AUTH_BIND,
|
||||
1,
|
||||
"auth-1",
|
||||
module::CONN_ACCESS,
|
||||
&body,
|
||||
));
|
||||
|
||||
conn.handle_auth_response(&msg)?;
|
||||
Ok(conn.account().connect_id)
|
||||
}
|
||||
|
||||
pub fn auth_response_rejects_status_for_test() -> String {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
|
||||
let mut head = Vec::new();
|
||||
encode_field_varint(1, cmd_type::RESPONSE as u64, &mut head);
|
||||
encode_field_string(2, cmd::AUTH_BIND, &mut head);
|
||||
encode_field_string(4, "auth-2", &mut head);
|
||||
encode_field_string(5, module::CONN_ACCESS, &mut head);
|
||||
encode_field_varint(10, 401, &mut head);
|
||||
|
||||
let mut frame = Vec::new();
|
||||
encode_field_bytes(1, &head, &mut frame);
|
||||
let msg = Message::Binary(frame);
|
||||
|
||||
match conn.handle_auth_response(&msg).unwrap_err() {
|
||||
YuanbaoError::AuthFailed(message) => message,
|
||||
other => format!("{other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn handle_binary_routes_builtin_and_push_frames_for_test() -> Vec<String> {
|
||||
let (tx, mut rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
|
||||
conn.handle_binary(encode_conn_msg(
|
||||
cmd_type::RESPONSE,
|
||||
biz_cmd::QUERY_GROUP_INFO,
|
||||
1,
|
||||
"orphan-response",
|
||||
module::BIZ_PKG,
|
||||
b"response",
|
||||
))
|
||||
.await;
|
||||
conn.handle_binary(encode_conn_msg(
|
||||
cmd_type::PUSH,
|
||||
cmd::UPDATE_META,
|
||||
2,
|
||||
"meta",
|
||||
module::CONN_ACCESS,
|
||||
b"ignored",
|
||||
))
|
||||
.await;
|
||||
conn.handle_binary(encode_conn_msg(
|
||||
cmd_type::REQUEST,
|
||||
biz_cmd::SEND_C2C_MESSAGE,
|
||||
3,
|
||||
"request",
|
||||
module::BIZ_PKG,
|
||||
b"not-a-push",
|
||||
))
|
||||
.await;
|
||||
conn.handle_binary(encode_conn_msg(
|
||||
cmd_type::PUSH,
|
||||
cmd::KICKOUT,
|
||||
4,
|
||||
"kick",
|
||||
module::CONN_ACCESS,
|
||||
b"logged out",
|
||||
))
|
||||
.await;
|
||||
conn.handle_binary(encode_conn_msg(
|
||||
cmd_type::PUSH,
|
||||
"incoming-message",
|
||||
5,
|
||||
"push-1",
|
||||
module::BIZ_PKG,
|
||||
b"payload",
|
||||
))
|
||||
.await;
|
||||
|
||||
let mut events = Vec::new();
|
||||
while let Ok(event) = rx.try_recv() {
|
||||
match event {
|
||||
InboundEvent::Kickout(reason) => events.push(format!("kickout:{reason}")),
|
||||
InboundEvent::Push(frame) => {
|
||||
events.push(format!("push:{}:{}", frame.cmd, frame.msg_id));
|
||||
}
|
||||
}
|
||||
}
|
||||
events
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn backoff_follows_schedule() {
|
||||
assert_eq!(backoff_seconds(1), 1);
|
||||
assert_eq!(backoff_seconds(2), 2);
|
||||
assert_eq!(backoff_seconds(3), 5);
|
||||
assert_eq!(backoff_seconds(6), 60);
|
||||
assert_eq!(backoff_seconds(100), 60);
|
||||
assert_eq!(backoff_seconds(0), 1);
|
||||
}
|
||||
|
||||
fn cfg() -> YuanbaoConfig {
|
||||
let mut c = YuanbaoConfig::default();
|
||||
c.app_key = "ak".into();
|
||||
c.ws_domain = "wss://example".into();
|
||||
c.token = "tok".into();
|
||||
c.bot_id = "bot1".into();
|
||||
c
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pending_correlator_times_out() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
let err = conn
|
||||
.send_and_wait("missing_id", vec![1, 2, 3], Duration::from_millis(20))
|
||||
.await
|
||||
.unwrap_err();
|
||||
// Without a connected socket, send_frame fails first → SendFailed/NotConnected.
|
||||
assert!(matches!(
|
||||
err,
|
||||
YuanbaoError::NotConnected | YuanbaoError::Timeout(_) | YuanbaoError::SendFailed(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_back_fill_picks_up_uid() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
assert_eq!(conn.account().uid, "bot1");
|
||||
conn.update_account(|a| a.connect_id = "cid_xyz".into());
|
||||
assert_eq!(conn.account().connect_id, "cid_xyz");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn next_msg_id_is_monotonic_and_prefixed() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
let a = conn.next_msg_id("pfx");
|
||||
let b = conn.next_msg_id("pfx");
|
||||
assert!(a.starts_with("pfx_"));
|
||||
assert!(b.starts_with("pfx_"));
|
||||
// Suffix is monotonically increasing.
|
||||
let na: u64 = a.strip_prefix("pfx_").unwrap().parse().unwrap();
|
||||
let nb: u64 = b.strip_prefix("pfx_").unwrap().parse().unwrap();
|
||||
assert!(nb > na);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initial_state_is_disconnected() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
assert_eq!(conn.state(), ConnectionState::Disconnected);
|
||||
assert!(!conn.is_connected());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn set_state_connected_flips_is_connected_flag() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
conn.set_state(ConnectionState::Connected);
|
||||
assert_eq!(conn.state(), ConnectionState::Connected);
|
||||
assert!(conn.is_connected());
|
||||
conn.set_state(ConnectionState::Reconnecting);
|
||||
assert_eq!(conn.state(), ConnectionState::Reconnecting);
|
||||
assert!(!conn.is_connected());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_frame_without_socket_returns_not_connected() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
let err = conn.send_frame(vec![1, 2, 3]).await.unwrap_err();
|
||||
assert!(matches!(err, YuanbaoError::NotConnected));
|
||||
let err2 = conn.send_conn_msg(vec![4]).await.unwrap_err();
|
||||
assert!(matches!(err2, YuanbaoError::NotConnected));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_clears_pending_and_sets_disconnected() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
conn.set_state(ConnectionState::Connected);
|
||||
// Drop a phantom pending entry then shutdown.
|
||||
let (phantom_tx, _phantom_rx) = oneshot::channel();
|
||||
conn.pending.lock().insert("ghost".into(), phantom_tx);
|
||||
conn.shutdown().await;
|
||||
assert_eq!(conn.state(), ConnectionState::Disconnected);
|
||||
assert!(!conn.is_connected());
|
||||
assert!(conn.pending.lock().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backoff_caps_at_last_entry_for_huge_attempts() {
|
||||
let last = *RECONNECT_DELAYS.last().unwrap();
|
||||
assert_eq!(backoff_seconds(RECONNECT_DELAYS.len() as u32 + 5), last);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_token_uses_static_token_when_present() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
let (token, bot_id, source) = conn.resolve_token().await.unwrap();
|
||||
assert_eq!(token, "tok");
|
||||
assert_eq!(bot_id, "bot1");
|
||||
assert_eq!(source, "");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_token_without_token_and_without_sign_manager_errors() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let mut c = cfg();
|
||||
c.token = String::new();
|
||||
let conn = YuanbaoConnection::new(c, tx, None);
|
||||
match conn.resolve_token().await.unwrap_err() {
|
||||
YuanbaoError::AuthFailed(m) => assert!(m.contains("no token"), "got {m}"),
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_token_with_sign_manager_but_no_app_secret_errors() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let mut c = cfg();
|
||||
c.token = String::new();
|
||||
c.app_secret = String::new();
|
||||
let mgr = SignManager::new(reqwest::Client::new());
|
||||
let conn = YuanbaoConnection::new(c, tx, Some(mgr));
|
||||
match conn.resolve_token().await.unwrap_err() {
|
||||
YuanbaoError::AuthFailed(m) => assert!(m.contains("app_secret"), "got {m}"),
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_auth_bind_without_socket_returns_not_connected() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
let err = conn.send_auth_bind("tok", "bot1", "bot").await.unwrap_err();
|
||||
assert!(matches!(err, YuanbaoError::NotConnected));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_auth_bind_falls_back_to_account_uid_when_bot_id_empty() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
// bot_id="" → reads from account.uid (which was seeded from cfg.bot_id="bot1")
|
||||
let err = conn.send_auth_bind("tok", "", "").await.unwrap_err();
|
||||
assert!(matches!(err, YuanbaoError::NotConnected));
|
||||
// Account uid still in place.
|
||||
assert_eq!(conn.account().uid, "bot1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_auth_response_rejects_non_binary_message() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
let msg = Message::Text("nope".into());
|
||||
match conn.handle_auth_response(&msg).unwrap_err() {
|
||||
YuanbaoError::AuthFailed(m) => {
|
||||
assert!(m.contains("binary"), "got {m}")
|
||||
}
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_auth_response_rejects_undecodable_binary() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
// Wholly invalid wire data — decode_conn_msg fails.
|
||||
let msg = Message::Binary(vec![0xFF, 0xFF, 0xFF, 0xFF]);
|
||||
let err = conn.handle_auth_response(&msg).unwrap_err();
|
||||
// Either Proto decode error or some other surface — must not be Ok.
|
||||
assert!(
|
||||
!matches!(err, YuanbaoError::AuthFailed(_) if format!("{err:?}").contains("binary"))
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression guard for the post-`connect_once` shutdown short-circuit:
|
||||
/// once shutdown is signaled, `run()` must not block on the reconnect
|
||||
/// backoff. We force connect_once to fail synchronously (invalid WS URL),
|
||||
/// then signal shutdown — total runtime must be well under the first
|
||||
/// backoff slot (`backoff_seconds(1) == 1s`).
|
||||
#[tokio::test]
|
||||
async fn run_exits_promptly_after_shutdown_signal() {
|
||||
use std::time::Instant;
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let mut c = cfg();
|
||||
// tokio-tungstenite rejects the URL synchronously — connect_once
|
||||
// returns Err in microseconds, putting `run()` on the post-connect
|
||||
// cleanup path that the fix targets.
|
||||
c.ws_domain = "not-a-valid-ws-url".to_string();
|
||||
c.max_reconnect_attempts = 100;
|
||||
let conn = YuanbaoConnection::new(c, tx, None);
|
||||
let (sd_tx, sd_rx) = watch::channel(false);
|
||||
|
||||
let handle = tokio::spawn(conn.clone().run(sd_rx));
|
||||
// Let `run()` enter the loop and attempt connect_once at least once.
|
||||
time::sleep(Duration::from_millis(20)).await;
|
||||
|
||||
let started = Instant::now();
|
||||
sd_tx.send(true).unwrap();
|
||||
|
||||
// The first reconnect backoff slot is 1s. Without responsive
|
||||
// shutdown handling, run() would sleep through it before checking
|
||||
// the flag. 500ms gives us comfortable headroom while staying
|
||||
// far enough below the backoff to detect a regression.
|
||||
let res = time::timeout(Duration::from_millis(500), handle).await;
|
||||
res.expect("run() did not exit within 500ms of shutdown signal")
|
||||
.expect("run() task panicked");
|
||||
assert!(
|
||||
started.elapsed() < Duration::from_millis(500),
|
||||
"run() took {:?} to exit after shutdown — backoff was not skipped",
|
||||
started.elapsed()
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handle_binary_with_garbage_does_not_panic() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let conn = YuanbaoConnection::new(cfg(), tx, None);
|
||||
// Should silently log + return — no panic.
|
||||
conn.handle_binary(vec![0xFF, 0xFF, 0xFF, 0xFF]).await;
|
||||
}
|
||||
}
|
||||
@@ -1,567 +0,0 @@
|
||||
//! Tencent COS upload — HMAC-SHA1 signing and `genUploadInfo` flow.
|
||||
//!
|
||||
//! Split out of `media.rs` to stay under the 500-line per-file ceiling.
|
||||
//! Reference: <https://cloud.tencent.com/document/product/436/7778>.
|
||||
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha1::{Digest, Sha1};
|
||||
use tracing::{debug, info};
|
||||
|
||||
use super::errors::YuanbaoError;
|
||||
use super::media::{guess_mime_type, is_image, parse_image_size, ImageDims};
|
||||
|
||||
const UPLOAD_INFO_PATH: &str = "/api/resource/genUploadInfo";
|
||||
const COS_USE_ACCELERATE: bool = true;
|
||||
|
||||
type HmacSha1 = Hmac<Sha1>;
|
||||
|
||||
fn hmac_sha1_hex(key: &[u8], msg: &[u8]) -> String {
|
||||
let mut mac = HmacSha1::new_from_slice(key).expect("HMAC accepts any key length");
|
||||
mac.update(msg);
|
||||
hex::encode(mac.finalize().into_bytes())
|
||||
}
|
||||
|
||||
fn sha1_hex(msg: &[u8]) -> String {
|
||||
let mut hasher = Sha1::new();
|
||||
hasher.update(msg);
|
||||
hex::encode(hasher.finalize())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CosSignInput<'a> {
|
||||
pub method: &'a str,
|
||||
/// URL-encoded path with leading `/`.
|
||||
pub path: &'a str,
|
||||
pub params: &'a [(&'a str, &'a str)],
|
||||
pub headers: &'a [(&'a str, &'a str)],
|
||||
pub secret_id: &'a str,
|
||||
pub secret_key: &'a str,
|
||||
pub start_time: u64,
|
||||
pub expire_seconds: u64,
|
||||
}
|
||||
|
||||
/// Build the COS `Authorization` header value.
|
||||
pub fn cos_sign(input: &CosSignInput<'_>) -> String {
|
||||
let q_sign_time = format!(
|
||||
"{};{}",
|
||||
input.start_time,
|
||||
input.start_time + input.expire_seconds
|
||||
);
|
||||
|
||||
// Step 1 — SignKey = HMAC-SHA1(SecretKey, q-sign-time).
|
||||
let sign_key = hmac_sha1_hex(input.secret_key.as_bytes(), q_sign_time.as_bytes());
|
||||
|
||||
// Step 2 — HttpString. Names lower-cased, values URL-encoded.
|
||||
let mut params: Vec<(String, String)> = input
|
||||
.params
|
||||
.iter()
|
||||
.map(|(k, v)| (k.to_ascii_lowercase(), url_encode(v)))
|
||||
.collect();
|
||||
params.sort();
|
||||
let mut headers: Vec<(String, String)> = input
|
||||
.headers
|
||||
.iter()
|
||||
.map(|(k, v)| (k.to_ascii_lowercase(), url_encode(v)))
|
||||
.collect();
|
||||
headers.sort();
|
||||
|
||||
let url_param_list = params
|
||||
.iter()
|
||||
.map(|(k, _)| k.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join(";");
|
||||
let url_params = params
|
||||
.iter()
|
||||
.map(|(k, v)| format!("{k}={v}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("&");
|
||||
let header_list = headers
|
||||
.iter()
|
||||
.map(|(k, _)| k.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join(";");
|
||||
let header_str = headers
|
||||
.iter()
|
||||
.map(|(k, v)| format!("{k}={v}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join("&");
|
||||
|
||||
let http_string = format!(
|
||||
"{}\n{}\n{}\n{}\n",
|
||||
input.method.to_ascii_lowercase(),
|
||||
input.path,
|
||||
url_params,
|
||||
header_str
|
||||
);
|
||||
|
||||
// Step 3 — StringToSign.
|
||||
let sha1_of_http = sha1_hex(http_string.as_bytes());
|
||||
let string_to_sign = format!("sha1\n{q_sign_time}\n{sha1_of_http}\n");
|
||||
|
||||
// Step 4 — Signature.
|
||||
let signature = hmac_sha1_hex(sign_key.as_bytes(), string_to_sign.as_bytes());
|
||||
|
||||
format!(
|
||||
"q-sign-algorithm=sha1&q-ak={sid}&q-sign-time={t}&q-key-time={t}\
|
||||
&q-header-list={hl}&q-url-param-list={pl}&q-signature={sig}",
|
||||
sid = input.secret_id,
|
||||
t = q_sign_time,
|
||||
hl = header_list,
|
||||
pl = url_param_list,
|
||||
sig = signature
|
||||
)
|
||||
}
|
||||
|
||||
fn url_encode(s: &str) -> String {
|
||||
urlencoding::encode(s).into_owned()
|
||||
}
|
||||
|
||||
fn encode_cos_key(key: &str) -> String {
|
||||
key.split('/')
|
||||
.map(|seg| urlencoding::encode(seg).into_owned())
|
||||
.collect::<Vec<_>>()
|
||||
.join("/")
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct CosCredentials {
|
||||
pub bucket: String,
|
||||
pub region: String,
|
||||
pub location: String,
|
||||
pub secret_id: String,
|
||||
pub secret_key: String,
|
||||
pub session_token: String,
|
||||
pub start_time: u64,
|
||||
pub expired_time: u64,
|
||||
pub resource_url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UploadResult {
|
||||
pub url: String,
|
||||
pub uuid: String,
|
||||
pub size: u64,
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
}
|
||||
|
||||
/// Fetch COS upload credentials from the yuanbao gateway.
|
||||
pub async fn get_cos_credentials(
|
||||
http: &reqwest::Client,
|
||||
api_domain: &str,
|
||||
app_key: &str,
|
||||
bot_id: &str,
|
||||
token: &str,
|
||||
route_env: &str,
|
||||
filename: &str,
|
||||
) -> Result<CosCredentials, YuanbaoError> {
|
||||
let upload_url = format!(
|
||||
"{}/{}",
|
||||
api_domain.trim_end_matches('/'),
|
||||
UPLOAD_INFO_PATH.trim_start_matches('/')
|
||||
);
|
||||
let body = serde_json::json!({
|
||||
"fileName": filename,
|
||||
"fileId": uuid::Uuid::new_v4().simple().to_string(),
|
||||
"docFrom": "localDoc",
|
||||
"docOpenId": "",
|
||||
});
|
||||
let mut req = http
|
||||
.post(&upload_url)
|
||||
.header("Content-Type", "application/json")
|
||||
.header("X-Token", token)
|
||||
.header("X-ID", if bot_id.is_empty() { app_key } else { bot_id })
|
||||
.header("X-Source", "web");
|
||||
if !route_env.is_empty() {
|
||||
req = req.header("X-Route-Env", route_env);
|
||||
}
|
||||
let resp = req
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| YuanbaoError::Connection(format!("genUploadInfo: {e}")))?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(YuanbaoError::Media(format!(
|
||||
"genUploadInfo HTTP {}",
|
||||
resp.status()
|
||||
)));
|
||||
}
|
||||
let payload: serde_json::Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| YuanbaoError::Media(format!("genUploadInfo body parse: {e}")))?;
|
||||
if let Some(code) = payload.get("code").and_then(|c| c.as_i64()) {
|
||||
if code != 0 {
|
||||
return Err(YuanbaoError::Media(format!(
|
||||
"genUploadInfo code={code}, msg={}",
|
||||
payload.get("msg").and_then(|m| m.as_str()).unwrap_or("")
|
||||
)));
|
||||
}
|
||||
}
|
||||
let data = payload.get("data").unwrap_or(&payload);
|
||||
let get_str = |k: &str| -> String {
|
||||
data.get(k)
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string()
|
||||
};
|
||||
let get_u64 = |k: &str| -> u64 { data.get(k).and_then(|v| v.as_u64()).unwrap_or(0) };
|
||||
|
||||
Ok(CosCredentials {
|
||||
bucket: get_str("bucketName"),
|
||||
region: get_str("region"),
|
||||
location: get_str("location"),
|
||||
secret_id: get_str("encryptTmpSecretId"),
|
||||
secret_key: get_str("encryptTmpSecretKey"),
|
||||
session_token: get_str("encryptToken"),
|
||||
start_time: get_u64("startTime"),
|
||||
expired_time: get_u64("expiredTime"),
|
||||
resource_url: get_str("resourceUrl"),
|
||||
})
|
||||
}
|
||||
|
||||
/// PUT a file to COS using credentials returned by `get_cos_credentials`.
|
||||
pub async fn upload_to_cos(
|
||||
http: &reqwest::Client,
|
||||
creds: &CosCredentials,
|
||||
data: &[u8],
|
||||
filename: &str,
|
||||
mut content_type: String,
|
||||
) -> Result<UploadResult, YuanbaoError> {
|
||||
if creds.secret_id.is_empty() || creds.secret_key.is_empty() || creds.location.is_empty() {
|
||||
return Err(YuanbaoError::Media(
|
||||
"COS credentials missing secret_id / secret_key / location".into(),
|
||||
));
|
||||
}
|
||||
if content_type.is_empty() || content_type == "application/octet-stream" {
|
||||
content_type = if is_image(filename, "") {
|
||||
guess_mime_type(filename).to_string()
|
||||
} else {
|
||||
"application/octet-stream".into()
|
||||
};
|
||||
}
|
||||
|
||||
let cos_host = if COS_USE_ACCELERATE {
|
||||
format!("{}.cos.accelerate.myqcloud.com", creds.bucket)
|
||||
} else {
|
||||
format!("{}.cos.{}.myqcloud.com", creds.bucket, creds.region)
|
||||
};
|
||||
let encoded_key = encode_cos_key(&creds.location);
|
||||
let cos_url = format!("https://{cos_host}/{}", encoded_key.trim_start_matches('/'));
|
||||
|
||||
let now = unix_now();
|
||||
let start = if creds.start_time != 0 {
|
||||
creds.start_time
|
||||
} else {
|
||||
now
|
||||
};
|
||||
let expire = if creds.expired_time > now {
|
||||
creds.expired_time - now
|
||||
} else {
|
||||
3600
|
||||
};
|
||||
|
||||
let headers_for_sign: Vec<(&str, &str)> = vec![
|
||||
("host", cos_host.as_str()),
|
||||
("content-type", content_type.as_str()),
|
||||
("x-cos-security-token", creds.session_token.as_str()),
|
||||
];
|
||||
let path = format!("/{}", encoded_key.trim_start_matches('/'));
|
||||
let sig = cos_sign(&CosSignInput {
|
||||
method: "put",
|
||||
path: &path,
|
||||
params: &[],
|
||||
headers: &headers_for_sign,
|
||||
secret_id: &creds.secret_id,
|
||||
secret_key: &creds.secret_key,
|
||||
start_time: start,
|
||||
expire_seconds: expire,
|
||||
});
|
||||
|
||||
info!(
|
||||
"[yuanbao] COS PUT bucket={} key={} size={}",
|
||||
creds.bucket,
|
||||
creds.location,
|
||||
data.len()
|
||||
);
|
||||
let resp = http
|
||||
.put(&cos_url)
|
||||
.header("Authorization", sig)
|
||||
.header("Content-Type", content_type.as_str())
|
||||
.header("x-cos-security-token", &creds.session_token)
|
||||
.body(data.to_vec())
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| YuanbaoError::Connection(format!("COS PUT: {e}")))?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(YuanbaoError::Media(format!(
|
||||
"COS PUT HTTP {}",
|
||||
resp.status()
|
||||
)));
|
||||
}
|
||||
|
||||
let dims = if content_type.starts_with("image/") {
|
||||
parse_image_size(data).unwrap_or(ImageDims {
|
||||
width: 0,
|
||||
height: 0,
|
||||
})
|
||||
} else {
|
||||
ImageDims {
|
||||
width: 0,
|
||||
height: 0,
|
||||
}
|
||||
};
|
||||
|
||||
let uuid = {
|
||||
let mut h = Sha1::new();
|
||||
h.update(data);
|
||||
hex::encode(h.finalize())
|
||||
};
|
||||
let url = if creds.resource_url.is_empty() {
|
||||
cos_url
|
||||
} else {
|
||||
creds.resource_url.clone()
|
||||
};
|
||||
debug!("[yuanbao] COS upload ok url={url}");
|
||||
Ok(UploadResult {
|
||||
url,
|
||||
uuid,
|
||||
size: data.len() as u64,
|
||||
width: dims.width,
|
||||
height: dims.height,
|
||||
})
|
||||
}
|
||||
|
||||
fn unix_now() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn cos_sign_is_deterministic() {
|
||||
let s = cos_sign(&CosSignInput {
|
||||
method: "put",
|
||||
path: "/test/file.bin",
|
||||
params: &[],
|
||||
headers: &[("host", "bucket.cos.example.com")],
|
||||
secret_id: "AKID",
|
||||
secret_key: "SK",
|
||||
start_time: 1_700_000_000,
|
||||
expire_seconds: 3600,
|
||||
});
|
||||
let s2 = cos_sign(&CosSignInput {
|
||||
method: "put",
|
||||
path: "/test/file.bin",
|
||||
params: &[],
|
||||
headers: &[("host", "bucket.cos.example.com")],
|
||||
secret_id: "AKID",
|
||||
secret_key: "SK",
|
||||
start_time: 1_700_000_000,
|
||||
expire_seconds: 3600,
|
||||
});
|
||||
assert_eq!(s, s2);
|
||||
assert!(s.starts_with("q-sign-algorithm=sha1"));
|
||||
assert!(s.contains("q-ak=AKID"));
|
||||
assert!(s.contains("q-sign-time=1700000000;1700003600"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cos_sign_changes_with_path() {
|
||||
let base = CosSignInput {
|
||||
method: "put",
|
||||
path: "/a",
|
||||
params: &[],
|
||||
headers: &[("host", "h")],
|
||||
secret_id: "AKID",
|
||||
secret_key: "SK",
|
||||
start_time: 1_700_000_000,
|
||||
expire_seconds: 3600,
|
||||
};
|
||||
let s1 = cos_sign(&base);
|
||||
let s2 = cos_sign(&CosSignInput { path: "/b", ..base });
|
||||
assert_ne!(s1, s2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cos_sign_lowercases_method_and_includes_url_params() {
|
||||
let s = cos_sign(&CosSignInput {
|
||||
method: "PUT", // mixed case → should be lowercased into sig
|
||||
path: "/k",
|
||||
params: &[("Foo", "Bar Baz")], // url-encoded value
|
||||
headers: &[("Host", "h")],
|
||||
secret_id: "AKID",
|
||||
secret_key: "SK",
|
||||
start_time: 1_700_000_000,
|
||||
expire_seconds: 600,
|
||||
});
|
||||
assert!(s.contains("q-url-param-list=foo"));
|
||||
// header list also lowercased
|
||||
assert!(s.contains("q-header-list=host"));
|
||||
}
|
||||
|
||||
fn ok_credentials_body(bucket: &str, location: &str) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"code": 0,
|
||||
"data": {
|
||||
"bucketName": bucket,
|
||||
"region": "ap-shanghai",
|
||||
"location": location,
|
||||
"encryptTmpSecretId": "AKID",
|
||||
"encryptTmpSecretKey": "SECRET",
|
||||
"encryptToken": "session-token",
|
||||
"startTime": 1_700_000_000u64,
|
||||
"expiredTime": 1_700_003_600u64,
|
||||
"resourceUrl": "https://cdn.example/r",
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_cos_credentials_parses_data_block() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::path(UPLOAD_INFO_PATH))
|
||||
.and(wiremock::matchers::header("X-Token", "tok"))
|
||||
.and(wiremock::matchers::header("X-Source", "web"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200)
|
||||
.set_body_json(ok_credentials_body("bkt-1", "k/v/file.png")),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let http = reqwest::Client::new();
|
||||
let creds = get_cos_credentials(&http, &server.uri(), "appk", "bot", "tok", "", "file.png")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(creds.bucket, "bkt-1");
|
||||
assert_eq!(creds.region, "ap-shanghai");
|
||||
assert_eq!(creds.location, "k/v/file.png");
|
||||
assert_eq!(creds.secret_id, "AKID");
|
||||
assert_eq!(creds.secret_key, "SECRET");
|
||||
assert_eq!(creds.session_token, "session-token");
|
||||
assert_eq!(creds.resource_url, "https://cdn.example/r");
|
||||
assert_eq!(creds.start_time, 1_700_000_000);
|
||||
assert_eq!(creds.expired_time, 1_700_003_600);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_cos_credentials_falls_back_to_app_key_for_xid_when_bot_id_empty() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::path(UPLOAD_INFO_PATH))
|
||||
.and(wiremock::matchers::header("X-ID", "appk"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200)
|
||||
.set_body_json(ok_credentials_body("bkt", "loc")),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let http = reqwest::Client::new();
|
||||
let creds = get_cos_credentials(&http, &server.uri(), "appk", "", "tok", "", "f")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(creds.bucket, "bkt");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_cos_credentials_sends_route_env_header_when_non_empty() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
// Bind the matcher to both the upload-info path AND the header so
|
||||
// this test fails if a future refactor routes the call elsewhere
|
||||
// but happens to still attach `X-Route-Env: canary` somewhere.
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::path(UPLOAD_INFO_PATH))
|
||||
.and(wiremock::matchers::header("X-Route-Env", "canary"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200)
|
||||
.set_body_json(ok_credentials_body("bkt", "loc")),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let http = reqwest::Client::new();
|
||||
get_cos_credentials(&http, &server.uri(), "appk", "bot", "tok", "canary", "f")
|
||||
.await
|
||||
.expect("should send canary header");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_cos_credentials_surfaces_http_error() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.respond_with(wiremock::ResponseTemplate::new(500).set_body_string("boom"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let http = reqwest::Client::new();
|
||||
let err = get_cos_credentials(&http, &server.uri(), "appk", "bot", "tok", "", "f")
|
||||
.await
|
||||
.unwrap_err();
|
||||
match err {
|
||||
YuanbaoError::Media(m) => assert!(m.contains("HTTP 500"), "got {m}"),
|
||||
other => panic!("expected Media error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_cos_credentials_surfaces_non_zero_business_code() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"code": 4001,
|
||||
"msg": "quota",
|
||||
})),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let http = reqwest::Client::new();
|
||||
let err = get_cos_credentials(&http, &server.uri(), "appk", "bot", "tok", "", "f")
|
||||
.await
|
||||
.unwrap_err();
|
||||
match err {
|
||||
YuanbaoError::Media(m) => {
|
||||
assert!(m.contains("code=4001"), "got {m}");
|
||||
assert!(m.contains("quota"), "got {m}");
|
||||
}
|
||||
other => panic!("expected Media error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upload_to_cos_rejects_missing_credentials() {
|
||||
let http = reqwest::Client::new();
|
||||
// empty credentials → fail without making any HTTP call
|
||||
let bad = CosCredentials::default();
|
||||
let err = upload_to_cos(
|
||||
&http,
|
||||
&bad,
|
||||
b"data",
|
||||
"f.bin",
|
||||
"application/octet-stream".into(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
match err {
|
||||
YuanbaoError::Media(m) => assert!(m.contains("credentials missing"), "got {m}"),
|
||||
other => panic!("expected Media error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
// NOTE: upload_to_cos always targets `<bucket>.cos.accelerate.myqcloud.com`
|
||||
// which we cannot redirect at the reqwest layer without DNS hacks, so we
|
||||
// only cover the guard branch (missing creds) above. The PUT body itself
|
||||
// is exercised by integration tests, not unit tests.
|
||||
|
||||
#[test]
|
||||
fn encode_cos_key_keeps_slashes_but_escapes_segments() {
|
||||
assert_eq!(encode_cos_key("plain/file.png"), "plain/file.png");
|
||||
assert_eq!(encode_cos_key("a b/c d.png"), "a%20b/c%20d.png");
|
||||
}
|
||||
}
|
||||
@@ -1,61 +0,0 @@
|
||||
//! Yuanbao channel error types.
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
/// Close codes from the yuanbao gateway that indicate the connection
|
||||
/// must **not** be retried (auth failure, kicked off, etc.).
|
||||
///
|
||||
/// Mirrors `NO_RECONNECT_CLOSE_CODES` in hermes-agent `yuanbao.py`.
|
||||
pub const NO_RECONNECT_CLOSE_CODES: &[u16] = &[4012, 4013, 4014, 4018, 4019, 4021];
|
||||
|
||||
/// Auth-related response codes that mean "credentials are bad" — surface
|
||||
/// to the user, don't auto-retry.
|
||||
pub const AUTH_FAILED_CODES: &[u32] = &[40001, 40002, 40003];
|
||||
|
||||
/// Auth-related codes that are transient — retry with backoff.
|
||||
pub const AUTH_RETRYABLE_CODES: &[u32] = &[40010, 40011];
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum YuanbaoError {
|
||||
#[error("protocol encode error: {0}")]
|
||||
ProtoEncode(String),
|
||||
|
||||
#[error("protocol decode error: {0}")]
|
||||
ProtoDecode(String),
|
||||
|
||||
#[error("not connected")]
|
||||
NotConnected,
|
||||
|
||||
#[error("connection closed: code={code}, reason={reason}")]
|
||||
ConnectionClosed { code: u16, reason: String },
|
||||
|
||||
#[error("WebSocket error: {0}")]
|
||||
WebSocket(String),
|
||||
|
||||
#[error("HTTP/connection error: {0}")]
|
||||
Connection(String),
|
||||
|
||||
#[error("auth-bind failed: {0}")]
|
||||
AuthFailed(String),
|
||||
|
||||
#[error("auth-bind timeout")]
|
||||
AuthTimeout,
|
||||
|
||||
#[error("login timeout")]
|
||||
LoginTimeout,
|
||||
|
||||
#[error("request timeout: {0}")]
|
||||
Timeout(String),
|
||||
|
||||
#[error("send-message failed: {0}")]
|
||||
SendFailed(String),
|
||||
|
||||
#[error("media error: {0}")]
|
||||
Media(String),
|
||||
|
||||
#[error("invalid message: {0}")]
|
||||
InvalidMessage(String),
|
||||
|
||||
#[error("config error: {0}")]
|
||||
Config(String),
|
||||
}
|
||||
@@ -1,115 +0,0 @@
|
||||
//! Account-id shortening for yuanbao.
|
||||
//!
|
||||
//! Yuanbao uids (`from_account`) are 64-char hashes assigned by the platform.
|
||||
//! The composite `ChannelMessage` thread_id that downstream consumers derive
|
||||
//! from `sender` and `reply_target` (`channel:yuanbao_<sender>_<reply_target>`)
|
||||
//! becomes ~145 chars. After the conversation store hex-encodes that for the
|
||||
//! per-thread JSONL filename it grows to ~296 chars, exceeding `NAME_MAX`
|
||||
//! (255 bytes) on ext4/HFS+/APFS/NTFS — writes fail with `ENAMETOOLONG` and
|
||||
//! channel history is lost.
|
||||
//!
|
||||
//! Rather than push the filesystem limit into shared `ConversationStore` code,
|
||||
//! we shorten yuanbao-specific ids at the channel boundary. Internal yuanbao
|
||||
//! state (echo guard, access control, owner-command check) keeps the original
|
||||
//! `from_account` — only the value emitted on `ChannelMessage.sender` /
|
||||
//! `ChannelMessage.reply_target` is shortened.
|
||||
//!
|
||||
//! Format: `<first 8 chars of uid>_<first 16 hex chars of sha256(uid)>`.
|
||||
//! The 8-char prefix keeps logs roughly groupable for the same user; the
|
||||
//! sha256 suffix guarantees uniqueness across uids that share a prefix.
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
/// Max raw account-id length before the shortening kicks in.
|
||||
///
|
||||
/// Anything shorter is passed through unchanged so short upstream-style ids
|
||||
/// (e.g. numeric ids, future protocol changes) keep their natural form.
|
||||
const ACCOUNT_ID_PASSTHROUGH_MAX: usize = 24;
|
||||
|
||||
/// Shorten a yuanbao account id for use in `ChannelMessage.sender` /
|
||||
/// `ChannelMessage.reply_target`. See module docs for rationale.
|
||||
pub(super) fn shorten_account_id(uid: &str) -> String {
|
||||
if uid.len() <= ACCOUNT_ID_PASSTHROUGH_MAX {
|
||||
return uid.to_string();
|
||||
}
|
||||
let prefix: String = uid.chars().take(8).collect();
|
||||
let digest = Sha256::digest(uid.as_bytes());
|
||||
format!("{prefix}_{:.16x}", digest)
|
||||
}
|
||||
|
||||
/// Shorten a yuanbao `reply_target`, preserving the `g:<group_code>` shape
|
||||
/// used for group chats. The `g:` discriminator is required by outbound
|
||||
/// routing (see [`super::types::InboundContext::reply_target`]).
|
||||
pub(super) fn shorten_reply_target(reply_target: &str) -> String {
|
||||
if let Some(group_code) = reply_target.strip_prefix("g:") {
|
||||
format!("g:{}", shorten_account_id(group_code))
|
||||
} else {
|
||||
shorten_account_id(reply_target)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn passes_short_ids_through_unchanged() {
|
||||
assert_eq!(shorten_account_id("123456"), "123456");
|
||||
assert_eq!(shorten_account_id(""), "");
|
||||
let exactly_max = "a".repeat(ACCOUNT_ID_PASSTHROUGH_MAX);
|
||||
assert_eq!(shorten_account_id(&exactly_max), exactly_max);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shortens_long_ids_to_prefix_plus_hash() {
|
||||
let long_uid = "a".repeat(64);
|
||||
let shortened = shorten_account_id(&long_uid);
|
||||
assert_eq!(shortened.len(), 8 + 1 + 16, "8 prefix + '_' + 16 hex");
|
||||
assert!(shortened.starts_with("aaaaaaaa_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shortening_is_deterministic_and_collision_resistant() {
|
||||
let a = "f".repeat(64);
|
||||
let mut b = a.clone();
|
||||
b.replace_range(63..64, "e"); // differ in last char only
|
||||
let sa = shorten_account_id(&a);
|
||||
let sb = shorten_account_id(&b);
|
||||
assert_eq!(sa, shorten_account_id(&a), "deterministic");
|
||||
assert_ne!(sa, sb, "different uids hash to different ids");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn group_reply_target_preserves_g_prefix() {
|
||||
let short_group = shorten_reply_target("g:short_group");
|
||||
assert_eq!(short_group, "g:short_group");
|
||||
|
||||
let long_code = "a".repeat(64);
|
||||
let long_group = format!("g:{long_code}");
|
||||
let shortened = shorten_reply_target(&long_group);
|
||||
assert!(shortened.starts_with("g:aaaaaaaa_"));
|
||||
assert_eq!(shortened.len(), 2 + 8 + 1 + 16);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dm_reply_target_shortens_like_account_id() {
|
||||
let uid = "z".repeat(64);
|
||||
assert_eq!(shorten_reply_target(&uid), shorten_account_id(&uid));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shortened_thread_id_fits_under_name_max() {
|
||||
// Simulate the worst case: long uid for sender + reply_target.
|
||||
let uid = "f".repeat(64);
|
||||
let sender = shorten_account_id(&uid);
|
||||
let reply_target = shorten_account_id(&uid);
|
||||
let thread_id = format!("channel:yuanbao_{sender}_{reply_target}");
|
||||
// hex-encoded filename used by ConversationStore (`<hex>.jsonl`).
|
||||
let hex_name_len = thread_id.len() * 2 + ".jsonl".len();
|
||||
// NAME_MAX on common filesystems is 255 bytes.
|
||||
assert!(
|
||||
hex_name_len <= 255,
|
||||
"shortened thread_id hex filename ({hex_name_len} bytes) must fit under NAME_MAX (255)"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,623 +0,0 @@
|
||||
//! Inbound message pipeline (17 stages).
|
||||
//!
|
||||
//! Mirrors `InboundPipeline` in hermes-agent `gateway/platforms/yuanbao.py`.
|
||||
//! Each stage runs in order; any of them can short-circuit by
|
||||
//! returning `Skip(reason)`. `Abort(_)` propagates an error.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::{debug, trace};
|
||||
|
||||
use super::config::YuanbaoConfig;
|
||||
use super::errors::YuanbaoError;
|
||||
use super::proto::{decode_inbound_json, decode_inbound_push};
|
||||
use super::proto_constants::*;
|
||||
use super::types::*;
|
||||
|
||||
/// Shared per-channel state that survives across messages.
|
||||
pub struct PipelineState {
|
||||
pub bot_id: String,
|
||||
pub bot_name: String,
|
||||
pub owner_id: String,
|
||||
pub dm_access: AccessPolicy,
|
||||
pub group_access: AccessPolicy,
|
||||
pub allowed_users: Vec<String>,
|
||||
pub allowed_groups: Vec<String>,
|
||||
pub group_at_required: bool,
|
||||
pub home_chat: RwLock<Option<String>>,
|
||||
pub dedup: RwLock<DedupCache>,
|
||||
}
|
||||
|
||||
impl PipelineState {
|
||||
pub fn new(cfg: &YuanbaoConfig, bot_id: String) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
bot_id,
|
||||
bot_name: cfg.bot_name.clone(),
|
||||
owner_id: cfg.owner_id.clone(),
|
||||
dm_access: AccessPolicy::parse(&cfg.dm_access),
|
||||
group_access: AccessPolicy::parse(&cfg.group_access),
|
||||
allowed_users: cfg.allowed_users.clone(),
|
||||
allowed_groups: cfg.allowed_groups.clone(),
|
||||
group_at_required: cfg.group_at_required,
|
||||
home_chat: RwLock::new(None),
|
||||
dedup: RwLock::new(DedupCache::new(DEDUP_CAPACITY, DEDUP_TTL_SECS)),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AccessPolicy {
|
||||
Open,
|
||||
Allowlist,
|
||||
Closed,
|
||||
}
|
||||
|
||||
impl AccessPolicy {
|
||||
fn parse(s: &str) -> Self {
|
||||
match s.to_ascii_lowercase().as_str() {
|
||||
"open" => Self::Open,
|
||||
"closed" | "disabled" | "none" => Self::Closed,
|
||||
_ => Self::Allowlist,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Mutable context passed through every inbound stage.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PipelineCtx {
|
||||
pub msg: InboundMessage,
|
||||
pub source: Source,
|
||||
pub text: String,
|
||||
pub image_urls: Vec<String>,
|
||||
pub is_at_bot: bool,
|
||||
pub is_owner_command: bool,
|
||||
pub kind: MessageKind,
|
||||
}
|
||||
|
||||
/// Outcome of a single inbound stage invocation.
|
||||
#[derive(Debug)]
|
||||
pub enum MwResult {
|
||||
Continue,
|
||||
Skip(&'static str),
|
||||
Abort(YuanbaoError),
|
||||
}
|
||||
|
||||
/// Final outcome of the whole pipeline.
|
||||
#[derive(Debug)]
|
||||
pub enum PipelineOutcome {
|
||||
Dispatch(PipelineCtx),
|
||||
Filtered(&'static str),
|
||||
Failed(YuanbaoError),
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait Middleware: Send + Sync {
|
||||
fn name(&self) -> &'static str;
|
||||
async fn process(&self, state: &PipelineState, ctx: &mut PipelineCtx) -> MwResult;
|
||||
}
|
||||
|
||||
/// LRU-like dedup cache with TTL.
|
||||
pub struct DedupCache {
|
||||
capacity: usize,
|
||||
ttl: Duration,
|
||||
order: VecDeque<(String, Instant)>,
|
||||
index: std::collections::HashSet<String>,
|
||||
}
|
||||
|
||||
impl DedupCache {
|
||||
pub fn new(capacity: usize, ttl_secs: u64) -> Self {
|
||||
Self {
|
||||
capacity,
|
||||
ttl: Duration::from_secs(ttl_secs),
|
||||
order: VecDeque::with_capacity(capacity),
|
||||
index: std::collections::HashSet::with_capacity(capacity),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns `true` if `id` has been seen within the TTL window. Inserts it otherwise.
|
||||
pub fn check_and_insert(&mut self, id: &str) -> bool {
|
||||
self.evict_expired();
|
||||
if self.index.contains(id) {
|
||||
return true;
|
||||
}
|
||||
if self.order.len() >= self.capacity {
|
||||
if let Some((old, _)) = self.order.pop_front() {
|
||||
self.index.remove(&old);
|
||||
}
|
||||
}
|
||||
self.order.push_back((id.to_string(), Instant::now()));
|
||||
self.index.insert(id.to_string());
|
||||
false
|
||||
}
|
||||
|
||||
fn evict_expired(&mut self) {
|
||||
let now = Instant::now();
|
||||
while let Some((_, ts)) = self.order.front() {
|
||||
if now.duration_since(*ts) > self.ttl {
|
||||
if let Some((old, _)) = self.order.pop_front() {
|
||||
self.index.remove(&old);
|
||||
}
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ───── Individual inbound stages ────────────────────────────────────
|
||||
|
||||
struct DecodeMw;
|
||||
struct ExtractFieldsMw;
|
||||
struct RecallGuardMw;
|
||||
struct DedupMw;
|
||||
struct SkipSelfMw;
|
||||
struct ChatRoutingMw;
|
||||
struct AccessGuardMw;
|
||||
struct AutoSetHomeMw;
|
||||
struct ExtractContentMw;
|
||||
struct PlaceholderFilterMw;
|
||||
struct OwnerCommandMw;
|
||||
struct BuildSourceMw;
|
||||
struct GroupAtGuardMw;
|
||||
struct GroupAttributionMw;
|
||||
struct ClassifyMsgTypeMw;
|
||||
struct QuoteContextMw;
|
||||
struct MediaResolveMw;
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for DecodeMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"decode"
|
||||
}
|
||||
async fn process(&self, _s: &PipelineState, _c: &mut PipelineCtx) -> MwResult {
|
||||
// Decoding happens before we build a PipelineCtx — this MW is a placeholder
|
||||
// so the stage list still has 17 entries (mirrors hermes-agent).
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for ExtractFieldsMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"extract_fields"
|
||||
}
|
||||
async fn process(&self, _s: &PipelineState, _c: &mut PipelineCtx) -> MwResult {
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for RecallGuardMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"recall_guard"
|
||||
}
|
||||
async fn process(&self, _s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
if c.msg.is_recall() {
|
||||
c.kind = MessageKind::Recall;
|
||||
return MwResult::Skip("recall_guard");
|
||||
}
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for DedupMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"dedup"
|
||||
}
|
||||
async fn process(&self, s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
if c.msg.msg_id.is_empty() {
|
||||
return MwResult::Continue; // nothing to dedup on
|
||||
}
|
||||
let mut cache = s.dedup.write().await;
|
||||
if cache.check_and_insert(&c.msg.msg_id) {
|
||||
return MwResult::Skip("dedup");
|
||||
}
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for SkipSelfMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"skip_self"
|
||||
}
|
||||
async fn process(&self, s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
if !s.bot_id.is_empty() && c.msg.from_account == s.bot_id {
|
||||
return MwResult::Skip("skip_self");
|
||||
}
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for ChatRoutingMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"chat_routing"
|
||||
}
|
||||
async fn process(&self, _s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
c.source.is_group = c.msg.is_group();
|
||||
c.source.group_code = c.msg.group_code.clone();
|
||||
c.source.from_account = c.msg.from_account.clone();
|
||||
c.source.sender_nickname = c.msg.sender_nickname.clone();
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for AccessGuardMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"access_guard"
|
||||
}
|
||||
async fn process(&self, s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
let (policy, allow_list, key) = if c.source.is_group {
|
||||
(s.group_access, &s.allowed_groups, &c.source.group_code)
|
||||
} else {
|
||||
(s.dm_access, &s.allowed_users, &c.source.from_account)
|
||||
};
|
||||
let pass = match policy {
|
||||
AccessPolicy::Open => true,
|
||||
AccessPolicy::Closed => false,
|
||||
AccessPolicy::Allowlist => allow_list.iter().any(|u| u == "*" || u == key),
|
||||
};
|
||||
if !pass {
|
||||
return MwResult::Skip("access_guard");
|
||||
}
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for AutoSetHomeMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"auto_set_home"
|
||||
}
|
||||
async fn process(&self, s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
if !c.source.is_group {
|
||||
let mut home = s.home_chat.write().await;
|
||||
if home.is_none() {
|
||||
*home = Some(c.source.reply_target());
|
||||
}
|
||||
}
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for ExtractContentMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"extract_content"
|
||||
}
|
||||
async fn process(&self, _s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
c.text = c.msg.extract_text();
|
||||
c.image_urls = c.msg.extract_image_urls();
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for PlaceholderFilterMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"placeholder_filter"
|
||||
}
|
||||
async fn process(&self, _s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
let trimmed = c.text.trim();
|
||||
let is_placeholder = trimmed == "[image]" || trimmed == "[file]" || trimmed == "[图片]";
|
||||
if (trimmed.is_empty() || is_placeholder) && c.image_urls.is_empty() {
|
||||
return MwResult::Skip("placeholder_filter");
|
||||
}
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for OwnerCommandMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"owner_command"
|
||||
}
|
||||
async fn process(&self, s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
if !s.owner_id.is_empty()
|
||||
&& c.msg.from_account == s.owner_id
|
||||
&& c.text.trim_start().starts_with('/')
|
||||
{
|
||||
c.is_owner_command = true;
|
||||
}
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for BuildSourceMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"build_source"
|
||||
}
|
||||
async fn process(&self, _s: &PipelineState, _c: &mut PipelineCtx) -> MwResult {
|
||||
// Source already populated by ChatRoutingMw.
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for GroupAtGuardMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"group_at_guard"
|
||||
}
|
||||
async fn process(&self, s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
if !c.source.is_group || !s.group_at_required {
|
||||
return MwResult::Continue;
|
||||
}
|
||||
let by_name = !s.bot_name.is_empty() && c.text.contains(&format!("@{}", s.bot_name));
|
||||
let by_id = !s.bot_id.is_empty() && c.text.contains(&format!("@{}", s.bot_id));
|
||||
let by_mention =
|
||||
!s.bot_id.is_empty() && c.text.contains(&format!("[at|userId:{}]", s.bot_id));
|
||||
c.is_at_bot = by_name || by_id || by_mention;
|
||||
if !c.is_at_bot && !c.is_owner_command {
|
||||
return MwResult::Skip("group_at_guard");
|
||||
}
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for GroupAttributionMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"group_attribution"
|
||||
}
|
||||
async fn process(&self, s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
// Strip `@bot` from text and the TIM `[at|userId:…]` markup.
|
||||
if c.source.is_group && c.is_at_bot {
|
||||
if !s.bot_name.is_empty() {
|
||||
c.text = c.text.replace(&format!("@{}", s.bot_name), "");
|
||||
}
|
||||
if !s.bot_id.is_empty() {
|
||||
c.text = c.text.replace(&format!("@{}", s.bot_id), "");
|
||||
c.text = c.text.replace(&format!("[at|userId:{}]", s.bot_id), "");
|
||||
}
|
||||
c.text = c.text.trim().to_string();
|
||||
}
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for ClassifyMsgTypeMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"classify_msg_type"
|
||||
}
|
||||
async fn process(&self, _s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
let has_text = !c.text.is_empty();
|
||||
let has_image = !c.image_urls.is_empty();
|
||||
let has_file = c.msg.msg_body.iter().any(|el| el.msg_type == tim::FILE);
|
||||
let has_sound = c.msg.msg_body.iter().any(|el| el.msg_type == tim::SOUND);
|
||||
c.kind = match (has_text, has_image, has_file, has_sound) {
|
||||
(_, true, _, _) if has_text => MessageKind::Mixed,
|
||||
(_, true, _, _) => MessageKind::Image,
|
||||
(_, _, true, _) => MessageKind::File,
|
||||
(_, _, _, true) => MessageKind::Voice,
|
||||
_ => MessageKind::Text,
|
||||
};
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for QuoteContextMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"quote_context"
|
||||
}
|
||||
async fn process(&self, _s: &PipelineState, c: &mut PipelineCtx) -> MwResult {
|
||||
// The cloud_custom_data field carries a JSON quote envelope; for
|
||||
// now we just leave the raw payload accessible via `msg.cloud_custom_data`
|
||||
// for downstream tools. Full parsing is intentionally deferred —
|
||||
// hermes-agent does it lazily too.
|
||||
let _ = c;
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Middleware for MediaResolveMw {
|
||||
fn name(&self) -> &'static str {
|
||||
"media_resolve"
|
||||
}
|
||||
async fn process(&self, _s: &PipelineState, _c: &mut PipelineCtx) -> MwResult {
|
||||
// ybres:// resource URLs would be resolved here. Currently URLs
|
||||
// arrive pre-resolved from the server; expand later if needed.
|
||||
MwResult::Continue
|
||||
}
|
||||
}
|
||||
|
||||
/// Composite pipeline = ordered Vec of inbound stages.
|
||||
pub struct InboundPipeline {
|
||||
state: Arc<PipelineState>,
|
||||
stages: Vec<Box<dyn Middleware>>,
|
||||
}
|
||||
|
||||
impl InboundPipeline {
|
||||
pub fn new(state: Arc<PipelineState>) -> Self {
|
||||
let stages: Vec<Box<dyn Middleware>> = vec![
|
||||
Box::new(DecodeMw),
|
||||
Box::new(ExtractFieldsMw),
|
||||
Box::new(RecallGuardMw),
|
||||
Box::new(DedupMw),
|
||||
Box::new(SkipSelfMw),
|
||||
Box::new(ChatRoutingMw),
|
||||
Box::new(AccessGuardMw),
|
||||
Box::new(AutoSetHomeMw),
|
||||
Box::new(ExtractContentMw),
|
||||
Box::new(PlaceholderFilterMw),
|
||||
Box::new(OwnerCommandMw),
|
||||
Box::new(BuildSourceMw),
|
||||
Box::new(GroupAtGuardMw),
|
||||
Box::new(GroupAttributionMw),
|
||||
Box::new(ClassifyMsgTypeMw),
|
||||
Box::new(QuoteContextMw),
|
||||
Box::new(MediaResolveMw),
|
||||
];
|
||||
Self { state, stages }
|
||||
}
|
||||
|
||||
/// Decode a biz push body, run it through every stage, return the outcome.
|
||||
///
|
||||
/// The yuanbao gateway may push the biz body as either protobuf
|
||||
/// (`InboundMessagePush`) or a JSON string with the same field shape
|
||||
/// (snake_case + `log_ext.trace_id`). We sniff the first non-whitespace
|
||||
/// byte to pick the decoder — `{` means JSON, anything else is treated
|
||||
/// as protobuf. Mirrors plugin gateway.ts::wsPushToInboundMessage
|
||||
/// (l. 288), which tries protobuf first and falls back to JSON.
|
||||
pub async fn process(&self, biz_body: &[u8]) -> PipelineOutcome {
|
||||
let is_json = biz_body
|
||||
.iter()
|
||||
.find(|b| !b.is_ascii_whitespace())
|
||||
.map(|b| *b == b'{')
|
||||
.unwrap_or(false);
|
||||
|
||||
let msg = if is_json {
|
||||
match decode_inbound_json(biz_body) {
|
||||
Ok(m) => m,
|
||||
Err(e) => return PipelineOutcome::Failed(e),
|
||||
}
|
||||
} else {
|
||||
match decode_inbound_push(biz_body) {
|
||||
Ok(m) => m,
|
||||
Err(e) => return PipelineOutcome::Failed(e),
|
||||
}
|
||||
};
|
||||
let mut ctx = PipelineCtx {
|
||||
msg,
|
||||
source: Source::default(),
|
||||
text: String::new(),
|
||||
image_urls: Vec::new(),
|
||||
is_at_bot: false,
|
||||
is_owner_command: false,
|
||||
kind: MessageKind::Text,
|
||||
};
|
||||
for stage in &self.stages {
|
||||
match stage.process(&self.state, &mut ctx).await {
|
||||
MwResult::Continue => {
|
||||
trace!("[yuanbao:inbound] {} pass", stage.name());
|
||||
}
|
||||
MwResult::Skip(reason) => {
|
||||
debug!("[yuanbao:inbound] {} filtered ({})", stage.name(), reason);
|
||||
return PipelineOutcome::Filtered(reason);
|
||||
}
|
||||
MwResult::Abort(err) => return PipelineOutcome::Failed(err),
|
||||
}
|
||||
}
|
||||
PipelineOutcome::Dispatch(ctx)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn cfg(bot_id: &str) -> YuanbaoConfig {
|
||||
let mut c = YuanbaoConfig::default();
|
||||
c.app_key = "ak".into();
|
||||
c.ws_domain = "wss://x".into();
|
||||
c.token = "tok".into();
|
||||
c.bot_id = bot_id.into();
|
||||
c.bot_name = "bot".into();
|
||||
c.dm_access = "open".into();
|
||||
c.group_access = "open".into();
|
||||
c
|
||||
}
|
||||
|
||||
fn ctx_with(msg: InboundMessage) -> PipelineCtx {
|
||||
PipelineCtx {
|
||||
msg,
|
||||
source: Source::default(),
|
||||
text: String::new(),
|
||||
image_urls: Vec::new(),
|
||||
is_at_bot: false,
|
||||
is_owner_command: false,
|
||||
kind: MessageKind::Text,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dedup_skips_repeat() {
|
||||
let state = PipelineState::new(&cfg("bot1"), "bot1".into());
|
||||
let mw = DedupMw;
|
||||
let msg = InboundMessage {
|
||||
msg_id: "m1".into(),
|
||||
..Default::default()
|
||||
};
|
||||
let mut c1 = ctx_with(msg.clone());
|
||||
assert!(matches!(
|
||||
mw.process(&state, &mut c1).await,
|
||||
MwResult::Continue
|
||||
));
|
||||
let mut c2 = ctx_with(msg);
|
||||
assert!(matches!(
|
||||
mw.process(&state, &mut c2).await,
|
||||
MwResult::Skip(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn access_guard_open() {
|
||||
let state = PipelineState::new(&cfg("bot1"), "bot1".into());
|
||||
let mw = AccessGuardMw;
|
||||
let mut c = ctx_with(InboundMessage {
|
||||
from_account: "alice".into(),
|
||||
..Default::default()
|
||||
});
|
||||
c.source.is_group = false;
|
||||
c.source.from_account = "alice".into();
|
||||
assert!(matches!(
|
||||
mw.process(&state, &mut c).await,
|
||||
MwResult::Continue
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn full_dm_dispatch() {
|
||||
let mut config = cfg("bot1");
|
||||
config.group_at_required = false;
|
||||
let state = PipelineState::new(&config, "bot1".into());
|
||||
let pipeline = InboundPipeline::new(state);
|
||||
let msg = InboundMessage {
|
||||
from_account: "alice".into(),
|
||||
to_account: "bot1".into(),
|
||||
msg_id: "hi".into(),
|
||||
msg_body: vec![MsgBodyElement {
|
||||
msg_type: "TIMTextElem".into(),
|
||||
msg_content: MsgContent {
|
||||
text: Some("hello".into()),
|
||||
..Default::default()
|
||||
},
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
let body = crate::openhuman::channels::providers::yuanbao::proto::encode_msg_body_element(
|
||||
&msg.msg_body[0],
|
||||
);
|
||||
// Synthesize an InboundMessagePush from scratch:
|
||||
use crate::openhuman::channels::providers::yuanbao::proto;
|
||||
let mut buf = Vec::new();
|
||||
let mut put_str = |fnum: u32, s: &str, b: &mut Vec<u8>| {
|
||||
proto::encode_varint(((fnum as u64) << 3) | 2, b);
|
||||
proto::encode_varint(s.len() as u64, b);
|
||||
b.extend_from_slice(s.as_bytes());
|
||||
};
|
||||
put_str(2, &msg.from_account, &mut buf);
|
||||
put_str(3, &msg.to_account, &mut buf);
|
||||
put_str(12, &msg.msg_id, &mut buf);
|
||||
proto::encode_varint(((13u64) << 3) | 2, &mut buf);
|
||||
proto::encode_varint(body.len() as u64, &mut buf);
|
||||
buf.extend_from_slice(&body);
|
||||
|
||||
let outcome = pipeline.process(&buf).await;
|
||||
assert!(
|
||||
matches!(outcome, PipelineOutcome::Dispatch(_)),
|
||||
"got {:?}",
|
||||
outcome
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,619 +0,0 @@
|
||||
//! Media helpers — MIME mapping, byte-level image dimension parsing,
|
||||
//! download with size cap, and TIM `msg_body` builders.
|
||||
//!
|
||||
//! Tencent COS upload lives in [`super::cos`] to keep both files under
|
||||
//! the 500-line ceiling.
|
||||
|
||||
use super::errors::YuanbaoError;
|
||||
use super::types::MsgBodyElement;
|
||||
|
||||
// ─── MIME / image-format mapping ───────────────────────────────────
|
||||
|
||||
pub fn guess_mime_type(filename: &str) -> &'static str {
|
||||
let ext = filename
|
||||
.rsplit_once('.')
|
||||
.map(|(_, e)| e.to_ascii_lowercase())
|
||||
.unwrap_or_default();
|
||||
match ext.as_str() {
|
||||
"jpg" | "jpeg" => "image/jpeg",
|
||||
"png" => "image/png",
|
||||
"gif" => "image/gif",
|
||||
"webp" => "image/webp",
|
||||
"bmp" => "image/bmp",
|
||||
"heic" => "image/heic",
|
||||
"tiff" => "image/tiff",
|
||||
"ico" => "image/x-icon",
|
||||
"pdf" => "application/pdf",
|
||||
"doc" => "application/msword",
|
||||
"docx" => "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
"xls" => "application/vnd.ms-excel",
|
||||
"xlsx" => "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
"ppt" => "application/vnd.ms-powerpoint",
|
||||
"pptx" => "application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
"txt" => "text/plain",
|
||||
"zip" => "application/zip",
|
||||
"tar" => "application/x-tar",
|
||||
"gz" => "application/gzip",
|
||||
"mp3" => "audio/mpeg",
|
||||
"mp4" => "video/mp4",
|
||||
"wav" => "audio/wav",
|
||||
"ogg" => "audio/ogg",
|
||||
"webm" => "video/webm",
|
||||
_ => "application/octet-stream",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_image(filename: &str, mime_type: &str) -> bool {
|
||||
if mime_type.starts_with("image/") {
|
||||
return true;
|
||||
}
|
||||
guess_mime_type(filename).starts_with("image/")
|
||||
}
|
||||
|
||||
/// Map a MIME type to the TIM `image_format` enum.
|
||||
pub fn image_format_code(mime: &str) -> u32 {
|
||||
match mime {
|
||||
"image/jpeg" | "image/jpg" => 1,
|
||||
"image/gif" => 2,
|
||||
"image/png" => 3,
|
||||
"image/bmp" => 4,
|
||||
"image/webp" | "image/heic" | "image/tiff" => 255,
|
||||
_ => 255,
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Pure-bytes image dimension parsing ─────────────────────────────
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct ImageDims {
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
}
|
||||
|
||||
pub fn parse_image_size(data: &[u8]) -> Option<ImageDims> {
|
||||
parse_png(data)
|
||||
.or_else(|| parse_jpeg(data))
|
||||
.or_else(|| parse_gif(data))
|
||||
.or_else(|| parse_webp(data))
|
||||
}
|
||||
|
||||
fn parse_png(buf: &[u8]) -> Option<ImageDims> {
|
||||
if buf.len() < 24 || &buf[..4] != b"\x89PNG" {
|
||||
return None;
|
||||
}
|
||||
let w = u32::from_be_bytes(buf[16..20].try_into().ok()?);
|
||||
let h = u32::from_be_bytes(buf[20..24].try_into().ok()?);
|
||||
Some(ImageDims {
|
||||
width: w,
|
||||
height: h,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_jpeg(buf: &[u8]) -> Option<ImageDims> {
|
||||
if buf.len() < 4 || buf[0] != 0xFF || buf[1] != 0xD8 {
|
||||
return None;
|
||||
}
|
||||
let mut i = 2usize;
|
||||
while i + 9 < buf.len() {
|
||||
if buf[i] != 0xFF {
|
||||
i += 1;
|
||||
continue;
|
||||
}
|
||||
let marker = buf[i + 1];
|
||||
if marker == 0xC0 || marker == 0xC2 {
|
||||
let h = u16::from_be_bytes(buf[i + 5..i + 7].try_into().ok()?);
|
||||
let w = u16::from_be_bytes(buf[i + 7..i + 9].try_into().ok()?);
|
||||
return Some(ImageDims {
|
||||
width: w as u32,
|
||||
height: h as u32,
|
||||
});
|
||||
}
|
||||
if i + 3 >= buf.len() {
|
||||
break;
|
||||
}
|
||||
let seg_len = u16::from_be_bytes(buf[i + 2..i + 4].try_into().ok()?) as usize;
|
||||
i += 2 + seg_len;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn parse_gif(buf: &[u8]) -> Option<ImageDims> {
|
||||
if buf.len() < 10 {
|
||||
return None;
|
||||
}
|
||||
let sig = &buf[..6];
|
||||
if sig != b"GIF87a" && sig != b"GIF89a" {
|
||||
return None;
|
||||
}
|
||||
let w = u16::from_le_bytes(buf[6..8].try_into().ok()?);
|
||||
let h = u16::from_le_bytes(buf[8..10].try_into().ok()?);
|
||||
Some(ImageDims {
|
||||
width: w as u32,
|
||||
height: h as u32,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_webp(buf: &[u8]) -> Option<ImageDims> {
|
||||
if buf.len() < 16 || &buf[..4] != b"RIFF" || &buf[8..12] != b"WEBP" {
|
||||
return None;
|
||||
}
|
||||
let chunk = &buf[12..16];
|
||||
if chunk == b"VP8 " {
|
||||
if buf.len() >= 30 && buf[23] == 0x9D && buf[24] == 0x01 && buf[25] == 0x2A {
|
||||
let w = u16::from_le_bytes(buf[26..28].try_into().ok()?) & 0x3FFF;
|
||||
let h = u16::from_le_bytes(buf[28..30].try_into().ok()?) & 0x3FFF;
|
||||
return Some(ImageDims {
|
||||
width: w as u32,
|
||||
height: h as u32,
|
||||
});
|
||||
}
|
||||
} else if chunk == b"VP8L" {
|
||||
if buf.len() >= 25 && buf[20] == 0x2F {
|
||||
let bits = u32::from_le_bytes(buf[21..25].try_into().ok()?);
|
||||
let w = (bits & 0x3FFF) + 1;
|
||||
let h = ((bits >> 14) & 0x3FFF) + 1;
|
||||
return Some(ImageDims {
|
||||
width: w,
|
||||
height: h,
|
||||
});
|
||||
}
|
||||
} else if chunk == b"VP8X" && buf.len() >= 30 {
|
||||
let w = (buf[24] as u32 | ((buf[25] as u32) << 8) | ((buf[26] as u32) << 16)) + 1;
|
||||
let h = (buf[27] as u32 | ((buf[28] as u32) << 8) | ((buf[29] as u32) << 16)) + 1;
|
||||
return Some(ImageDims {
|
||||
width: w,
|
||||
height: h,
|
||||
});
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
// ─── HTTP download with size cap ────────────────────────────────────
|
||||
|
||||
pub async fn download_url(
|
||||
http: &reqwest::Client,
|
||||
url: &str,
|
||||
max_size_mb: u64,
|
||||
) -> Result<(Vec<u8>, String), YuanbaoError> {
|
||||
let limit = max_size_mb.saturating_mul(1024 * 1024);
|
||||
|
||||
if let Ok(head) = http.head(url).send().await {
|
||||
if let Some(len) = head
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_LENGTH)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
{
|
||||
if len > limit {
|
||||
return Err(YuanbaoError::Media(format!(
|
||||
"remote file too large: {len} > limit {limit}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let resp = http
|
||||
.get(url)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| YuanbaoError::Connection(format!("download {url}: {e}")))?;
|
||||
if !resp.status().is_success() {
|
||||
return Err(YuanbaoError::Media(format!(
|
||||
"download HTTP {} for {url}",
|
||||
resp.status()
|
||||
)));
|
||||
}
|
||||
let ct = resp
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.unwrap_or("")
|
||||
.split(';')
|
||||
.next()
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.to_string();
|
||||
|
||||
let bytes = resp
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|e| YuanbaoError::Media(format!("read body: {e}")))?;
|
||||
if bytes.len() as u64 > limit {
|
||||
return Err(YuanbaoError::Media(format!(
|
||||
"downloaded file exceeds limit: {} > {}",
|
||||
bytes.len(),
|
||||
limit
|
||||
)));
|
||||
}
|
||||
Ok((bytes.to_vec(), ct))
|
||||
}
|
||||
|
||||
// ─── TIM msg_body builders ──────────────────────────────────────────
|
||||
|
||||
/// Build a TIM `TIMImageElem` `msg_body` ready to send.
|
||||
pub fn build_image_msg_body(
|
||||
url: &str,
|
||||
uuid: Option<&str>,
|
||||
filename: Option<&str>,
|
||||
size: u32,
|
||||
width: u32,
|
||||
height: u32,
|
||||
mime_type: &str,
|
||||
) -> Vec<MsgBodyElement> {
|
||||
use super::types::{ImageInfo, MsgContent};
|
||||
let uuid_str = uuid
|
||||
.map(|s| s.to_string())
|
||||
.or_else(|| filename.map(|s| s.to_string()))
|
||||
.unwrap_or_else(|| "image".to_string());
|
||||
let format = if mime_type.is_empty() {
|
||||
255
|
||||
} else {
|
||||
image_format_code(mime_type)
|
||||
};
|
||||
vec![MsgBodyElement {
|
||||
msg_type: "TIMImageElem".into(),
|
||||
msg_content: MsgContent {
|
||||
uuid: Some(uuid_str),
|
||||
image_format: Some(format),
|
||||
image_info_array: vec![ImageInfo {
|
||||
image_type: 1,
|
||||
size,
|
||||
width,
|
||||
height,
|
||||
url: url.to_string(),
|
||||
}],
|
||||
..Default::default()
|
||||
},
|
||||
}]
|
||||
}
|
||||
|
||||
/// Build a TIM `TIMFileElem` `msg_body` ready to send.
|
||||
pub fn build_file_msg_body(
|
||||
url: &str,
|
||||
filename: &str,
|
||||
uuid: Option<&str>,
|
||||
size: u32,
|
||||
) -> Vec<MsgBodyElement> {
|
||||
use super::types::MsgContent;
|
||||
let uuid_str = uuid
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| filename.to_string());
|
||||
vec![MsgBodyElement {
|
||||
msg_type: "TIMFileElem".into(),
|
||||
msg_content: MsgContent {
|
||||
uuid: Some(uuid_str),
|
||||
file_name: Some(filename.to_string()),
|
||||
file_size: Some(size),
|
||||
url: Some(url.to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
}]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn png_dims_parse() {
|
||||
let png = [
|
||||
0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48,
|
||||
0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06,
|
||||
];
|
||||
let d = parse_image_size(&png).expect("png parse");
|
||||
assert_eq!(d.width, 1);
|
||||
assert_eq!(d.height, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gif_dims_parse() {
|
||||
let gif = b"GIF89a\x40\x01\xF0\x00rest";
|
||||
let d = parse_image_size(gif).expect("gif parse");
|
||||
assert_eq!(d.width, 320);
|
||||
assert_eq!(d.height, 240);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn guess_mime_basic() {
|
||||
assert_eq!(guess_mime_type("foo.png"), "image/png");
|
||||
assert_eq!(guess_mime_type("doc.pdf"), "application/pdf");
|
||||
assert_eq!(guess_mime_type("blob"), "application/octet-stream");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_image_works() {
|
||||
assert!(is_image("a.jpeg", ""));
|
||||
assert!(is_image("noext", "image/png"));
|
||||
assert!(!is_image("a.pdf", ""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_format_code_matrix() {
|
||||
assert_eq!(image_format_code("image/png"), 3);
|
||||
assert_eq!(image_format_code("image/jpeg"), 1);
|
||||
assert_eq!(image_format_code("image/gif"), 2);
|
||||
assert_eq!(image_format_code("image/bmp"), 4);
|
||||
assert_eq!(image_format_code("image/webp"), 255);
|
||||
assert_eq!(image_format_code("application/pdf"), 255);
|
||||
}
|
||||
|
||||
// ─── extended MIME / image-format tests ─────────────────────
|
||||
|
||||
#[test]
|
||||
fn guess_mime_handles_uppercase_extension() {
|
||||
assert_eq!(guess_mime_type("PHOTO.JPG"), "image/jpeg");
|
||||
assert_eq!(guess_mime_type("Doc.PDF"), "application/pdf");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn guess_mime_covers_office_audio_video_archive_types() {
|
||||
assert_eq!(guess_mime_type("file.jpg"), "image/jpeg");
|
||||
assert_eq!(guess_mime_type("file.jpeg"), "image/jpeg");
|
||||
assert_eq!(guess_mime_type("file.gif"), "image/gif");
|
||||
assert_eq!(guess_mime_type("file.webp"), "image/webp");
|
||||
assert_eq!(guess_mime_type("file.bmp"), "image/bmp");
|
||||
assert_eq!(guess_mime_type("file.heic"), "image/heic");
|
||||
assert_eq!(guess_mime_type("file.tiff"), "image/tiff");
|
||||
assert_eq!(guess_mime_type("file.ico"), "image/x-icon");
|
||||
assert_eq!(guess_mime_type("file.doc"), "application/msword");
|
||||
assert_eq!(
|
||||
guess_mime_type("file.docx"),
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
);
|
||||
assert_eq!(guess_mime_type("file.xls"), "application/vnd.ms-excel");
|
||||
assert_eq!(
|
||||
guess_mime_type("file.xlsx"),
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
|
||||
);
|
||||
assert_eq!(guess_mime_type("file.ppt"), "application/vnd.ms-powerpoint");
|
||||
assert_eq!(
|
||||
guess_mime_type("file.pptx"),
|
||||
"application/vnd.openxmlformats-officedocument.presentationml.presentation"
|
||||
);
|
||||
assert_eq!(guess_mime_type("file.txt"), "text/plain");
|
||||
assert_eq!(guess_mime_type("file.zip"), "application/zip");
|
||||
assert_eq!(guess_mime_type("file.tar"), "application/x-tar");
|
||||
assert_eq!(guess_mime_type("file.gz"), "application/gzip");
|
||||
assert_eq!(guess_mime_type("file.mp3"), "audio/mpeg");
|
||||
assert_eq!(guess_mime_type("file.mp4"), "video/mp4");
|
||||
assert_eq!(guess_mime_type("file.wav"), "audio/wav");
|
||||
assert_eq!(guess_mime_type("file.ogg"), "audio/ogg");
|
||||
assert_eq!(guess_mime_type("file.webm"), "video/webm");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_format_code_jpg_alias_and_heic_tiff() {
|
||||
assert_eq!(image_format_code("image/jpg"), 1);
|
||||
assert_eq!(image_format_code("image/heic"), 255);
|
||||
assert_eq!(image_format_code("image/tiff"), 255);
|
||||
assert_eq!(image_format_code(""), 255);
|
||||
}
|
||||
|
||||
// ─── parse_image_size — JPEG / WEBP / negative paths ────────
|
||||
|
||||
#[test]
|
||||
fn jpeg_dims_from_sof0_marker() {
|
||||
// SOI + filler APP0 segment + SOF0 marker carrying h=2, w=3.
|
||||
// parse_jpeg's read uses buf[i+5..i+9] and gates on `i + 9 < buf.len()`,
|
||||
// so trailing pad bytes are required (one tail byte makes 17 < 18 true).
|
||||
let mut jpeg = vec![0xFF, 0xD8];
|
||||
// APP0 (0xFFE0) with len=4 → 2 bytes payload
|
||||
jpeg.extend_from_slice(&[0xFF, 0xE0, 0x00, 0x04, 0x00, 0x00]);
|
||||
// SOF0: 0xFF 0xC0 len(11) precision(8) height(2 BE) width(3 BE)
|
||||
jpeg.extend_from_slice(&[0xFF, 0xC0, 0x00, 0x0B, 0x08, 0x00, 0x02, 0x00, 0x03]);
|
||||
// trailing pad so the `i + 9 < buf.len()` loop guard accepts the SOF0
|
||||
// entry on the second iteration.
|
||||
jpeg.push(0xFF);
|
||||
let d = parse_image_size(&jpeg).expect("jpeg parse");
|
||||
assert_eq!(d.width, 3);
|
||||
assert_eq!(d.height, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jpeg_too_short_returns_none() {
|
||||
assert!(parse_image_size(&[0xFF, 0xD8]).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn jpeg_wrong_magic_returns_none() {
|
||||
let buf = [0xCA, 0xFE, 0xBA, 0xBE, 0, 0, 0, 0, 0, 0];
|
||||
assert!(parse_image_size(&buf).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn webp_vp8x_dims_parse() {
|
||||
// RIFF size WEBP VP8X flags(4) padding(3) w-1(LE24) h-1(LE24)
|
||||
// w=320 → 319 little-endian = [0x3F, 0x01, 0x00]; h=240 → 239 = [0xEF, 0x00, 0x00]
|
||||
let mut buf = b"RIFF\x00\x00\x00\x00WEBPVP8X".to_vec();
|
||||
buf.extend_from_slice(&[0u8; 8]); // flags + reserved
|
||||
buf.extend_from_slice(&[0x3F, 0x01, 0x00, 0xEF, 0x00, 0x00]);
|
||||
let d = parse_image_size(&buf).expect("webp vp8x parse");
|
||||
assert_eq!(d.width, 320);
|
||||
assert_eq!(d.height, 240);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn webp_too_short_returns_none() {
|
||||
let buf = b"RIFF\0\0\0\0WEBPVP8";
|
||||
assert!(parse_image_size(buf).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn webp_unsupported_chunk_returns_none() {
|
||||
let mut buf = b"RIFF\0\0\0\0WEBPXXXX".to_vec();
|
||||
buf.extend_from_slice(&[0u8; 30]);
|
||||
assert!(parse_image_size(&buf).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn png_short_or_wrong_magic_returns_none() {
|
||||
assert!(parse_image_size(&[0x89, 0x50, 0x4E, 0x47]).is_none()); // too short
|
||||
let mut buf = vec![0xFF; 24];
|
||||
buf[..4].copy_from_slice(&[0x89, 0x50, 0x4F, 0x47]); // wrong magic
|
||||
assert!(parse_image_size(&buf).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gif_too_short_or_wrong_sig_returns_none() {
|
||||
assert!(parse_image_size(b"GIF87").is_none());
|
||||
assert!(parse_image_size(b"NOTGIFEXT").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_image_size_empty_returns_none() {
|
||||
assert!(parse_image_size(&[]).is_none());
|
||||
}
|
||||
|
||||
// ─── msg_body builders ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn build_image_msg_body_uses_uuid_when_present() {
|
||||
let body = build_image_msg_body(
|
||||
"https://x/cat.png",
|
||||
Some("uuid-1"),
|
||||
Some("cat.png"),
|
||||
1024,
|
||||
800,
|
||||
600,
|
||||
"image/png",
|
||||
);
|
||||
assert_eq!(body.len(), 1);
|
||||
let el = &body[0];
|
||||
assert_eq!(el.msg_type, "TIMImageElem");
|
||||
assert_eq!(el.msg_content.uuid.as_deref(), Some("uuid-1"));
|
||||
assert_eq!(el.msg_content.image_format, Some(3)); // png
|
||||
assert_eq!(el.msg_content.image_info_array.len(), 1);
|
||||
let info = &el.msg_content.image_info_array[0];
|
||||
assert_eq!(info.image_type, 1);
|
||||
assert_eq!(info.size, 1024);
|
||||
assert_eq!(info.width, 800);
|
||||
assert_eq!(info.height, 600);
|
||||
assert_eq!(info.url, "https://x/cat.png");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_image_msg_body_falls_back_to_filename_then_default_uuid() {
|
||||
let with_filename = build_image_msg_body(
|
||||
"https://x/",
|
||||
None,
|
||||
Some("only-name.png"),
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
"image/png",
|
||||
);
|
||||
assert_eq!(
|
||||
with_filename[0].msg_content.uuid.as_deref(),
|
||||
Some("only-name.png")
|
||||
);
|
||||
|
||||
let default_id = build_image_msg_body("https://x/", None, None, 0, 0, 0, "image/png");
|
||||
assert_eq!(default_id[0].msg_content.uuid.as_deref(), Some("image"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_image_msg_body_treats_empty_mime_as_format_255() {
|
||||
let body = build_image_msg_body("https://x/cat.jpg", None, None, 0, 0, 0, "");
|
||||
assert_eq!(body[0].msg_content.image_format, Some(255));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_file_msg_body_uses_filename_when_uuid_missing() {
|
||||
let body = build_file_msg_body("https://x/doc.pdf", "doc.pdf", None, 2048);
|
||||
assert_eq!(body.len(), 1);
|
||||
let el = &body[0];
|
||||
assert_eq!(el.msg_type, "TIMFileElem");
|
||||
assert_eq!(el.msg_content.uuid.as_deref(), Some("doc.pdf"));
|
||||
assert_eq!(el.msg_content.file_name.as_deref(), Some("doc.pdf"));
|
||||
assert_eq!(el.msg_content.file_size, Some(2048));
|
||||
assert_eq!(el.msg_content.url.as_deref(), Some("https://x/doc.pdf"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_file_msg_body_prefers_explicit_uuid() {
|
||||
let body = build_file_msg_body("https://x/y.pdf", "y.pdf", Some("uuid-y"), 0);
|
||||
assert_eq!(body[0].msg_content.uuid.as_deref(), Some("uuid-y"));
|
||||
}
|
||||
|
||||
// ─── download_url (wiremock) ────────────────────────────────
|
||||
|
||||
#[tokio::test]
|
||||
async fn download_url_returns_bytes_and_content_type() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("HEAD"))
|
||||
.respond_with(wiremock::ResponseTemplate::new(200).insert_header("Content-Length", "3"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("GET"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200)
|
||||
.insert_header("Content-Type", "image/png; charset=binary")
|
||||
.set_body_bytes(vec![1u8, 2, 3]),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let http = reqwest::Client::new();
|
||||
let (bytes, mime) = download_url(&http, &server.uri(), 10).await.unwrap();
|
||||
assert_eq!(bytes, vec![1, 2, 3]);
|
||||
assert_eq!(mime, "image/png");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn download_url_rejects_oversize_from_head_content_length() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
// HEAD reports a very large file → reject BEFORE GET.
|
||||
wiremock::Mock::given(wiremock::matchers::method("HEAD"))
|
||||
.respond_with(wiremock::ResponseTemplate::new(200).insert_header(
|
||||
"Content-Length",
|
||||
(10u64 * 1024 * 1024 + 1).to_string().as_str(),
|
||||
))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let http = reqwest::Client::new();
|
||||
let err = download_url(&http, &server.uri(), 10).await.unwrap_err();
|
||||
match err {
|
||||
YuanbaoError::Media(m) => assert!(m.contains("too large"), "got {m}"),
|
||||
other => panic!("expected Media error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn download_url_rejects_when_body_exceeds_limit() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("HEAD"))
|
||||
.respond_with(wiremock::ResponseTemplate::new(200))
|
||||
.mount(&server)
|
||||
.await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("GET"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_bytes(vec![0u8; 2 * 1024 * 1024]), // 2 MiB
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let http = reqwest::Client::new();
|
||||
let err = download_url(&http, &server.uri(), 1).await.unwrap_err();
|
||||
match err {
|
||||
YuanbaoError::Media(m) => assert!(m.contains("exceeds limit"), "got {m}"),
|
||||
other => panic!("expected Media error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn download_url_surfaces_http_error_status() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("HEAD"))
|
||||
.respond_with(wiremock::ResponseTemplate::new(404))
|
||||
.mount(&server)
|
||||
.await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("GET"))
|
||||
.respond_with(wiremock::ResponseTemplate::new(404))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let http = reqwest::Client::new();
|
||||
let err = download_url(&http, &server.uri(), 10).await.unwrap_err();
|
||||
match err {
|
||||
YuanbaoError::Media(m) => assert!(m.contains("HTTP 404"), "got {m}"),
|
||||
other => panic!("expected Media error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,29 +0,0 @@
|
||||
//! Yuanbao (元宝) channel provider.
|
||||
//!
|
||||
//! This module is intentionally export-focused. Operational code lives in
|
||||
//! sibling modules:
|
||||
//! - [`channel`] wires the provider into the generic OpenHuman `Channel` trait.
|
||||
//! - [`connection`] owns the WebSocket transport and request correlator.
|
||||
//! - [`inbound`] owns inbound filtering/extraction.
|
||||
//! - [`outbound`] owns Yuanbao send/query calls.
|
||||
//! - [`proto`] / [`proto_biz`] / [`wire`] own hand-written protobuf codecs.
|
||||
|
||||
pub mod channel;
|
||||
pub mod config;
|
||||
pub mod connection;
|
||||
pub mod cos;
|
||||
pub mod errors;
|
||||
pub mod ids;
|
||||
pub mod inbound;
|
||||
pub mod media;
|
||||
pub mod outbound;
|
||||
pub mod proto;
|
||||
pub mod proto_biz;
|
||||
pub mod proto_constants;
|
||||
pub mod sign;
|
||||
pub mod splitter;
|
||||
pub mod types;
|
||||
pub mod wire;
|
||||
|
||||
pub use channel::YuanbaoChannel;
|
||||
pub use config::YuanbaoConfig;
|
||||
@@ -1,449 +0,0 @@
|
||||
//! Outbound message sender.
|
||||
//!
|
||||
//! Translates high-level `send_text` / `send_image` / heartbeat calls
|
||||
//! into encoded ConnMsg frames and pushes them through the shared
|
||||
//! `YuanbaoConnection`. The recipient string uses the convention
|
||||
//! `g:<group_code>` for groups, raw `<uid>` for DMs.
|
||||
//!
|
||||
//! For the few request kinds where we care about the response body
|
||||
//! (notably `QueryGroupInfo`, `GetGroupMemberList`, and `SendXxxMessage`'s
|
||||
//! `code/msg_id` echo) we use the connection-level pending-acks
|
||||
//! correlator and return parsed results to the caller. Heartbeats are
|
||||
//! fire-and-forget — the response is never inspected.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use tracing::debug;
|
||||
|
||||
use super::connection::YuanbaoConnection;
|
||||
use super::cos::{get_cos_credentials, upload_to_cos};
|
||||
use super::errors::YuanbaoError;
|
||||
use super::media::{build_file_msg_body, build_image_msg_body, download_url, parse_image_size};
|
||||
use super::proto_biz::{
|
||||
decode_get_group_member_list_rsp, decode_query_group_info_rsp, encode_get_group_member_list,
|
||||
encode_query_group_info, encode_send_c2c_message, encode_send_group_heartbeat,
|
||||
encode_send_group_message, encode_send_private_heartbeat,
|
||||
};
|
||||
use super::proto_constants::{ws_heartbeat, DEFAULT_SEND_TIMEOUT_SECS};
|
||||
use super::sign::SignManager;
|
||||
use super::types::{GroupInfo, GroupMemberListPage, MsgBodyElement};
|
||||
|
||||
const GROUP_PREFIX: &str = "g:";
|
||||
/// Wait-for-response timeout on queries like `QueryGroupInfo`.
|
||||
const QUERY_TIMEOUT_SECS: u64 = 10;
|
||||
|
||||
/// Parsed addressing target.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum Target<'a> {
|
||||
Dm(&'a str),
|
||||
Group(&'a str),
|
||||
}
|
||||
|
||||
impl<'a> Target<'a> {
|
||||
pub fn parse(recipient: &'a str) -> Self {
|
||||
if let Some(rest) = recipient.strip_prefix(GROUP_PREFIX) {
|
||||
Self::Group(rest)
|
||||
} else {
|
||||
Self::Dm(recipient)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct OutboundSender {
|
||||
conn: Arc<YuanbaoConnection>,
|
||||
/// Sign-token cache holding the server-issued `bot_id`. Populated as a
|
||||
/// side effect of `connection`'s sign+auth-bind flow. The `bot_id` here
|
||||
/// is what `yuanbao_openclaw_proxy` expects in the outbound
|
||||
/// `from_account` field — config-only fallbacks like `app_key` get
|
||||
/// silently accepted (status=0) but never routed to a real conv id.
|
||||
sign_manager: Option<Arc<SignManager>>,
|
||||
/// Lookup key for `sign_manager.cached(app_key)`.
|
||||
app_key: String,
|
||||
/// User-supplied bot id override; empty when not set. Only used when
|
||||
/// the sign cache hasn't been primed yet (e.g. send-before-auth races).
|
||||
config_bot_id: String,
|
||||
http: reqwest::Client,
|
||||
}
|
||||
|
||||
impl OutboundSender {
|
||||
pub fn new(
|
||||
conn: Arc<YuanbaoConnection>,
|
||||
sign_manager: Option<Arc<SignManager>>,
|
||||
app_key: String,
|
||||
config_bot_id: String,
|
||||
) -> Self {
|
||||
Self {
|
||||
conn,
|
||||
sign_manager,
|
||||
app_key,
|
||||
config_bot_id,
|
||||
http: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the `from_account` to put on the next outbound frame.
|
||||
/// Prefers the server-issued `bot_id` cached after sign-token / auth-bind
|
||||
/// — matches hermes-agent `_bot_id = token_data["bot_id"]` (yuanbao.py:400).
|
||||
async fn resolve_from_account(&self) -> String {
|
||||
if let Some(sign) = &self.sign_manager {
|
||||
if let Some(entry) = sign.cached(&self.app_key).await {
|
||||
if !entry.bot_id.is_empty() {
|
||||
return entry.bot_id;
|
||||
}
|
||||
}
|
||||
}
|
||||
self.config_bot_id.clone()
|
||||
}
|
||||
|
||||
/// Send a plain-text message. Returns the client-side `msg_id`.
|
||||
pub async fn send_text(
|
||||
&self,
|
||||
recipient: &str,
|
||||
text: &str,
|
||||
ref_msg_id: Option<&str>,
|
||||
) -> Result<String, YuanbaoError> {
|
||||
let body = vec![MsgBodyElement {
|
||||
msg_type: "TIMTextElem".into(),
|
||||
msg_content: super::types::MsgContent {
|
||||
text: Some(text.to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
}];
|
||||
self.send_body(recipient, body, ref_msg_id).await
|
||||
}
|
||||
|
||||
/// Send an image by an already-uploaded (COS or other) URL.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn send_image_url(
|
||||
&self,
|
||||
recipient: &str,
|
||||
url: &str,
|
||||
size: u32,
|
||||
width: u32,
|
||||
height: u32,
|
||||
mime_type: &str,
|
||||
) -> Result<String, YuanbaoError> {
|
||||
let body = build_image_msg_body(url, None, None, size, width, height, mime_type);
|
||||
self.send_body(recipient, body, None).await
|
||||
}
|
||||
|
||||
/// End-to-end image send: download from URL → upload to COS → send
|
||||
/// as a `TIMImageElem`. Returns the outbound `msg_id`.
|
||||
///
|
||||
/// `app_key` / `bot_id` / `token` / `api_domain` / `route_env` come
|
||||
/// from the channel config; pass them in rather than reaching back
|
||||
/// through the conn to keep this fn easy to unit-test.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub async fn send_image_from_url(
|
||||
&self,
|
||||
recipient: &str,
|
||||
source_url: &str,
|
||||
app_key: &str,
|
||||
bot_id: &str,
|
||||
token: &str,
|
||||
api_domain: &str,
|
||||
route_env: &str,
|
||||
max_size_mb: u64,
|
||||
) -> Result<String, YuanbaoError> {
|
||||
let (bytes, mime) = download_url(&self.http, source_url, max_size_mb).await?;
|
||||
let dims = parse_image_size(&bytes);
|
||||
let width = dims.as_ref().map(|d| d.width).unwrap_or(0);
|
||||
let height = dims.as_ref().map(|d| d.height).unwrap_or(0);
|
||||
|
||||
let filename = extract_filename(source_url);
|
||||
let creds = get_cos_credentials(
|
||||
&self.http, api_domain, app_key, bot_id, token, route_env, &filename,
|
||||
)
|
||||
.await?;
|
||||
let upload = upload_to_cos(&self.http, &creds, &bytes, &filename, mime.clone()).await?;
|
||||
|
||||
let final_width = if upload.width > 0 {
|
||||
upload.width
|
||||
} else {
|
||||
width
|
||||
};
|
||||
let final_height = if upload.height > 0 {
|
||||
upload.height
|
||||
} else {
|
||||
height
|
||||
};
|
||||
let body = build_image_msg_body(
|
||||
&upload.url,
|
||||
Some(&upload.uuid),
|
||||
Some(&filename),
|
||||
upload.size as u32,
|
||||
final_width,
|
||||
final_height,
|
||||
&mime,
|
||||
);
|
||||
self.send_body(recipient, body, None).await
|
||||
}
|
||||
|
||||
/// Send a file by URL.
|
||||
pub async fn send_file_url(
|
||||
&self,
|
||||
recipient: &str,
|
||||
url: &str,
|
||||
file_name: &str,
|
||||
size: u32,
|
||||
) -> Result<String, YuanbaoError> {
|
||||
let body = build_file_msg_body(url, file_name, None, size);
|
||||
self.send_body(recipient, body, None).await
|
||||
}
|
||||
|
||||
/// Send a pre-built `msg_body`. Waits up to `DEFAULT_SEND_TIMEOUT_SECS`
|
||||
/// for the server response so the caller learns about delivery
|
||||
/// failures (rate-limit, banned content, etc.) instead of getting a
|
||||
/// silent drop.
|
||||
pub async fn send_body(
|
||||
&self,
|
||||
recipient: &str,
|
||||
msg_body: Vec<MsgBodyElement>,
|
||||
ref_msg_id: Option<&str>,
|
||||
) -> Result<String, YuanbaoError> {
|
||||
let msg_id = self.next_msg_id();
|
||||
let target = Target::parse(recipient);
|
||||
let from_account = self.resolve_from_account().await;
|
||||
let frame = match target {
|
||||
Target::Dm(uid) => encode_send_c2c_message(
|
||||
uid,
|
||||
&from_account,
|
||||
&msg_body,
|
||||
&msg_id,
|
||||
random_u32(),
|
||||
"",
|
||||
"",
|
||||
),
|
||||
Target::Group(group_code) => {
|
||||
let random = format!("{}", random_u32());
|
||||
encode_send_group_message(
|
||||
group_code,
|
||||
&from_account,
|
||||
&msg_body,
|
||||
&msg_id,
|
||||
"",
|
||||
&random,
|
||||
ref_msg_id.unwrap_or(""),
|
||||
"",
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
let timeout = Duration::from_secs(DEFAULT_SEND_TIMEOUT_SECS);
|
||||
match self.conn.send_and_wait(&msg_id, frame, timeout).await {
|
||||
Ok(resp) => {
|
||||
if resp.status != 0 {
|
||||
return Err(YuanbaoError::SendFailed(format!(
|
||||
"server status={} cmd={}",
|
||||
resp.status, resp.cmd
|
||||
)));
|
||||
}
|
||||
debug!("[outbound] ack msg_id={msg_id} target={:?}", target);
|
||||
Ok(msg_id)
|
||||
}
|
||||
// If the correlator isn't usable yet (NotConnected etc.) bubble up.
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a "thinking" heartbeat (RUNNING) — fire-and-forget.
|
||||
pub async fn start_heartbeat(&self, recipient: &str) -> Result<(), YuanbaoError> {
|
||||
self.send_heartbeat(recipient, ws_heartbeat::RUNNING).await
|
||||
}
|
||||
|
||||
/// Send a "done" heartbeat (FINISH) — fire-and-forget.
|
||||
pub async fn stop_heartbeat(&self, recipient: &str) -> Result<(), YuanbaoError> {
|
||||
self.send_heartbeat(recipient, ws_heartbeat::FINISH).await
|
||||
}
|
||||
|
||||
async fn send_heartbeat(&self, recipient: &str, heartbeat: u32) -> Result<(), YuanbaoError> {
|
||||
let req_id = self.conn.next_msg_id("hb");
|
||||
let from_account = self.resolve_from_account().await;
|
||||
let frame = match Target::parse(recipient) {
|
||||
Target::Dm(uid) => {
|
||||
encode_send_private_heartbeat(&req_id, &from_account, uid, heartbeat)
|
||||
}
|
||||
Target::Group(group_code) => {
|
||||
let now_ms = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as u64)
|
||||
.unwrap_or(0);
|
||||
encode_send_group_heartbeat(&req_id, &from_account, group_code, heartbeat, now_ms)
|
||||
}
|
||||
};
|
||||
// Fire-and-forget — we don't care about the heartbeat ack.
|
||||
self.conn.send_conn_msg(frame).await
|
||||
}
|
||||
|
||||
/// Query group info and wait for the server's reply.
|
||||
pub async fn query_group_info(&self, group_code: &str) -> Result<GroupInfo, YuanbaoError> {
|
||||
let req_id = self.conn.next_msg_id("qgi");
|
||||
let frame = encode_query_group_info(&req_id, group_code);
|
||||
let resp = self
|
||||
.conn
|
||||
.send_and_wait(&req_id, frame, Duration::from_secs(QUERY_TIMEOUT_SECS))
|
||||
.await?;
|
||||
decode_query_group_info_rsp(&resp.data)
|
||||
}
|
||||
|
||||
/// Fetch one page of group members. Use `offset=0, limit=100` for the
|
||||
/// first page; the response carries `next_offset` for pagination.
|
||||
pub async fn query_group_members(
|
||||
&self,
|
||||
group_code: &str,
|
||||
offset: u32,
|
||||
limit: u32,
|
||||
) -> Result<GroupMemberListPage, YuanbaoError> {
|
||||
let req_id = self.conn.next_msg_id("qgm");
|
||||
let frame = encode_get_group_member_list(&req_id, group_code, offset, limit);
|
||||
let resp = self
|
||||
.conn
|
||||
.send_and_wait(&req_id, frame, Duration::from_secs(QUERY_TIMEOUT_SECS))
|
||||
.await?;
|
||||
decode_get_group_member_list_rsp(&resp.data)
|
||||
}
|
||||
|
||||
fn next_msg_id(&self) -> String {
|
||||
// Use a stable prefix so logs can be grepped across send paths.
|
||||
self.conn.next_msg_id("om")
|
||||
}
|
||||
}
|
||||
|
||||
fn random_u32() -> u32 {
|
||||
rand::random::<u32>()
|
||||
}
|
||||
|
||||
/// Best-effort file name extraction from a URL — uses the URL's path
|
||||
/// component (so the host is never picked up as a filename) and falls
|
||||
/// back to "file" if there's nothing usable.
|
||||
fn extract_filename(url_str: &str) -> String {
|
||||
if let Ok(parsed) = url::Url::parse(url_str) {
|
||||
if let Some(segments) = parsed.path_segments() {
|
||||
if let Some(last) = segments.filter(|s| !s.is_empty()).last() {
|
||||
return last.to_string();
|
||||
}
|
||||
}
|
||||
return "file".to_string();
|
||||
}
|
||||
// Non-URL input (relative path, raw filename, etc.) — fall back to
|
||||
// last non-empty `/`-delimited segment.
|
||||
let without_query = url_str.split('?').next().unwrap_or(url_str);
|
||||
let last = without_query
|
||||
.rsplit('/')
|
||||
.find(|s| !s.is_empty())
|
||||
.unwrap_or("");
|
||||
if last.is_empty() {
|
||||
"file".to_string()
|
||||
} else {
|
||||
last.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn target_parse_dm() {
|
||||
match Target::parse("user_42") {
|
||||
Target::Dm(uid) => assert_eq!(uid, "user_42"),
|
||||
_ => panic!("should be DM"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn target_parse_group() {
|
||||
match Target::parse("g:room_99") {
|
||||
Target::Group(c) => assert_eq!(c, "room_99"),
|
||||
_ => panic!("should be group"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn target_parse_empty_dm() {
|
||||
assert!(matches!(Target::parse(""), Target::Dm("")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_filename_strips_query() {
|
||||
assert_eq!(extract_filename("https://x.com/a/b/cat.png"), "cat.png");
|
||||
assert_eq!(
|
||||
extract_filename("https://x.com/a/b/cat.png?sig=abc"),
|
||||
"cat.png"
|
||||
);
|
||||
assert_eq!(extract_filename("https://x.com/"), "file");
|
||||
assert_eq!(extract_filename(""), "file");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_filename_from_bare_path() {
|
||||
// Not a valid URL → fall back to last non-empty `/`-segment.
|
||||
assert_eq!(extract_filename("/var/log/foo.bin"), "foo.bin");
|
||||
// Trailing slash gets skipped; last non-empty segment wins.
|
||||
assert_eq!(extract_filename("/var/log/"), "log");
|
||||
// Plain filename with no slashes.
|
||||
assert_eq!(extract_filename("plain.txt"), "plain.txt");
|
||||
}
|
||||
|
||||
fn make_conn(cfg: super::super::config::YuanbaoConfig) -> Arc<YuanbaoConnection> {
|
||||
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
YuanbaoConnection::new(cfg, tx, None)
|
||||
}
|
||||
|
||||
fn base_cfg() -> super::super::config::YuanbaoConfig {
|
||||
let mut c = super::super::config::YuanbaoConfig::default();
|
||||
c.app_key = "ak".into();
|
||||
c.ws_domain = "wss://x".into();
|
||||
c.token = "tok".into();
|
||||
c.bot_id = "cfg-bot".into();
|
||||
c
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_from_account_uses_config_bot_id_when_no_sign_manager() {
|
||||
let conn = make_conn(base_cfg());
|
||||
let sender = OutboundSender::new(conn, None, "ak".into(), "cfg-bot".into());
|
||||
assert_eq!(sender.resolve_from_account().await, "cfg-bot");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_from_account_uses_sign_cache_when_bot_id_present() {
|
||||
let conn = make_conn(base_cfg());
|
||||
let mgr = super::super::sign::SignManager::new(reqwest::Client::new());
|
||||
// Seed the cache with a bot_id keyed on the same app_key.
|
||||
mgr.set_cached_for_test(
|
||||
"ak",
|
||||
super::super::sign::TokenEntry {
|
||||
token: "tok".into(),
|
||||
bot_id: "server-bot".into(),
|
||||
product: String::new(),
|
||||
source: "bot".into(),
|
||||
expire_ts: u64::MAX / 2,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let sender = OutboundSender::new(conn, Some(mgr), "ak".into(), "fallback-bot".into());
|
||||
// Sign cache hit → use server bot_id, not the fallback.
|
||||
assert_eq!(sender.resolve_from_account().await, "server-bot");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_from_account_falls_back_when_sign_cache_bot_id_empty() {
|
||||
let conn = make_conn(base_cfg());
|
||||
let mgr = super::super::sign::SignManager::new(reqwest::Client::new());
|
||||
mgr.set_cached_for_test(
|
||||
"ak",
|
||||
super::super::sign::TokenEntry {
|
||||
token: "tok".into(),
|
||||
bot_id: String::new(),
|
||||
product: String::new(),
|
||||
source: String::new(),
|
||||
expire_ts: u64::MAX / 2,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let sender = OutboundSender::new(conn, Some(mgr), "ak".into(), "fallback-bot".into());
|
||||
assert_eq!(sender.resolve_from_account().await, "fallback-bot");
|
||||
}
|
||||
}
|
||||
@@ -1,916 +0,0 @@
|
||||
//! Yuanbao WebSocket ConnMsg envelope + built-in protocol commands
|
||||
//! (auth-bind, ping, push-ack) + TIM `MsgBodyElement` codecs.
|
||||
//!
|
||||
//! Each WebSocket binary frame carries one full `ConnMsg` protobuf
|
||||
//! message; **no extra length prefix is needed** (the WS frame boundary
|
||||
//! delimits one message). Verified against the hermes-agent Python
|
||||
//! reference (yuanbao_proto.py) and the TypeScript openclaw plugin.
|
||||
//!
|
||||
//! Business-layer codecs (send-message / heartbeat / group query) live
|
||||
//! in [`super::proto_biz`]. Wire-format primitives (varint, FieldValue,
|
||||
//! parse_fields) live in [`super::wire`].
|
||||
|
||||
use super::errors::YuanbaoError;
|
||||
use super::proto_constants::*;
|
||||
use super::types::*;
|
||||
use super::wire::{
|
||||
encode_field_bytes, encode_field_string, encode_field_varint, get_bytes, get_repeated_bytes,
|
||||
get_string, get_varint, next_seq_no, parse_fields, FieldValue,
|
||||
};
|
||||
|
||||
// Re-export wire primitives for downstream callers (tests, tools).
|
||||
pub use super::wire::{decode_varint, encode_varint};
|
||||
|
||||
// ─── ConnMsg envelope ──────────────────────────────────────────────
|
||||
//
|
||||
// message Head {
|
||||
// uint32 cmd_type = 1;
|
||||
// string cmd = 2;
|
||||
// uint32 seq_no = 3;
|
||||
// string msg_id = 4;
|
||||
// string module = 5;
|
||||
// bool need_ack = 6;
|
||||
// int32 status = 10;
|
||||
// }
|
||||
// message ConnMsg {
|
||||
// Head head = 1;
|
||||
// bytes data = 2;
|
||||
// }
|
||||
|
||||
fn encode_head(
|
||||
cmd_type: u32,
|
||||
cmd: &str,
|
||||
seq_no: u32,
|
||||
msg_id: &str,
|
||||
module: &str,
|
||||
need_ack: bool,
|
||||
status: u32,
|
||||
) -> Vec<u8> {
|
||||
let mut buf = Vec::with_capacity(64);
|
||||
if cmd_type != 0 {
|
||||
encode_field_varint(1, cmd_type as u64, &mut buf);
|
||||
}
|
||||
if !cmd.is_empty() {
|
||||
encode_field_string(2, cmd, &mut buf);
|
||||
}
|
||||
if seq_no != 0 {
|
||||
encode_field_varint(3, seq_no as u64, &mut buf);
|
||||
}
|
||||
if !msg_id.is_empty() {
|
||||
encode_field_string(4, msg_id, &mut buf);
|
||||
}
|
||||
if !module.is_empty() {
|
||||
encode_field_string(5, module, &mut buf);
|
||||
}
|
||||
if need_ack {
|
||||
encode_field_varint(6, 1, &mut buf);
|
||||
}
|
||||
if status != 0 {
|
||||
encode_field_varint(10, status as u64, &mut buf);
|
||||
}
|
||||
buf
|
||||
}
|
||||
|
||||
/// Encode a full `ConnMsg` frame (ready to send as a binary WS frame).
|
||||
pub fn encode_conn_msg(
|
||||
cmd_type: u32,
|
||||
cmd: &str,
|
||||
seq_no: u32,
|
||||
msg_id: &str,
|
||||
module: &str,
|
||||
data: &[u8],
|
||||
) -> Vec<u8> {
|
||||
let head = encode_head(cmd_type, cmd, seq_no, msg_id, module, false, 0);
|
||||
let mut buf = Vec::with_capacity(head.len() + data.len() + 16);
|
||||
encode_field_bytes(1, &head, &mut buf);
|
||||
if !data.is_empty() {
|
||||
encode_field_bytes(2, data, &mut buf);
|
||||
}
|
||||
buf
|
||||
}
|
||||
|
||||
/// Decode a `ConnMsg` frame received from the gateway.
|
||||
pub fn decode_conn_msg(data: &[u8]) -> Result<ConnFrame, YuanbaoError> {
|
||||
let fields = parse_fields(data)?;
|
||||
let head_bytes = get_bytes(&fields, 1);
|
||||
let payload = get_bytes(&fields, 2);
|
||||
let head_fields = if head_bytes.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
parse_fields(&head_bytes)?
|
||||
};
|
||||
Ok(ConnFrame {
|
||||
cmd_type: get_varint(&head_fields, 1) as u32,
|
||||
cmd: get_string(&head_fields, 2),
|
||||
seq_no: get_varint(&head_fields, 3) as u32,
|
||||
msg_id: get_string(&head_fields, 4),
|
||||
module: get_string(&head_fields, 5),
|
||||
need_ack: get_varint(&head_fields, 6) != 0,
|
||||
status: get_varint(&head_fields, 10) as u32,
|
||||
data: payload,
|
||||
})
|
||||
}
|
||||
|
||||
// ─── Built-in protocol commands ────────────────────────────────────
|
||||
|
||||
/// `AuthBindReq` — first request after the WebSocket opens.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn encode_auth_bind(
|
||||
biz_id: &str,
|
||||
uid: &str,
|
||||
source: &str,
|
||||
token: &str,
|
||||
msg_id: &str,
|
||||
app_version: &str,
|
||||
operation_system: &str,
|
||||
bot_version: &str,
|
||||
route_env: &str,
|
||||
) -> Vec<u8> {
|
||||
let mut auth_buf = Vec::with_capacity(uid.len() + source.len() + token.len() + 16);
|
||||
encode_field_string(1, uid, &mut auth_buf);
|
||||
encode_field_string(2, source, &mut auth_buf);
|
||||
encode_field_string(3, token, &mut auth_buf);
|
||||
|
||||
let mut dev_buf = Vec::with_capacity(64);
|
||||
if !app_version.is_empty() {
|
||||
encode_field_string(1, app_version, &mut dev_buf);
|
||||
}
|
||||
if !operation_system.is_empty() {
|
||||
encode_field_string(2, operation_system, &mut dev_buf);
|
||||
}
|
||||
encode_field_string(10, OPENHUMAN_INSTANCE_ID, &mut dev_buf);
|
||||
if !bot_version.is_empty() {
|
||||
encode_field_string(24, bot_version, &mut dev_buf);
|
||||
}
|
||||
|
||||
let mut req_buf = Vec::with_capacity(auth_buf.len() + dev_buf.len() + biz_id.len() + 16);
|
||||
encode_field_string(1, biz_id, &mut req_buf);
|
||||
encode_field_bytes(2, &auth_buf, &mut req_buf);
|
||||
encode_field_bytes(3, &dev_buf, &mut req_buf);
|
||||
if !route_env.is_empty() {
|
||||
encode_field_string(5, route_env, &mut req_buf);
|
||||
}
|
||||
|
||||
encode_conn_msg(
|
||||
cmd_type::REQUEST,
|
||||
cmd::AUTH_BIND,
|
||||
next_seq_no(),
|
||||
msg_id,
|
||||
module::CONN_ACCESS,
|
||||
&req_buf,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn encode_ping(msg_id: &str) -> Vec<u8> {
|
||||
encode_conn_msg(
|
||||
cmd_type::REQUEST,
|
||||
cmd::PING,
|
||||
next_seq_no(),
|
||||
msg_id,
|
||||
module::CONN_ACCESS,
|
||||
&[],
|
||||
)
|
||||
}
|
||||
|
||||
/// Decoded `AuthBindRsp` body.
|
||||
///
|
||||
/// message AuthBindRsp {
|
||||
/// int32 code = 1;
|
||||
/// string message = 2;
|
||||
/// string connect_id = 3;
|
||||
/// }
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AuthBindRsp {
|
||||
pub code: i32,
|
||||
pub message: String,
|
||||
pub connect_id: String,
|
||||
}
|
||||
|
||||
/// Parse an `AuthBindRsp` from the biz payload (`ConnMsg.data`).
|
||||
pub fn decode_auth_bind_rsp(data: &[u8]) -> Result<AuthBindRsp, YuanbaoError> {
|
||||
let fields = parse_fields(data)?;
|
||||
Ok(AuthBindRsp {
|
||||
code: get_varint(&fields, 1) as i32,
|
||||
message: get_string(&fields, 2),
|
||||
connect_id: get_string(&fields, 3),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn encode_push_ack(original: &ConnFrame) -> Vec<u8> {
|
||||
encode_conn_msg(
|
||||
cmd_type::PUSH_ACK,
|
||||
&original.cmd,
|
||||
next_seq_no(),
|
||||
&original.msg_id,
|
||||
&original.module,
|
||||
&[],
|
||||
)
|
||||
}
|
||||
|
||||
// ─── MsgBodyElement (TIM) encoding ─────────────────────────────────
|
||||
|
||||
pub fn encode_msg_content(c: &MsgContent) -> Vec<u8> {
|
||||
let mut buf = Vec::with_capacity(64);
|
||||
if let Some(ref v) = c.text {
|
||||
if !v.is_empty() {
|
||||
encode_field_string(1, v, &mut buf);
|
||||
}
|
||||
}
|
||||
if let Some(ref v) = c.uuid {
|
||||
if !v.is_empty() {
|
||||
encode_field_string(2, v, &mut buf);
|
||||
}
|
||||
}
|
||||
if let Some(v) = c.image_format {
|
||||
if v != 0 {
|
||||
encode_field_varint(3, v as u64, &mut buf);
|
||||
}
|
||||
}
|
||||
if let Some(ref v) = c.data {
|
||||
if !v.is_empty() {
|
||||
encode_field_string(4, v, &mut buf);
|
||||
}
|
||||
}
|
||||
if let Some(ref v) = c.desc {
|
||||
if !v.is_empty() {
|
||||
encode_field_string(5, v, &mut buf);
|
||||
}
|
||||
}
|
||||
if let Some(ref v) = c.ext {
|
||||
if !v.is_empty() {
|
||||
encode_field_string(6, v, &mut buf);
|
||||
}
|
||||
}
|
||||
if let Some(ref v) = c.sound {
|
||||
if !v.is_empty() {
|
||||
encode_field_string(7, v, &mut buf);
|
||||
}
|
||||
}
|
||||
for img in &c.image_info_array {
|
||||
let mut ib = Vec::with_capacity(48);
|
||||
if img.image_type != 0 {
|
||||
encode_field_varint(1, img.image_type as u64, &mut ib);
|
||||
}
|
||||
if img.size != 0 {
|
||||
encode_field_varint(2, img.size as u64, &mut ib);
|
||||
}
|
||||
if img.width != 0 {
|
||||
encode_field_varint(3, img.width as u64, &mut ib);
|
||||
}
|
||||
if img.height != 0 {
|
||||
encode_field_varint(4, img.height as u64, &mut ib);
|
||||
}
|
||||
if !img.url.is_empty() {
|
||||
encode_field_string(5, &img.url, &mut ib);
|
||||
}
|
||||
encode_field_bytes(8, &ib, &mut buf);
|
||||
}
|
||||
if let Some(v) = c.index {
|
||||
if v != 0 {
|
||||
encode_field_varint(9, v as u64, &mut buf);
|
||||
}
|
||||
}
|
||||
if let Some(ref v) = c.url {
|
||||
if !v.is_empty() {
|
||||
encode_field_string(10, v, &mut buf);
|
||||
}
|
||||
}
|
||||
if let Some(v) = c.file_size {
|
||||
if v != 0 {
|
||||
encode_field_varint(11, v as u64, &mut buf);
|
||||
}
|
||||
}
|
||||
if let Some(ref v) = c.file_name {
|
||||
if !v.is_empty() {
|
||||
encode_field_string(12, v, &mut buf);
|
||||
}
|
||||
}
|
||||
buf
|
||||
}
|
||||
|
||||
fn decode_msg_content(data: &[u8]) -> Result<MsgContent, YuanbaoError> {
|
||||
let fields = parse_fields(data)?;
|
||||
let mut c = MsgContent::default();
|
||||
for (n, v) in &fields {
|
||||
match (*n, v) {
|
||||
(1, FieldValue::Bytes(b)) => c.text = Some(String::from_utf8_lossy(b).into_owned()),
|
||||
(2, FieldValue::Bytes(b)) => c.uuid = Some(String::from_utf8_lossy(b).into_owned()),
|
||||
(3, FieldValue::Varint(x)) => c.image_format = Some(*x as u32),
|
||||
(4, FieldValue::Bytes(b)) => c.data = Some(String::from_utf8_lossy(b).into_owned()),
|
||||
(5, FieldValue::Bytes(b)) => c.desc = Some(String::from_utf8_lossy(b).into_owned()),
|
||||
(6, FieldValue::Bytes(b)) => c.ext = Some(String::from_utf8_lossy(b).into_owned()),
|
||||
(7, FieldValue::Bytes(b)) => c.sound = Some(String::from_utf8_lossy(b).into_owned()),
|
||||
(8, FieldValue::Bytes(b)) => {
|
||||
let ifields = parse_fields(b)?;
|
||||
let mut info = ImageInfo {
|
||||
image_type: get_varint(&ifields, 1) as u32,
|
||||
size: get_varint(&ifields, 2) as u32,
|
||||
width: get_varint(&ifields, 3) as u32,
|
||||
height: get_varint(&ifields, 4) as u32,
|
||||
url: get_string(&ifields, 5),
|
||||
};
|
||||
if info.image_type != 0 || !info.url.is_empty() {
|
||||
if info.image_type == 0 {
|
||||
info.image_type = 1;
|
||||
}
|
||||
c.image_info_array.push(info);
|
||||
}
|
||||
}
|
||||
(9, FieldValue::Varint(x)) => c.index = Some(*x as u32),
|
||||
(10, FieldValue::Bytes(b)) => c.url = Some(String::from_utf8_lossy(b).into_owned()),
|
||||
(11, FieldValue::Varint(x)) => c.file_size = Some(*x as u32),
|
||||
(12, FieldValue::Bytes(b)) => {
|
||||
c.file_name = Some(String::from_utf8_lossy(b).into_owned())
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(c)
|
||||
}
|
||||
|
||||
pub fn encode_msg_body_element(el: &MsgBodyElement) -> Vec<u8> {
|
||||
let mut buf = Vec::with_capacity(64);
|
||||
if !el.msg_type.is_empty() {
|
||||
encode_field_string(1, &el.msg_type, &mut buf);
|
||||
}
|
||||
let content = encode_msg_content(&el.msg_content);
|
||||
if !content.is_empty() {
|
||||
encode_field_bytes(2, &content, &mut buf);
|
||||
}
|
||||
buf
|
||||
}
|
||||
|
||||
fn decode_msg_body_element(data: &[u8]) -> Result<MsgBodyElement, YuanbaoError> {
|
||||
let fields = parse_fields(data)?;
|
||||
let content_bytes = get_bytes(&fields, 2);
|
||||
let content = if content_bytes.is_empty() {
|
||||
MsgContent::default()
|
||||
} else {
|
||||
decode_msg_content(&content_bytes)?
|
||||
};
|
||||
Ok(MsgBodyElement {
|
||||
msg_type: get_string(&fields, 1),
|
||||
msg_content: content,
|
||||
})
|
||||
}
|
||||
|
||||
// ─── PushMsg envelope (cmd_type=Push inner wrapper) ────────────────
|
||||
//
|
||||
// message PushMsg {
|
||||
// string cmd = 1;
|
||||
// string module = 2;
|
||||
// string msg_id = 3;
|
||||
// bytes data = 4; // ← actual biz body (e.g. InboundMessagePush)
|
||||
// }
|
||||
//
|
||||
// The yuanbao gateway wraps every downstream push in this envelope
|
||||
// *inside* `ConnMsg.data`. Mirrors plugin client.ts::onPush which
|
||||
// decodes PushMsg before handing `data` to the business decoder.
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct PushMsg {
|
||||
pub cmd: String,
|
||||
pub module: String,
|
||||
pub msg_id: String,
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
pub fn decode_push_msg(data: &[u8]) -> Result<PushMsg, YuanbaoError> {
|
||||
let fields = parse_fields(data)?;
|
||||
Ok(PushMsg {
|
||||
cmd: get_string(&fields, 1),
|
||||
module: get_string(&fields, 2),
|
||||
msg_id: get_string(&fields, 3),
|
||||
data: get_bytes(&fields, 4),
|
||||
})
|
||||
}
|
||||
|
||||
// ─── InboundMessagePush decode ─────────────────────────────────────
|
||||
|
||||
pub fn decode_inbound_push(data: &[u8]) -> Result<InboundMessage, YuanbaoError> {
|
||||
let fields = parse_fields(data)?;
|
||||
|
||||
let mut msg_body = Vec::new();
|
||||
for b in get_repeated_bytes(&fields, 13) {
|
||||
msg_body.push(decode_msg_body_element(&b)?);
|
||||
}
|
||||
|
||||
let mut recalls = Vec::new();
|
||||
for b in get_repeated_bytes(&fields, 17) {
|
||||
let f = parse_fields(&b)?;
|
||||
recalls.push(ImMsgSeq {
|
||||
msg_seq: get_varint(&f, 1) as u32,
|
||||
msg_id: get_string(&f, 2),
|
||||
});
|
||||
}
|
||||
|
||||
let log_ext_bytes = get_bytes(&fields, 20);
|
||||
let trace_id = if log_ext_bytes.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
get_string(&parse_fields(&log_ext_bytes)?, 1)
|
||||
};
|
||||
|
||||
Ok(InboundMessage {
|
||||
callback_command: get_string(&fields, 1),
|
||||
from_account: get_string(&fields, 2),
|
||||
to_account: get_string(&fields, 3),
|
||||
sender_nickname: get_string(&fields, 4),
|
||||
group_id: get_string(&fields, 5),
|
||||
group_code: get_string(&fields, 6),
|
||||
group_name: get_string(&fields, 7),
|
||||
msg_seq: get_varint(&fields, 8) as u32,
|
||||
msg_random: get_varint(&fields, 9) as u32,
|
||||
msg_time: get_varint(&fields, 10) as u32,
|
||||
msg_key: get_string(&fields, 11),
|
||||
msg_id: get_string(&fields, 12),
|
||||
msg_body,
|
||||
cloud_custom_data: get_string(&fields, 14),
|
||||
event_time: get_varint(&fields, 15) as u32,
|
||||
bot_owner_id: get_string(&fields, 16),
|
||||
recall_msg_seq_list: recalls,
|
||||
claw_msg_type: get_varint(&fields, 18) as u32,
|
||||
private_from_group_code: get_string(&fields, 19),
|
||||
trace_id,
|
||||
})
|
||||
}
|
||||
|
||||
// ─── InboundMessagePush JSON decode ────────────────────────────────
|
||||
//
|
||||
// The yuanbao gateway sometimes (depending on backend account config /
|
||||
// source channel) pushes `inbound_message` as a JSON string instead of
|
||||
// protobuf. The shape matches `InboundMessagePush` field-for-field
|
||||
// (snake_case), with `log_ext.trace_id` nested. Mirrors plugin
|
||||
// gateway.ts::decodeFromRawDataJson (l. 238).
|
||||
|
||||
pub fn decode_inbound_json(data: &[u8]) -> Result<InboundMessage, YuanbaoError> {
|
||||
let v: serde_json::Value = serde_json::from_slice(data)
|
||||
.map_err(|e| YuanbaoError::ProtoDecode(format!("json parse failed: {e}")))?;
|
||||
|
||||
let obj = v
|
||||
.as_object()
|
||||
.ok_or_else(|| YuanbaoError::ProtoDecode("json root is not an object".into()))?;
|
||||
|
||||
let get_str = |k: &str| -> String {
|
||||
obj.get(k)
|
||||
.and_then(|x| x.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string()
|
||||
};
|
||||
let get_u32 = |k: &str| -> u32 { obj.get(k).and_then(|x| x.as_u64()).unwrap_or(0) as u32 };
|
||||
|
||||
let msg_body = obj
|
||||
.get("msg_body")
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| arr.iter().map(decode_msg_body_element_json).collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
let recall_msg_seq_list = obj
|
||||
.get("recall_msg_seq_list")
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.map(|e| ImMsgSeq {
|
||||
msg_seq: e.get("msg_seq").and_then(|v| v.as_u64()).unwrap_or(0) as u32,
|
||||
msg_id: e
|
||||
.get("msg_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let trace_id = obj
|
||||
.get("log_ext")
|
||||
.and_then(|v| v.get("trace_id"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
|
||||
Ok(InboundMessage {
|
||||
callback_command: get_str("callback_command"),
|
||||
from_account: get_str("from_account"),
|
||||
to_account: get_str("to_account"),
|
||||
sender_nickname: get_str("sender_nickname"),
|
||||
group_id: get_str("group_id"),
|
||||
group_code: get_str("group_code"),
|
||||
group_name: get_str("group_name"),
|
||||
msg_seq: get_u32("msg_seq"),
|
||||
msg_random: get_u32("msg_random"),
|
||||
msg_time: get_u32("msg_time"),
|
||||
msg_key: get_str("msg_key"),
|
||||
msg_id: get_str("msg_id"),
|
||||
msg_body,
|
||||
cloud_custom_data: get_str("cloud_custom_data"),
|
||||
event_time: get_u32("event_time"),
|
||||
bot_owner_id: get_str("bot_owner_id"),
|
||||
recall_msg_seq_list,
|
||||
claw_msg_type: get_u32("claw_msg_type"),
|
||||
private_from_group_code: get_str("private_from_group_code"),
|
||||
trace_id,
|
||||
})
|
||||
}
|
||||
|
||||
fn decode_msg_body_element_json(v: &serde_json::Value) -> MsgBodyElement {
|
||||
let msg_type = v
|
||||
.get("msg_type")
|
||||
.and_then(|x| x.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
let mc = v.get("msg_content").and_then(|x| x.as_object());
|
||||
|
||||
let str_field = |k: &str| -> Option<String> {
|
||||
mc.and_then(|m| m.get(k))
|
||||
.and_then(|x| x.as_str())
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(|s| s.to_string())
|
||||
};
|
||||
let u32_field = |k: &str| -> Option<u32> {
|
||||
mc.and_then(|m| m.get(k))
|
||||
.and_then(|x| x.as_u64())
|
||||
.map(|n| n as u32)
|
||||
};
|
||||
|
||||
let image_info_array = mc
|
||||
.and_then(|m| m.get("image_info_array"))
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.map(|e| ImageInfo {
|
||||
image_type: e
|
||||
.get("type")
|
||||
.or_else(|| e.get("image_type"))
|
||||
.and_then(|x| x.as_u64())
|
||||
.unwrap_or(0) as u32,
|
||||
size: e.get("size").and_then(|x| x.as_u64()).unwrap_or(0) as u32,
|
||||
width: e.get("width").and_then(|x| x.as_u64()).unwrap_or(0) as u32,
|
||||
height: e.get("height").and_then(|x| x.as_u64()).unwrap_or(0) as u32,
|
||||
url: e
|
||||
.get("url")
|
||||
.and_then(|x| x.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
MsgBodyElement {
|
||||
msg_type,
|
||||
msg_content: MsgContent {
|
||||
text: str_field("text"),
|
||||
uuid: str_field("uuid"),
|
||||
image_format: u32_field("image_format"),
|
||||
data: str_field("data"),
|
||||
desc: str_field("desc"),
|
||||
ext: str_field("ext"),
|
||||
sound: str_field("sound"),
|
||||
image_info_array,
|
||||
index: u32_field("index"),
|
||||
url: str_field("url"),
|
||||
file_size: u32_field("file_size"),
|
||||
file_name: str_field("file_name"),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn conn_msg_roundtrip() {
|
||||
let buf = encode_conn_msg(
|
||||
cmd_type::REQUEST,
|
||||
cmd::PING,
|
||||
42,
|
||||
"mid-1",
|
||||
module::CONN_ACCESS,
|
||||
b"payload",
|
||||
);
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd_type, cmd_type::REQUEST);
|
||||
assert_eq!(frame.cmd, cmd::PING);
|
||||
assert_eq!(frame.seq_no, 42);
|
||||
assert_eq!(frame.msg_id, "mid-1");
|
||||
assert_eq!(frame.module, module::CONN_ACCESS);
|
||||
assert_eq!(frame.data, b"payload");
|
||||
}
|
||||
|
||||
/// Smoke-test that [`encode_auth_bind`] produces a frame round-trippable
|
||||
/// via [`decode_conn_msg`] and that the `app_version` / `bot_version`
|
||||
/// arguments land in the expected `DeviceInfo` fields (regression guard
|
||||
/// for the plugin_version/bot_version swap).
|
||||
#[test]
|
||||
fn auth_bind_smoke() {
|
||||
let buf = encode_auth_bind(
|
||||
"biz", "uid", "openclaw", "tok", "mid", "0.1.0", "linux", "1.0", "",
|
||||
);
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd, cmd::AUTH_BIND);
|
||||
assert_eq!(frame.module, module::CONN_ACCESS);
|
||||
assert!(!frame.data.is_empty());
|
||||
|
||||
let req_fields = parse_fields(&frame.data).unwrap();
|
||||
let dev_buf = get_bytes(&req_fields, 3);
|
||||
let dev_fields = parse_fields(&dev_buf).unwrap();
|
||||
assert_eq!(get_string(&dev_fields, 1), "0.1.0", "app_version");
|
||||
assert_eq!(get_string(&dev_fields, 24), "1.0", "bot_version");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn push_ack_mirrors_original() {
|
||||
let original = ConnFrame {
|
||||
cmd_type: cmd_type::PUSH,
|
||||
cmd: "some_push".into(),
|
||||
module: "yuanbao_openclaw_proxy".into(),
|
||||
seq_no: 99,
|
||||
msg_id: "mid-abc".into(),
|
||||
need_ack: true,
|
||||
status: 0,
|
||||
data: vec![],
|
||||
};
|
||||
let buf = encode_push_ack(&original);
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd_type, cmd_type::PUSH_ACK);
|
||||
assert_eq!(frame.cmd, original.cmd);
|
||||
assert_eq!(frame.module, original.module);
|
||||
assert_eq!(frame.msg_id, original.msg_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn msg_body_element_roundtrip() {
|
||||
let el = MsgBodyElement {
|
||||
msg_type: "TIMTextElem".into(),
|
||||
msg_content: MsgContent {
|
||||
text: Some("hello 元宝".into()),
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
let buf = encode_msg_body_element(&el);
|
||||
let got = decode_msg_body_element(&buf).unwrap();
|
||||
assert_eq!(got, el);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_element_roundtrip() {
|
||||
let el = MsgBodyElement {
|
||||
msg_type: "TIMImageElem".into(),
|
||||
msg_content: MsgContent {
|
||||
uuid: Some("abc123".into()),
|
||||
image_format: Some(3),
|
||||
image_info_array: vec![ImageInfo {
|
||||
image_type: 1,
|
||||
size: 1024,
|
||||
width: 800,
|
||||
height: 600,
|
||||
url: "https://example/img.png".into(),
|
||||
}],
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
let buf = encode_msg_body_element(&el);
|
||||
let got = decode_msg_body_element(&buf).unwrap();
|
||||
assert_eq!(got, el);
|
||||
}
|
||||
|
||||
// ─── decode_auth_bind_rsp ─────────────────────────────────────
|
||||
|
||||
fn build_auth_bind_rsp_bytes(code: u64, message: &str, connect_id: &str) -> Vec<u8> {
|
||||
let mut buf = Vec::new();
|
||||
if code != 0 {
|
||||
encode_field_varint(1, code, &mut buf);
|
||||
}
|
||||
if !message.is_empty() {
|
||||
encode_field_string(2, message, &mut buf);
|
||||
}
|
||||
if !connect_id.is_empty() {
|
||||
encode_field_string(3, connect_id, &mut buf);
|
||||
}
|
||||
buf
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_auth_bind_rsp_happy_path() {
|
||||
let body = build_auth_bind_rsp_bytes(0, "ok", "conn-42");
|
||||
let r = decode_auth_bind_rsp(&body).unwrap();
|
||||
assert_eq!(r.code, 0);
|
||||
assert_eq!(r.message, "ok");
|
||||
assert_eq!(r.connect_id, "conn-42");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_auth_bind_rsp_with_error_code() {
|
||||
let body = build_auth_bind_rsp_bytes(40011, "rejected", "");
|
||||
let r = decode_auth_bind_rsp(&body).unwrap();
|
||||
assert_eq!(r.code, 40011);
|
||||
assert_eq!(r.message, "rejected");
|
||||
assert!(r.connect_id.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_auth_bind_rsp_on_empty_returns_default() {
|
||||
let r = decode_auth_bind_rsp(&[]).unwrap();
|
||||
assert_eq!(r, AuthBindRsp::default());
|
||||
}
|
||||
|
||||
// ─── decode_push_msg ──────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn decode_push_msg_extracts_all_fields() {
|
||||
let inner_payload = vec![0xCA, 0xFE, 0xBA, 0xBE];
|
||||
let mut buf = Vec::new();
|
||||
encode_field_string(1, "inbound_message", &mut buf);
|
||||
encode_field_string(2, "yuanbao_openclaw_proxy", &mut buf);
|
||||
encode_field_string(3, "pm-1", &mut buf);
|
||||
encode_field_bytes(4, &inner_payload, &mut buf);
|
||||
|
||||
let pm = decode_push_msg(&buf).unwrap();
|
||||
assert_eq!(pm.cmd, "inbound_message");
|
||||
assert_eq!(pm.module, "yuanbao_openclaw_proxy");
|
||||
assert_eq!(pm.msg_id, "pm-1");
|
||||
assert_eq!(pm.data, inner_payload);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_push_msg_on_empty_returns_defaults() {
|
||||
let pm = decode_push_msg(&[]).unwrap();
|
||||
assert!(pm.cmd.is_empty());
|
||||
assert!(pm.module.is_empty());
|
||||
assert!(pm.msg_id.is_empty());
|
||||
assert!(pm.data.is_empty());
|
||||
}
|
||||
|
||||
// ─── decode_inbound_push (protobuf) ───────────────────────────
|
||||
|
||||
#[test]
|
||||
fn decode_inbound_push_dm_with_text_body() {
|
||||
// Build a minimal DM push: from/to/sender_nickname + one TIMTextElem.
|
||||
let text_elem = MsgBodyElement {
|
||||
msg_type: "TIMTextElem".into(),
|
||||
msg_content: MsgContent {
|
||||
text: Some("hello".into()),
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
let elem_bytes = encode_msg_body_element(&text_elem);
|
||||
|
||||
let mut log_ext = Vec::new();
|
||||
encode_field_string(1, "trace-123", &mut log_ext);
|
||||
|
||||
let mut buf = Vec::new();
|
||||
encode_field_string(1, "C2CMsg", &mut buf);
|
||||
encode_field_string(2, "user_42", &mut buf);
|
||||
encode_field_string(3, "bot_1", &mut buf);
|
||||
encode_field_string(4, "Alice", &mut buf);
|
||||
encode_field_varint(8, 7, &mut buf);
|
||||
encode_field_varint(9, 123, &mut buf);
|
||||
encode_field_varint(10, 1_700_000_000, &mut buf);
|
||||
encode_field_string(12, "mid-abc", &mut buf);
|
||||
encode_field_bytes(13, &elem_bytes, &mut buf);
|
||||
encode_field_varint(15, 1_700_000_001, &mut buf);
|
||||
encode_field_bytes(20, &log_ext, &mut buf);
|
||||
|
||||
let m = decode_inbound_push(&buf).unwrap();
|
||||
assert_eq!(m.callback_command, "C2CMsg");
|
||||
assert_eq!(m.from_account, "user_42");
|
||||
assert_eq!(m.to_account, "bot_1");
|
||||
assert_eq!(m.sender_nickname, "Alice");
|
||||
assert_eq!(m.msg_seq, 7);
|
||||
assert_eq!(m.msg_random, 123);
|
||||
assert_eq!(m.msg_time, 1_700_000_000);
|
||||
assert_eq!(m.msg_id, "mid-abc");
|
||||
assert_eq!(m.event_time, 1_700_000_001);
|
||||
assert_eq!(m.trace_id, "trace-123");
|
||||
assert_eq!(m.msg_body.len(), 1);
|
||||
assert_eq!(m.msg_body[0].msg_content.text.as_deref(), Some("hello"));
|
||||
assert!(m.recall_msg_seq_list.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_inbound_push_group_with_recall_list() {
|
||||
let mut recall_entry = Vec::new();
|
||||
encode_field_varint(1, 99, &mut recall_entry);
|
||||
encode_field_string(2, "old-msg-id", &mut recall_entry);
|
||||
|
||||
let mut buf = Vec::new();
|
||||
encode_field_string(1, "GroupSysMsg", &mut buf);
|
||||
encode_field_string(5, "gid-x", &mut buf);
|
||||
encode_field_string(6, "gcode-y", &mut buf);
|
||||
encode_field_string(7, "Room", &mut buf);
|
||||
encode_field_bytes(17, &recall_entry, &mut buf);
|
||||
encode_field_string(19, "g-private-code", &mut buf);
|
||||
|
||||
let m = decode_inbound_push(&buf).unwrap();
|
||||
assert_eq!(m.callback_command, "GroupSysMsg");
|
||||
assert_eq!(m.group_id, "gid-x");
|
||||
assert_eq!(m.group_code, "gcode-y");
|
||||
assert_eq!(m.group_name, "Room");
|
||||
assert_eq!(m.private_from_group_code, "g-private-code");
|
||||
assert_eq!(m.recall_msg_seq_list.len(), 1);
|
||||
assert_eq!(m.recall_msg_seq_list[0].msg_seq, 99);
|
||||
assert_eq!(m.recall_msg_seq_list[0].msg_id, "old-msg-id");
|
||||
assert!(m.trace_id.is_empty(), "no log_ext => empty trace_id");
|
||||
}
|
||||
|
||||
// ─── decode_inbound_json ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn decode_inbound_json_full_dm_shape() {
|
||||
let json = serde_json::json!({
|
||||
"callback_command": "C2CMsg",
|
||||
"from_account": "user_42",
|
||||
"to_account": "bot_1",
|
||||
"sender_nickname": "Alice",
|
||||
"msg_seq": 7,
|
||||
"msg_random": 123,
|
||||
"msg_time": 1_700_000_000,
|
||||
"msg_id": "mid-1",
|
||||
"msg_body": [
|
||||
{
|
||||
"msg_type": "TIMTextElem",
|
||||
"msg_content": { "text": "hi" }
|
||||
},
|
||||
{
|
||||
"msg_type": "TIMImageElem",
|
||||
"msg_content": {
|
||||
"uuid": "u-1",
|
||||
"image_format": 1,
|
||||
"image_info_array": [
|
||||
{ "type": 1, "size": 100, "width": 10, "height": 20, "url": "https://x/i.png" }
|
||||
]
|
||||
}
|
||||
}
|
||||
],
|
||||
"recall_msg_seq_list": [{ "msg_seq": 9, "msg_id": "old" }],
|
||||
"log_ext": { "trace_id": "trace-json" }
|
||||
});
|
||||
let m = decode_inbound_json(json.to_string().as_bytes()).unwrap();
|
||||
assert_eq!(m.callback_command, "C2CMsg");
|
||||
assert_eq!(m.from_account, "user_42");
|
||||
assert_eq!(m.msg_id, "mid-1");
|
||||
assert_eq!(m.msg_body.len(), 2);
|
||||
assert_eq!(m.msg_body[0].msg_content.text.as_deref(), Some("hi"));
|
||||
let img = &m.msg_body[1].msg_content;
|
||||
assert_eq!(img.uuid.as_deref(), Some("u-1"));
|
||||
assert_eq!(img.image_info_array.len(), 1);
|
||||
assert_eq!(img.image_info_array[0].url, "https://x/i.png");
|
||||
assert_eq!(m.recall_msg_seq_list.len(), 1);
|
||||
assert_eq!(m.recall_msg_seq_list[0].msg_seq, 9);
|
||||
assert_eq!(m.trace_id, "trace-json");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_inbound_json_rejects_non_object_root() {
|
||||
let err = decode_inbound_json(b"[1,2,3]").unwrap_err();
|
||||
match err {
|
||||
YuanbaoError::ProtoDecode(m) => assert!(m.contains("not an object"), "got {m}"),
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_inbound_json_rejects_invalid_json() {
|
||||
let err = decode_inbound_json(b"not json").unwrap_err();
|
||||
match err {
|
||||
YuanbaoError::ProtoDecode(m) => assert!(m.contains("json parse"), "got {m}"),
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_msg_body_element_json_handles_image_type_alias() {
|
||||
// Some payloads use `image_type` (snake_case) instead of `type`.
|
||||
let v = serde_json::json!({
|
||||
"msg_type": "TIMImageElem",
|
||||
"msg_content": {
|
||||
"image_info_array": [
|
||||
{ "image_type": 2, "size": 50, "width": 5, "height": 6, "url": "u" }
|
||||
]
|
||||
}
|
||||
});
|
||||
let el = decode_msg_body_element_json(&v);
|
||||
assert_eq!(el.msg_type, "TIMImageElem");
|
||||
assert_eq!(el.msg_content.image_info_array.len(), 1);
|
||||
assert_eq!(el.msg_content.image_info_array[0].image_type, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_msg_content_image_info_with_only_image_type_zero_defaults_to_one() {
|
||||
// When `image_type` is 0 but url is present, decoder bumps to 1.
|
||||
let mut ib = Vec::new();
|
||||
encode_field_varint(2, 64, &mut ib);
|
||||
encode_field_string(5, "https://x/y.png", &mut ib);
|
||||
let mut content = Vec::new();
|
||||
encode_field_bytes(8, &ib, &mut content);
|
||||
let mut elem = Vec::new();
|
||||
encode_field_string(1, "TIMImageElem", &mut elem);
|
||||
encode_field_bytes(2, &content, &mut elem);
|
||||
let got = decode_msg_body_element(&elem).unwrap();
|
||||
assert_eq!(got.msg_content.image_info_array.len(), 1);
|
||||
assert_eq!(got.msg_content.image_info_array[0].image_type, 1);
|
||||
assert_eq!(got.msg_content.image_info_array[0].url, "https://x/y.png");
|
||||
}
|
||||
}
|
||||
@@ -1,638 +0,0 @@
|
||||
//! Business-layer protobuf codecs (biz payloads inside `ConnMsg.data`).
|
||||
//!
|
||||
//! Kept separate from `proto.rs` to stay under the 500-line ceiling and
|
||||
//! to isolate the "openclaw biz protocol" surface from the lower-level
|
||||
//! ConnMsg envelope.
|
||||
|
||||
use super::errors::YuanbaoError;
|
||||
use super::proto::{decode_conn_msg, encode_conn_msg, encode_msg_body_element};
|
||||
use super::proto_constants::*;
|
||||
use super::types::*;
|
||||
use super::wire::{
|
||||
encode_field_bytes as put_bytes_field, encode_field_string as put_string_field,
|
||||
encode_field_varint as put_varint_field, get_bytes, get_repeated_bytes, get_string, get_varint,
|
||||
next_seq_no, parse_fields,
|
||||
};
|
||||
|
||||
// ─── SendC2CMessageReq ────────────────────────────────────────────
|
||||
//
|
||||
// 1: msg_id (string) 5: msg_body (repeated MsgBodyElement)
|
||||
// 2: to_account 6: group_code (DM-from-group)
|
||||
// 3: from_account 7: msg_seq
|
||||
// 4: msg_random 8: log_ext
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn encode_send_c2c_req(
|
||||
msg_id: &str,
|
||||
to_account: &str,
|
||||
from_account: &str,
|
||||
msg_random: u32,
|
||||
msg_body: &[MsgBodyElement],
|
||||
group_code: &str,
|
||||
msg_seq: Option<u64>,
|
||||
trace_id: &str,
|
||||
) -> Vec<u8> {
|
||||
let mut buf = Vec::with_capacity(128);
|
||||
if !msg_id.is_empty() {
|
||||
put_string_field(1, msg_id, &mut buf);
|
||||
}
|
||||
put_string_field(2, to_account, &mut buf);
|
||||
if !from_account.is_empty() {
|
||||
put_string_field(3, from_account, &mut buf);
|
||||
}
|
||||
if msg_random != 0 {
|
||||
put_varint_field(4, msg_random as u64, &mut buf);
|
||||
}
|
||||
for el in msg_body {
|
||||
let el_bytes = encode_msg_body_element(el);
|
||||
put_bytes_field(5, &el_bytes, &mut buf);
|
||||
}
|
||||
if !group_code.is_empty() {
|
||||
put_string_field(6, group_code, &mut buf);
|
||||
}
|
||||
if let Some(seq) = msg_seq {
|
||||
put_varint_field(7, seq, &mut buf);
|
||||
}
|
||||
if !trace_id.is_empty() {
|
||||
// log_ext is field 8 with a nested {1: trace_id}
|
||||
let mut log = Vec::new();
|
||||
put_string_field(1, trace_id, &mut log);
|
||||
put_bytes_field(8, &log, &mut buf);
|
||||
}
|
||||
buf
|
||||
}
|
||||
|
||||
/// Encode a full C2C send request as a `ConnMsg` ready to send over WS.
|
||||
pub fn encode_send_c2c_message(
|
||||
to_account: &str,
|
||||
from_account: &str,
|
||||
msg_body: &[MsgBodyElement],
|
||||
msg_id: &str,
|
||||
msg_random: u32,
|
||||
group_code: &str,
|
||||
trace_id: &str,
|
||||
) -> Vec<u8> {
|
||||
let body = encode_send_c2c_req(
|
||||
msg_id,
|
||||
to_account,
|
||||
from_account,
|
||||
msg_random,
|
||||
msg_body,
|
||||
group_code,
|
||||
None,
|
||||
trace_id,
|
||||
);
|
||||
let req_id = if msg_id.is_empty() {
|
||||
format!("c2c_{}", next_seq_no())
|
||||
} else {
|
||||
msg_id.to_string()
|
||||
};
|
||||
encode_conn_msg(
|
||||
cmd_type::REQUEST,
|
||||
biz_cmd::SEND_C2C_MESSAGE,
|
||||
next_seq_no(),
|
||||
&req_id,
|
||||
module::BIZ_PKG,
|
||||
&body,
|
||||
)
|
||||
}
|
||||
|
||||
// ─── SendGroupMessageReq ───────────────────────────────────────────
|
||||
//
|
||||
// 1: msg_id 5: random (string)
|
||||
// 2: group_code 6: msg_body (repeated)
|
||||
// 3: from_account 7: ref_msg_id
|
||||
// 4: to_account 8: msg_seq
|
||||
// 9: log_ext
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn encode_send_group_req(
|
||||
msg_id: &str,
|
||||
group_code: &str,
|
||||
from_account: &str,
|
||||
to_account: &str,
|
||||
random: &str,
|
||||
msg_body: &[MsgBodyElement],
|
||||
ref_msg_id: &str,
|
||||
trace_id: &str,
|
||||
) -> Vec<u8> {
|
||||
let mut buf = Vec::with_capacity(128);
|
||||
if !msg_id.is_empty() {
|
||||
put_string_field(1, msg_id, &mut buf);
|
||||
}
|
||||
put_string_field(2, group_code, &mut buf);
|
||||
if !from_account.is_empty() {
|
||||
put_string_field(3, from_account, &mut buf);
|
||||
}
|
||||
if !to_account.is_empty() {
|
||||
put_string_field(4, to_account, &mut buf);
|
||||
}
|
||||
if !random.is_empty() {
|
||||
put_string_field(5, random, &mut buf);
|
||||
}
|
||||
for el in msg_body {
|
||||
let el_bytes = encode_msg_body_element(el);
|
||||
put_bytes_field(6, &el_bytes, &mut buf);
|
||||
}
|
||||
if !ref_msg_id.is_empty() {
|
||||
put_string_field(7, ref_msg_id, &mut buf);
|
||||
}
|
||||
if !trace_id.is_empty() {
|
||||
let mut log = Vec::new();
|
||||
put_string_field(1, trace_id, &mut log);
|
||||
put_bytes_field(9, &log, &mut buf);
|
||||
}
|
||||
buf
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn encode_send_group_message(
|
||||
group_code: &str,
|
||||
from_account: &str,
|
||||
msg_body: &[MsgBodyElement],
|
||||
msg_id: &str,
|
||||
to_account: &str,
|
||||
random: &str,
|
||||
ref_msg_id: &str,
|
||||
trace_id: &str,
|
||||
) -> Vec<u8> {
|
||||
let body = encode_send_group_req(
|
||||
msg_id,
|
||||
group_code,
|
||||
from_account,
|
||||
to_account,
|
||||
random,
|
||||
msg_body,
|
||||
ref_msg_id,
|
||||
trace_id,
|
||||
);
|
||||
let req_id = if msg_id.is_empty() {
|
||||
format!("grp_{}", next_seq_no())
|
||||
} else {
|
||||
msg_id.to_string()
|
||||
};
|
||||
encode_conn_msg(
|
||||
cmd_type::REQUEST,
|
||||
biz_cmd::SEND_GROUP_MESSAGE,
|
||||
next_seq_no(),
|
||||
&req_id,
|
||||
module::BIZ_PKG,
|
||||
&body,
|
||||
)
|
||||
}
|
||||
|
||||
// ─── Heartbeats ────────────────────────────────────────────────────
|
||||
|
||||
pub fn encode_send_private_heartbeat(
|
||||
req_id: &str,
|
||||
from_account: &str,
|
||||
to_account: &str,
|
||||
heartbeat: u32,
|
||||
) -> Vec<u8> {
|
||||
let mut body = Vec::with_capacity(48);
|
||||
put_string_field(1, from_account, &mut body);
|
||||
put_string_field(2, to_account, &mut body);
|
||||
put_varint_field(3, heartbeat as u64, &mut body);
|
||||
encode_conn_msg(
|
||||
cmd_type::REQUEST,
|
||||
biz_cmd::SEND_PRIVATE_HEARTBEAT,
|
||||
next_seq_no(),
|
||||
req_id,
|
||||
module::BIZ_PKG,
|
||||
&body,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn encode_send_group_heartbeat(
|
||||
req_id: &str,
|
||||
from_account: &str,
|
||||
group_code: &str,
|
||||
heartbeat: u32,
|
||||
send_time_ms: u64,
|
||||
) -> Vec<u8> {
|
||||
let mut body = Vec::with_capacity(64);
|
||||
put_string_field(1, from_account, &mut body);
|
||||
put_string_field(2, "", &mut body); // to_account empty for group
|
||||
put_string_field(3, group_code, &mut body);
|
||||
put_varint_field(4, send_time_ms, &mut body);
|
||||
put_varint_field(5, heartbeat as u64, &mut body);
|
||||
encode_conn_msg(
|
||||
cmd_type::REQUEST,
|
||||
biz_cmd::SEND_GROUP_HEARTBEAT,
|
||||
next_seq_no(),
|
||||
req_id,
|
||||
module::BIZ_PKG,
|
||||
&body,
|
||||
)
|
||||
}
|
||||
|
||||
// ─── QueryGroupInfo ────────────────────────────────────────────────
|
||||
|
||||
pub fn encode_query_group_info(req_id: &str, group_code: &str) -> Vec<u8> {
|
||||
let mut body = Vec::with_capacity(16 + group_code.len());
|
||||
put_string_field(1, group_code, &mut body);
|
||||
encode_conn_msg(
|
||||
cmd_type::REQUEST,
|
||||
biz_cmd::QUERY_GROUP_INFO,
|
||||
next_seq_no(),
|
||||
req_id,
|
||||
module::BIZ_PKG,
|
||||
&body,
|
||||
)
|
||||
}
|
||||
|
||||
/// Try to narrow a varint into a smaller integer type, returning
|
||||
/// `YuanbaoError::ProtoDecode` (instead of silently truncating) when
|
||||
/// the upstream value is out of range. Used to harden response decoders
|
||||
/// against malformed / adversarial input.
|
||||
fn varint_to_i32(value: u64, field_label: &str) -> Result<i32, YuanbaoError> {
|
||||
i32::try_from(value)
|
||||
.map_err(|_| YuanbaoError::ProtoDecode(format!("{field_label} out of i32 range: {value}")))
|
||||
}
|
||||
|
||||
fn varint_to_u32(value: u64, field_label: &str) -> Result<u32, YuanbaoError> {
|
||||
u32::try_from(value)
|
||||
.map_err(|_| YuanbaoError::ProtoDecode(format!("{field_label} out of u32 range: {value}")))
|
||||
}
|
||||
|
||||
pub fn decode_query_group_info_rsp(data: &[u8]) -> Result<GroupInfo, YuanbaoError> {
|
||||
let fields = parse_fields(data)?;
|
||||
let mut info = GroupInfo {
|
||||
code: varint_to_i32(get_varint(&fields, 1), "GroupInfoRsp.code")?,
|
||||
message: get_string(&fields, 2),
|
||||
..Default::default()
|
||||
};
|
||||
let gi_bytes = get_bytes(&fields, 3);
|
||||
if !gi_bytes.is_empty() {
|
||||
let gi = parse_fields(&gi_bytes)?;
|
||||
info.group_name = get_string(&gi, 1);
|
||||
info.owner_id = get_string(&gi, 2);
|
||||
info.owner_nickname = get_string(&gi, 3);
|
||||
info.member_count = varint_to_u32(get_varint(&gi, 4), "GroupInfo.member_count")?;
|
||||
}
|
||||
Ok(info)
|
||||
}
|
||||
|
||||
// ─── GetGroupMemberList ────────────────────────────────────────────
|
||||
|
||||
pub fn encode_get_group_member_list(
|
||||
req_id: &str,
|
||||
group_code: &str,
|
||||
offset: u32,
|
||||
limit: u32,
|
||||
) -> Vec<u8> {
|
||||
let mut body = Vec::with_capacity(32 + group_code.len());
|
||||
put_string_field(1, group_code, &mut body);
|
||||
if offset != 0 {
|
||||
put_varint_field(2, offset as u64, &mut body);
|
||||
}
|
||||
put_varint_field(3, limit as u64, &mut body);
|
||||
encode_conn_msg(
|
||||
cmd_type::REQUEST,
|
||||
biz_cmd::GET_GROUP_MEMBER_LIST,
|
||||
next_seq_no(),
|
||||
req_id,
|
||||
module::BIZ_PKG,
|
||||
&body,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn decode_get_group_member_list_rsp(data: &[u8]) -> Result<GroupMemberListPage, YuanbaoError> {
|
||||
let fields = parse_fields(data)?;
|
||||
let mut members = Vec::new();
|
||||
for b in get_repeated_bytes(&fields, 3) {
|
||||
let m = parse_fields(&b)?;
|
||||
members.push(GroupMember {
|
||||
user_id: get_string(&m, 1),
|
||||
nickname: get_string(&m, 2),
|
||||
role: varint_to_u32(get_varint(&m, 3), "GroupMember.role")?,
|
||||
join_time: varint_to_u32(get_varint(&m, 4), "GroupMember.join_time")?,
|
||||
name_card: get_string(&m, 5),
|
||||
});
|
||||
}
|
||||
Ok(GroupMemberListPage {
|
||||
code: varint_to_i32(get_varint(&fields, 1), "GroupMemberListRsp.code")?,
|
||||
message: get_string(&fields, 2),
|
||||
members,
|
||||
next_offset: varint_to_u32(get_varint(&fields, 4), "GroupMemberListRsp.next_offset")?,
|
||||
is_complete: get_varint(&fields, 5) != 0,
|
||||
})
|
||||
}
|
||||
|
||||
// ─── Generic biz response code helper ──────────────────────────────
|
||||
|
||||
/// Decode the `code` and `message` from a biz response.
|
||||
///
|
||||
/// All biz responses share the convention: field 1 = code, field 2 = message.
|
||||
pub fn decode_biz_rsp_code(data: &[u8]) -> Result<(i32, String), YuanbaoError> {
|
||||
let fields = parse_fields(data)?;
|
||||
Ok((
|
||||
varint_to_i32(get_varint(&fields, 1), "BizRsp.code")?,
|
||||
get_string(&fields, 2),
|
||||
))
|
||||
}
|
||||
|
||||
/// Decode a `ConnMsg` and return the typed biz response code + frame for
|
||||
/// the request/response correlator.
|
||||
pub fn decode_response_envelope(frame_bytes: &[u8]) -> Result<ConnFrame, YuanbaoError> {
|
||||
decode_conn_msg(frame_bytes)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn text_body(s: &str) -> Vec<MsgBodyElement> {
|
||||
vec![MsgBodyElement {
|
||||
msg_type: "TIMTextElem".into(),
|
||||
msg_content: MsgContent {
|
||||
text: Some(s.into()),
|
||||
..Default::default()
|
||||
},
|
||||
}]
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn c2c_encode_smoke() {
|
||||
let buf = encode_send_c2c_message("uid_alice", "uid_bot", &text_body("hi"), "", 0, "", "");
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd, biz_cmd::SEND_C2C_MESSAGE);
|
||||
assert_eq!(frame.module, module::BIZ_PKG);
|
||||
assert!(!frame.data.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn group_encode_smoke() {
|
||||
let buf = encode_send_group_message(
|
||||
"group_42",
|
||||
"uid_bot",
|
||||
&text_body("hello"),
|
||||
"",
|
||||
"",
|
||||
"rand",
|
||||
"",
|
||||
"",
|
||||
);
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd, biz_cmd::SEND_GROUP_MESSAGE);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn private_heartbeat_smoke() {
|
||||
let buf =
|
||||
encode_send_private_heartbeat("hb_1", "uid_bot", "uid_user", ws_heartbeat::RUNNING);
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd, biz_cmd::SEND_PRIVATE_HEARTBEAT);
|
||||
assert_eq!(frame.msg_id, "hb_1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn query_group_info_roundtrip() {
|
||||
let buf = encode_query_group_info("qgi_1", "group_99");
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd, biz_cmd::QUERY_GROUP_INFO);
|
||||
assert_eq!(frame.msg_id, "qgi_1");
|
||||
|
||||
// Simulate response payload: code=0, message="ok", group_name="g", owner=…
|
||||
let mut gi = Vec::new();
|
||||
put_string_field(1, "TestGroup", &mut gi);
|
||||
put_string_field(2, "owner_uid", &mut gi);
|
||||
put_string_field(3, "OwnerNick", &mut gi);
|
||||
put_varint_field(4, 42, &mut gi);
|
||||
let mut rsp = Vec::new();
|
||||
put_varint_field(1, 0, &mut rsp);
|
||||
put_string_field(2, "ok", &mut rsp);
|
||||
put_bytes_field(3, &gi, &mut rsp);
|
||||
|
||||
let parsed = decode_query_group_info_rsp(&rsp).unwrap();
|
||||
assert_eq!(parsed.code, 0);
|
||||
assert_eq!(parsed.group_name, "TestGroup");
|
||||
assert_eq!(parsed.owner_id, "owner_uid");
|
||||
assert_eq!(parsed.member_count, 42);
|
||||
}
|
||||
|
||||
// ─── encode_send_c2c branches ──────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn c2c_encode_with_msg_id_msg_random_group_code_trace_id() {
|
||||
// Hit the branches: msg_id non-empty, msg_random != 0, group_code
|
||||
// non-empty, trace_id non-empty.
|
||||
let buf = encode_send_c2c_message(
|
||||
"uid_alice",
|
||||
"uid_bot",
|
||||
&text_body("hi"),
|
||||
"mid-1",
|
||||
42,
|
||||
"gcode-x",
|
||||
"trace-1",
|
||||
);
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd, biz_cmd::SEND_C2C_MESSAGE);
|
||||
assert_eq!(frame.msg_id, "mid-1");
|
||||
// Re-parse the biz body and check the fields we encoded show up.
|
||||
let f = parse_fields(&frame.data).unwrap();
|
||||
assert_eq!(get_string(&f, 1), "mid-1");
|
||||
assert_eq!(get_string(&f, 2), "uid_alice");
|
||||
assert_eq!(get_string(&f, 3), "uid_bot");
|
||||
assert_eq!(get_varint(&f, 4), 42);
|
||||
assert_eq!(get_string(&f, 6), "gcode-x");
|
||||
// log_ext (field 8) carries nested {1: trace_id}
|
||||
let log_ext = get_bytes(&f, 8);
|
||||
assert!(!log_ext.is_empty());
|
||||
let inner = parse_fields(&log_ext).unwrap();
|
||||
assert_eq!(get_string(&inner, 1), "trace-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn c2c_encode_generates_synthetic_req_id_when_msg_id_empty() {
|
||||
// msg_id empty branch — req_id falls back to `c2c_<seq>`.
|
||||
let buf = encode_send_c2c_message("uid_alice", "uid_bot", &text_body("hi"), "", 0, "", "");
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert!(
|
||||
frame.msg_id.starts_with("c2c_"),
|
||||
"expected synthetic req_id starting with c2c_, got {}",
|
||||
frame.msg_id
|
||||
);
|
||||
}
|
||||
|
||||
// ─── encode_send_group branches ────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn group_encode_with_all_optional_fields() {
|
||||
let buf = encode_send_group_message(
|
||||
"group_42",
|
||||
"uid_bot",
|
||||
&text_body("hello"),
|
||||
"mid-g",
|
||||
"uid_to",
|
||||
"rand_x",
|
||||
"ref-msg-99",
|
||||
"trace-g",
|
||||
);
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd, biz_cmd::SEND_GROUP_MESSAGE);
|
||||
assert_eq!(frame.msg_id, "mid-g");
|
||||
let f = parse_fields(&frame.data).unwrap();
|
||||
assert_eq!(get_string(&f, 1), "mid-g");
|
||||
assert_eq!(get_string(&f, 2), "group_42");
|
||||
assert_eq!(get_string(&f, 3), "uid_bot");
|
||||
assert_eq!(get_string(&f, 4), "uid_to");
|
||||
assert_eq!(get_string(&f, 5), "rand_x");
|
||||
assert_eq!(get_string(&f, 7), "ref-msg-99");
|
||||
let log_ext = get_bytes(&f, 9);
|
||||
let inner = parse_fields(&log_ext).unwrap();
|
||||
assert_eq!(get_string(&inner, 1), "trace-g");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn group_encode_generates_synthetic_req_id_when_msg_id_empty() {
|
||||
let buf =
|
||||
encode_send_group_message("group_x", "uid_bot", &text_body("hi"), "", "", "", "", "");
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert!(
|
||||
frame.msg_id.starts_with("grp_"),
|
||||
"expected synthetic req_id starting with grp_, got {}",
|
||||
frame.msg_id
|
||||
);
|
||||
}
|
||||
|
||||
// ─── encode_send_group_heartbeat ───────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn group_heartbeat_encodes_send_time_and_heartbeat() {
|
||||
let buf = encode_send_group_heartbeat(
|
||||
"hb_g_1",
|
||||
"uid_bot",
|
||||
"group_42",
|
||||
ws_heartbeat::RUNNING,
|
||||
1_700_000_123,
|
||||
);
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd, biz_cmd::SEND_GROUP_HEARTBEAT);
|
||||
assert_eq!(frame.msg_id, "hb_g_1");
|
||||
let f = parse_fields(&frame.data).unwrap();
|
||||
assert_eq!(get_string(&f, 1), "uid_bot");
|
||||
assert_eq!(get_string(&f, 2), ""); // to_account empty for group
|
||||
assert_eq!(get_string(&f, 3), "group_42");
|
||||
assert_eq!(get_varint(&f, 4), 1_700_000_123);
|
||||
assert_eq!(get_varint(&f, 5), ws_heartbeat::RUNNING as u64);
|
||||
}
|
||||
|
||||
// ─── encode_get_group_member_list ──────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn get_group_member_list_omits_offset_when_zero() {
|
||||
let buf = encode_get_group_member_list("qgm_1", "group_42", 0, 100);
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
assert_eq!(frame.cmd, biz_cmd::GET_GROUP_MEMBER_LIST);
|
||||
let f = parse_fields(&frame.data).unwrap();
|
||||
assert_eq!(get_string(&f, 1), "group_42");
|
||||
// offset (field 2) skipped when 0
|
||||
assert_eq!(get_varint(&f, 2), 0);
|
||||
assert_eq!(get_varint(&f, 3), 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_group_member_list_includes_offset_when_nonzero() {
|
||||
let buf = encode_get_group_member_list("qgm_2", "group_42", 200, 50);
|
||||
let frame = decode_conn_msg(&buf).unwrap();
|
||||
let f = parse_fields(&frame.data).unwrap();
|
||||
assert_eq!(get_varint(&f, 2), 200);
|
||||
assert_eq!(get_varint(&f, 3), 50);
|
||||
}
|
||||
|
||||
// ─── decode_biz_rsp_code + decode_response_envelope ────────────
|
||||
|
||||
#[test]
|
||||
fn decode_biz_rsp_code_reads_code_and_message() {
|
||||
let mut buf = Vec::new();
|
||||
put_varint_field(1, 4002, &mut buf);
|
||||
put_string_field(2, "rate limited", &mut buf);
|
||||
let (code, msg) = decode_biz_rsp_code(&buf).unwrap();
|
||||
assert_eq!(code, 4002);
|
||||
assert_eq!(msg, "rate limited");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_biz_rsp_code_on_empty_returns_defaults() {
|
||||
let (code, msg) = decode_biz_rsp_code(&[]).unwrap();
|
||||
assert_eq!(code, 0);
|
||||
assert!(msg.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_response_envelope_extracts_frame() {
|
||||
let original = encode_conn_msg(
|
||||
cmd_type::RESPONSE,
|
||||
biz_cmd::SEND_C2C_MESSAGE,
|
||||
1,
|
||||
"mid-r",
|
||||
module::BIZ_PKG,
|
||||
&[0xAA, 0xBB],
|
||||
);
|
||||
let frame = decode_response_envelope(&original).unwrap();
|
||||
assert_eq!(frame.cmd, biz_cmd::SEND_C2C_MESSAGE);
|
||||
assert_eq!(frame.msg_id, "mid-r");
|
||||
assert_eq!(frame.data, vec![0xAA, 0xBB]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn group_member_list_decode() {
|
||||
let mut m1 = Vec::new();
|
||||
put_string_field(1, "uid_a", &mut m1);
|
||||
put_string_field(2, "Alice", &mut m1);
|
||||
put_varint_field(3, 2, &mut m1);
|
||||
let mut rsp = Vec::new();
|
||||
put_varint_field(1, 0, &mut rsp);
|
||||
put_string_field(2, "ok", &mut rsp);
|
||||
put_bytes_field(3, &m1, &mut rsp);
|
||||
put_varint_field(4, 100, &mut rsp);
|
||||
put_varint_field(5, 1, &mut rsp);
|
||||
|
||||
let page = decode_get_group_member_list_rsp(&rsp).unwrap();
|
||||
assert_eq!(page.members.len(), 1);
|
||||
assert_eq!(page.members[0].user_id, "uid_a");
|
||||
assert_eq!(page.members[0].role, 2);
|
||||
assert_eq!(page.next_offset, 100);
|
||||
assert!(page.is_complete);
|
||||
}
|
||||
|
||||
/// Adversarial input: a varint that overflows i32. The decoder must
|
||||
/// surface `YuanbaoError::ProtoDecode` instead of silently truncating
|
||||
/// (which would corrupt the `code` field returned to callers).
|
||||
#[test]
|
||||
fn decode_biz_rsp_code_rejects_varint_out_of_i32_range() {
|
||||
let mut buf = Vec::new();
|
||||
put_varint_field(1, u64::MAX, &mut buf);
|
||||
put_string_field(2, "ok", &mut buf);
|
||||
match decode_biz_rsp_code(&buf).unwrap_err() {
|
||||
YuanbaoError::ProtoDecode(m) => {
|
||||
assert!(
|
||||
m.contains("out of i32 range"),
|
||||
"expected i32 overflow message, got: {m}"
|
||||
);
|
||||
}
|
||||
other => panic!("expected ProtoDecode, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Same guard applied to the group-member-list `next_offset` field —
|
||||
/// an oversized varint must produce a structured decode error, not a
|
||||
/// silent `as u32` wrap that would mis-paginate subsequent fetches.
|
||||
#[test]
|
||||
fn decode_group_member_list_rejects_varint_out_of_u32_range() {
|
||||
let mut rsp = Vec::new();
|
||||
put_varint_field(1, 0, &mut rsp);
|
||||
put_string_field(2, "ok", &mut rsp);
|
||||
put_varint_field(4, u64::from(u32::MAX) + 1, &mut rsp);
|
||||
match decode_get_group_member_list_rsp(&rsp).unwrap_err() {
|
||||
YuanbaoError::ProtoDecode(m) => {
|
||||
assert!(
|
||||
m.contains("out of u32 range"),
|
||||
"expected u32 overflow message, got: {m}"
|
||||
);
|
||||
}
|
||||
other => panic!("expected ProtoDecode, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,90 +0,0 @@
|
||||
//! Yuanbao WebSocket protocol constants.
|
||||
//!
|
||||
//! Values mirror `gateway/platforms/yuanbao_proto.py` in hermes-agent
|
||||
//! (the authoritative reference implementation).
|
||||
|
||||
/// `ConnMsg.Head.cmd_type` enum.
|
||||
pub mod cmd_type {
|
||||
/// Upstream request.
|
||||
pub const REQUEST: u32 = 0;
|
||||
/// Response to a previous upstream request.
|
||||
pub const RESPONSE: u32 = 1;
|
||||
/// Downstream push from server.
|
||||
pub const PUSH: u32 = 2;
|
||||
/// ACK reply to a downstream push.
|
||||
pub const PUSH_ACK: u32 = 3;
|
||||
}
|
||||
|
||||
/// Built-in command words used in `ConnMsg.Head.cmd`.
|
||||
pub mod cmd {
|
||||
pub const AUTH_BIND: &str = "auth-bind";
|
||||
pub const PING: &str = "ping";
|
||||
pub const KICKOUT: &str = "kickout";
|
||||
pub const UPDATE_META: &str = "update-meta";
|
||||
}
|
||||
|
||||
/// Module / service names used in `ConnMsg.Head.module`.
|
||||
pub mod module {
|
||||
pub const CONN_ACCESS: &str = "conn_access";
|
||||
/// Short name of the openclaw biz module (matches TS client).
|
||||
pub const BIZ_PKG: &str = "yuanbao_openclaw_proxy";
|
||||
}
|
||||
|
||||
/// Business command words (`ConnMsg.Head.cmd` when module=BIZ_PKG).
|
||||
///
|
||||
/// Note: there is intentionally no constant for the inbound push cmd —
|
||||
/// the yuanbao gateway uses several cmd words for inbound messages and
|
||||
/// the routing is purely by `cmd_type=Push` (see `connection.rs` /
|
||||
/// `mod.rs::dispatch_push`).
|
||||
pub mod biz_cmd {
|
||||
pub const SEND_C2C_MESSAGE: &str = "send_c2c_message";
|
||||
pub const SEND_GROUP_MESSAGE: &str = "send_group_message";
|
||||
pub const SEND_PRIVATE_HEARTBEAT: &str = "send_private_heartbeat";
|
||||
pub const SEND_GROUP_HEARTBEAT: &str = "send_group_heartbeat";
|
||||
pub const QUERY_GROUP_INFO: &str = "query_group_info";
|
||||
pub const GET_GROUP_MEMBER_LIST: &str = "get_group_member_list";
|
||||
}
|
||||
|
||||
/// Reply Heartbeat status enum (`heartbeat` field of `Send*HeartbeatReq`).
|
||||
pub mod ws_heartbeat {
|
||||
/// Bot is currently producing output.
|
||||
pub const RUNNING: u32 = 1;
|
||||
/// Bot has finished its turn.
|
||||
pub const FINISH: u32 = 2;
|
||||
}
|
||||
|
||||
/// TIM `msg_type` string constants for `MsgBodyElement.msg_type`.
|
||||
pub mod tim {
|
||||
pub const TEXT: &str = "TIMTextElem";
|
||||
pub const IMAGE: &str = "TIMImageElem";
|
||||
pub const FILE: &str = "TIMFileElem";
|
||||
pub const SOUND: &str = "TIMSoundElem";
|
||||
pub const VIDEO: &str = "TIMVideoFileElem";
|
||||
pub const FACE: &str = "TIMFaceElem";
|
||||
pub const CUSTOM: &str = "TIMCustomElem";
|
||||
}
|
||||
|
||||
/// Fixed instance id reported in `AuthBindReq.DeviceInfo.instance_id` and
|
||||
/// the `X-Instance-Id` HTTP header. Mirrors `OPENCLAW_ID = 20` used by
|
||||
/// `yuanbao-openclaw-plugin` (`src/access/ws/conn-codec.ts`) — the server
|
||||
/// keys some checks off this value, so it must match the value the sign
|
||||
/// endpoint sees when the token is minted.
|
||||
pub const OPENHUMAN_INSTANCE_ID: &str = "20";
|
||||
|
||||
/// Reconnect backoff schedule (seconds). Mirrors hermes-agent.
|
||||
pub const RECONNECT_DELAYS: &[u64] = &[1, 2, 5, 10, 30, 60];
|
||||
pub const MAX_RECONNECT_ATTEMPTS: u32 = 100;
|
||||
|
||||
/// Ping interval (seconds). Server-driven; this is the upper bound.
|
||||
pub const PING_INTERVAL_SECS: u64 = 30;
|
||||
/// Number of consecutive ping timeouts before the connection is dropped.
|
||||
pub const HEARTBEAT_TIMEOUT_THRESHOLD: u32 = 2;
|
||||
/// Per-request biz timeout (seconds).
|
||||
pub const DEFAULT_SEND_TIMEOUT_SECS: u64 = 30;
|
||||
/// Auth-bind handshake timeout (seconds).
|
||||
pub const AUTH_TIMEOUT_SECS: u64 = 15;
|
||||
|
||||
/// Inbound dedup TTL — drop a `msg_id` we've already seen within this window.
|
||||
pub const DEDUP_TTL_SECS: u64 = 300;
|
||||
/// LRU-style cap on the dedup table.
|
||||
pub const DEDUP_CAPACITY: usize = 10_000;
|
||||
@@ -1,629 +0,0 @@
|
||||
//! Token sign manager — talks to `/api/v5/robotLogic/sign-token` to
|
||||
//! exchange `(app_key, app_secret)` for a short-lived WS token + bot_id.
|
||||
//!
|
||||
//! Mirrors hermes-agent `SignManager` (yuanbao.py 641-881). Implements:
|
||||
//! - per-app_key tokio `Mutex` to coalesce concurrent refresh attempts
|
||||
//! - 60-second early-refresh margin to avoid using a token that's
|
||||
//! about to expire mid-handshake
|
||||
//! - retry on `code=10099` up to 3 times
|
||||
//!
|
||||
//! Signature scheme (TS plugin compatible):
|
||||
//! plain = nonce + timestamp + app_key + app_secret
|
||||
//! signature = HMAC-SHA256(key = app_secret, msg = plain) as lower-hex
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use chrono::FixedOffset;
|
||||
use hmac::{Hmac, Mac};
|
||||
use sha2::Sha256;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{info, warn};
|
||||
|
||||
use super::errors::YuanbaoError;
|
||||
|
||||
const SIGN_PATH: &str = "/api/v5/robotLogic/sign-token";
|
||||
const RETRYABLE_CODE: i64 = 10099;
|
||||
const MAX_RETRIES: usize = 3;
|
||||
const RETRY_DELAY_MS: u64 = 1_000;
|
||||
/// Treat as expiring this many seconds before actual expiry so a fresh
|
||||
/// token is fetched before the running one dies mid-request.
|
||||
const CACHE_REFRESH_MARGIN_SECS: u64 = 60;
|
||||
const HTTP_TIMEOUT_SECS: u64 = 10;
|
||||
const DEFAULT_DURATION_SECS: u64 = 3600;
|
||||
|
||||
/// One cached token entry.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TokenEntry {
|
||||
pub token: String,
|
||||
pub bot_id: String,
|
||||
pub product: String,
|
||||
pub source: String,
|
||||
/// Seconds-since-epoch when this token expires server-side.
|
||||
pub expire_ts: u64,
|
||||
}
|
||||
|
||||
impl TokenEntry {
|
||||
pub fn is_valid(&self) -> bool {
|
||||
let now = unix_now();
|
||||
self.expire_ts > now + CACHE_REFRESH_MARGIN_SECS
|
||||
}
|
||||
|
||||
pub fn seconds_remaining(&self) -> i64 {
|
||||
self.expire_ts as i64 - unix_now() as i64
|
||||
}
|
||||
}
|
||||
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
|
||||
/// Compute the `signature` field for the sign-token API.
|
||||
pub fn compute_signature(nonce: &str, timestamp: &str, app_key: &str, app_secret: &str) -> String {
|
||||
let plain = format!("{nonce}{timestamp}{app_key}{app_secret}");
|
||||
let mut mac =
|
||||
HmacSha256::new_from_slice(app_secret.as_bytes()).expect("HMAC accepts any key length");
|
||||
mac.update(plain.as_bytes());
|
||||
hex::encode(mac.finalize().into_bytes())
|
||||
}
|
||||
|
||||
/// Build Beijing-time ISO-8601 timestamp without milliseconds.
|
||||
/// Format: `2006-01-02T15:04:05+08:00`.
|
||||
pub fn build_timestamp() -> String {
|
||||
let bj_offset = FixedOffset::east_opt(8 * 3600).expect("valid offset");
|
||||
let now = chrono::Utc::now().with_timezone(&bj_offset);
|
||||
now.format("%Y-%m-%dT%H:%M:%S+08:00").to_string()
|
||||
}
|
||||
|
||||
/// Generate a 32-char hex nonce.
|
||||
pub fn generate_nonce() -> String {
|
||||
let mut bytes = [0u8; 16];
|
||||
for b in &mut bytes {
|
||||
*b = rand::random::<u8>();
|
||||
}
|
||||
hex::encode(bytes)
|
||||
}
|
||||
|
||||
/// Process-wide token manager. One instance is built per `YuanbaoChannel`
|
||||
/// and shared with the connection layer; the per-app_key Mutex makes it
|
||||
/// safe to have multiple connections sharing this manager.
|
||||
pub struct SignManager {
|
||||
http: reqwest::Client,
|
||||
/// Per-app_key refresh mutexes — coalesce concurrent refresh attempts.
|
||||
locks: Mutex<HashMap<String, Arc<Mutex<()>>>>,
|
||||
/// Token cache keyed by app_key.
|
||||
cache: Mutex<HashMap<String, TokenEntry>>,
|
||||
}
|
||||
|
||||
impl SignManager {
|
||||
pub fn new(http: reqwest::Client) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
http,
|
||||
locks: Mutex::new(HashMap::new()),
|
||||
cache: Mutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Look up a cached token without touching the network.
|
||||
pub async fn cached(&self, app_key: &str) -> Option<TokenEntry> {
|
||||
let cache = self.cache.lock().await;
|
||||
cache.get(app_key).cloned().filter(|e| e.is_valid())
|
||||
}
|
||||
|
||||
/// Test-only: inject a cache entry without touching the sign endpoint.
|
||||
#[cfg(test)]
|
||||
pub(crate) async fn set_cached_for_test(&self, app_key: &str, entry: TokenEntry) {
|
||||
self.cache.lock().await.insert(app_key.to_string(), entry);
|
||||
}
|
||||
|
||||
/// Get a valid token, fetching one if the cache is empty or stale.
|
||||
pub async fn get_token(
|
||||
&self,
|
||||
app_key: &str,
|
||||
app_secret: &str,
|
||||
api_domain: &str,
|
||||
route_env: &str,
|
||||
) -> Result<TokenEntry, YuanbaoError> {
|
||||
if let Some(entry) = self.cached(app_key).await {
|
||||
info!(
|
||||
"[yuanbao/sign] using cached token ({}s remaining)",
|
||||
entry.seconds_remaining()
|
||||
);
|
||||
return Ok(entry);
|
||||
}
|
||||
self.refresh(app_key, app_secret, api_domain, route_env)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Force-refresh: drop the cache entry and re-fetch.
|
||||
pub async fn force_refresh(
|
||||
&self,
|
||||
app_key: &str,
|
||||
app_secret: &str,
|
||||
api_domain: &str,
|
||||
route_env: &str,
|
||||
) -> Result<TokenEntry, YuanbaoError> {
|
||||
{
|
||||
let mut cache = self.cache.lock().await;
|
||||
cache.remove(app_key);
|
||||
}
|
||||
warn!(
|
||||
"[yuanbao/sign] force-refresh app_key=****{}",
|
||||
suffix(app_key)
|
||||
);
|
||||
self.refresh(app_key, app_secret, api_domain, route_env)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn refresh(
|
||||
&self,
|
||||
app_key: &str,
|
||||
app_secret: &str,
|
||||
api_domain: &str,
|
||||
route_env: &str,
|
||||
) -> Result<TokenEntry, YuanbaoError> {
|
||||
let lock = self.get_refresh_lock(app_key).await;
|
||||
let _g = lock.lock().await;
|
||||
|
||||
// Double-checked locking: another task may have refreshed while we waited.
|
||||
if let Some(entry) = self.cached(app_key).await {
|
||||
return Ok(entry);
|
||||
}
|
||||
|
||||
let entry = self
|
||||
.fetch_with_retry(app_key, app_secret, api_domain, route_env)
|
||||
.await?;
|
||||
let mut cache = self.cache.lock().await;
|
||||
cache.insert(app_key.to_string(), entry.clone());
|
||||
Ok(entry)
|
||||
}
|
||||
|
||||
async fn get_refresh_lock(&self, app_key: &str) -> Arc<Mutex<()>> {
|
||||
let mut locks = self.locks.lock().await;
|
||||
locks
|
||||
.entry(app_key.to_string())
|
||||
.or_insert_with(|| Arc::new(Mutex::new(())))
|
||||
.clone()
|
||||
}
|
||||
|
||||
async fn fetch_with_retry(
|
||||
&self,
|
||||
app_key: &str,
|
||||
app_secret: &str,
|
||||
api_domain: &str,
|
||||
route_env: &str,
|
||||
) -> Result<TokenEntry, YuanbaoError> {
|
||||
let url = format!("{}{}", api_domain.trim_end_matches('/'), SIGN_PATH);
|
||||
let mut last_err: Option<YuanbaoError> = None;
|
||||
|
||||
for attempt in 0..=MAX_RETRIES {
|
||||
let nonce = generate_nonce();
|
||||
let timestamp = build_timestamp();
|
||||
let signature = compute_signature(&nonce, ×tamp, app_key, app_secret);
|
||||
let payload = serde_json::json!({
|
||||
"app_key": app_key,
|
||||
"nonce": nonce,
|
||||
"signature": signature,
|
||||
"timestamp": timestamp,
|
||||
});
|
||||
|
||||
let mut req = self
|
||||
.http
|
||||
.post(&url)
|
||||
.timeout(Duration::from_secs(HTTP_TIMEOUT_SECS))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("X-AppVersion", super::config::DEFAULT_PLUGIN_VERSION)
|
||||
.header("X-OperationSystem", "linux")
|
||||
.header(
|
||||
"X-Instance-Id",
|
||||
super::proto_constants::OPENHUMAN_INSTANCE_ID,
|
||||
)
|
||||
.header("X-Bot-Version", env!("CARGO_PKG_VERSION"));
|
||||
if !route_env.is_empty() {
|
||||
req = req.header("X-Route-Env", route_env);
|
||||
}
|
||||
|
||||
info!(
|
||||
"[yuanbao/sign] POST {}{}",
|
||||
url,
|
||||
if attempt > 0 {
|
||||
format!(" (retry {attempt}/{MAX_RETRIES})")
|
||||
} else {
|
||||
String::new()
|
||||
}
|
||||
);
|
||||
|
||||
let resp = match req.json(&payload).send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
last_err = Some(YuanbaoError::Connection(format!("sign-token: {e}")));
|
||||
if attempt < MAX_RETRIES {
|
||||
tokio::time::sleep(Duration::from_millis(RETRY_DELAY_MS)).await;
|
||||
continue;
|
||||
}
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(YuanbaoError::AuthFailed(format!(
|
||||
"sign-token HTTP {status}: {}",
|
||||
&body.chars().take(200).collect::<String>()
|
||||
)));
|
||||
}
|
||||
|
||||
let json: serde_json::Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| YuanbaoError::AuthFailed(format!("sign-token body: {e}")))?;
|
||||
|
||||
let code = json.get("code").and_then(|c| c.as_i64()).unwrap_or(0);
|
||||
if code == 0 {
|
||||
let data = match json.get("data") {
|
||||
Some(d) if d.is_object() => d,
|
||||
_ => {
|
||||
return Err(YuanbaoError::AuthFailed(
|
||||
"sign-token response missing 'data'".into(),
|
||||
))
|
||||
}
|
||||
};
|
||||
let duration = data
|
||||
.get("duration")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(DEFAULT_DURATION_SECS);
|
||||
let entry = TokenEntry {
|
||||
token: data
|
||||
.get("token")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
bot_id: data
|
||||
.get("bot_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
product: data
|
||||
.get("product")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
source: data
|
||||
.get("source")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
expire_ts: unix_now() + duration,
|
||||
};
|
||||
info!(
|
||||
"[yuanbao/sign] success: bot_id={} duration={}s",
|
||||
entry.bot_id, duration
|
||||
);
|
||||
return Ok(entry);
|
||||
}
|
||||
|
||||
if code == RETRYABLE_CODE && attempt < MAX_RETRIES {
|
||||
warn!(
|
||||
"[yuanbao/sign] retryable code={code}, retrying in {RETRY_DELAY_MS}ms (attempt {})",
|
||||
attempt + 1
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(RETRY_DELAY_MS)).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
let msg = json
|
||||
.get("msg")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string();
|
||||
return Err(YuanbaoError::AuthFailed(format!(
|
||||
"sign-token code={code} msg={msg}"
|
||||
)));
|
||||
}
|
||||
|
||||
Err(last_err.unwrap_or(YuanbaoError::AuthFailed(
|
||||
"sign-token max retries exceeded".into(),
|
||||
)))
|
||||
}
|
||||
|
||||
/// Drop all per-app_key locks. Called on channel shutdown to avoid
|
||||
/// leaking entries across reconnects within the same process.
|
||||
pub async fn clear_locks(&self) {
|
||||
let mut locks = self.locks.lock().await;
|
||||
locks.clear();
|
||||
}
|
||||
}
|
||||
|
||||
fn unix_now() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn suffix(s: &str) -> &str {
|
||||
if s.len() <= 4 {
|
||||
s
|
||||
} else {
|
||||
&s[s.len() - 4..]
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn signature_matches_python_reference() {
|
||||
// Reproducible vector — hand-computed:
|
||||
// plain = "n123" + "2026-05-19T22:00:00+08:00" + "app_k" + "secret"
|
||||
// sig = HMAC-SHA256(key="secret", msg=plain) as lower hex
|
||||
let sig = compute_signature("n123", "2026-05-19T22:00:00+08:00", "app_k", "secret");
|
||||
// We don't pin the exact bytes (would require running Python to confirm) —
|
||||
// instead verify the contract: same inputs → same output, 64-char hex.
|
||||
assert_eq!(sig.len(), 64);
|
||||
assert!(sig.chars().all(|c| c.is_ascii_hexdigit()));
|
||||
let sig2 = compute_signature("n123", "2026-05-19T22:00:00+08:00", "app_k", "secret");
|
||||
assert_eq!(sig, sig2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signature_varies_with_inputs() {
|
||||
let s1 = compute_signature("n1", "t", "ak", "sk");
|
||||
let s2 = compute_signature("n2", "t", "ak", "sk");
|
||||
let s3 = compute_signature("n1", "t2", "ak", "sk");
|
||||
let s4 = compute_signature("n1", "t", "ak2", "sk");
|
||||
let s5 = compute_signature("n1", "t", "ak", "sk2");
|
||||
let all = [&s1, &s2, &s3, &s4, &s5];
|
||||
for (i, a) in all.iter().enumerate() {
|
||||
for (j, b) in all.iter().enumerate() {
|
||||
if i != j {
|
||||
assert_ne!(a, b, "inputs {i} vs {j} should differ");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nonce_is_32_char_hex() {
|
||||
let n = generate_nonce();
|
||||
assert_eq!(n.len(), 32);
|
||||
assert!(n.chars().all(|c| c.is_ascii_hexdigit()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn timestamp_matches_beijing_format() {
|
||||
let t = build_timestamp();
|
||||
// 2006-01-02T15:04:05+08:00 → length 25
|
||||
assert_eq!(t.len(), 25);
|
||||
assert!(t.ends_with("+08:00"));
|
||||
assert_eq!(&t[4..5], "-");
|
||||
assert_eq!(&t[7..8], "-");
|
||||
assert_eq!(&t[10..11], "T");
|
||||
assert_eq!(&t[13..14], ":");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_entry_is_valid_only_with_margin() {
|
||||
let mut e = TokenEntry {
|
||||
token: "t".into(),
|
||||
bot_id: "b".into(),
|
||||
product: String::new(),
|
||||
source: String::new(),
|
||||
expire_ts: unix_now() + 120,
|
||||
};
|
||||
assert!(e.is_valid());
|
||||
e.expire_ts = unix_now() + 30; // less than 60s margin
|
||||
assert!(!e.is_valid());
|
||||
e.expire_ts = unix_now().saturating_sub(10);
|
||||
assert!(!e.is_valid());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cache_returns_entry_when_valid() {
|
||||
let mgr = SignManager::new(reqwest::Client::new());
|
||||
let entry = TokenEntry {
|
||||
token: "tok".into(),
|
||||
bot_id: "bot".into(),
|
||||
product: String::new(),
|
||||
source: String::new(),
|
||||
expire_ts: unix_now() + 600,
|
||||
};
|
||||
mgr.cache.lock().await.insert("ak".into(), entry.clone());
|
||||
let got = mgr.cached("ak").await.expect("cache hit");
|
||||
assert_eq!(got.token, "tok");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cache_drops_expired_entry() {
|
||||
let mgr = SignManager::new(reqwest::Client::new());
|
||||
mgr.cache.lock().await.insert(
|
||||
"ak".into(),
|
||||
TokenEntry {
|
||||
token: "tok".into(),
|
||||
bot_id: "bot".into(),
|
||||
product: String::new(),
|
||||
source: String::new(),
|
||||
expire_ts: unix_now() + 10, // under margin
|
||||
},
|
||||
);
|
||||
assert!(mgr.cached("ak").await.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_entry_seconds_remaining_is_signed() {
|
||||
let e_future = TokenEntry {
|
||||
token: "t".into(),
|
||||
bot_id: "b".into(),
|
||||
product: String::new(),
|
||||
source: String::new(),
|
||||
expire_ts: unix_now() + 300,
|
||||
};
|
||||
assert!(e_future.seconds_remaining() >= 290);
|
||||
let e_past = TokenEntry {
|
||||
expire_ts: unix_now().saturating_sub(60),
|
||||
..e_future
|
||||
};
|
||||
assert!(e_past.seconds_remaining() <= 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suffix_redacts_to_last_4_chars() {
|
||||
assert_eq!(suffix(""), "");
|
||||
assert_eq!(suffix("a"), "a");
|
||||
assert_eq!(suffix("abcd"), "abcd");
|
||||
assert_eq!(suffix("abcdef"), "cdef");
|
||||
assert_eq!(suffix("0123456789"), "6789");
|
||||
}
|
||||
|
||||
// ─── refresh / fetch_with_retry via wiremock ────────────────
|
||||
|
||||
fn ok_body(token: &str, bot_id: &str, duration_secs: u64) -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"code": 0,
|
||||
"msg": "ok",
|
||||
"data": {
|
||||
"token": token,
|
||||
"bot_id": bot_id,
|
||||
"product": "prod1",
|
||||
"source": "src1",
|
||||
"duration": duration_secs,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_token_fetches_and_caches_on_first_call() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::path(SIGN_PATH))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(ok_body("tok-1", "bot-1", 7200)),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let mgr = SignManager::new(reqwest::Client::new());
|
||||
let e = mgr
|
||||
.get_token("ak", "sk", &server.uri(), "")
|
||||
.await
|
||||
.expect("token");
|
||||
assert_eq!(e.token, "tok-1");
|
||||
assert_eq!(e.bot_id, "bot-1");
|
||||
assert!(e.expire_ts > unix_now() + 7000);
|
||||
|
||||
// Second call should hit the cache (still works even if server stops).
|
||||
let cached = mgr.cached("ak").await.expect("cached");
|
||||
assert_eq!(cached.token, "tok-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_token_retries_on_code_10099_then_succeeds() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
// First two requests return code=10099, third returns code=0.
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::path(SIGN_PATH))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"code": 10099,
|
||||
"msg": "try again",
|
||||
})),
|
||||
)
|
||||
.up_to_n_times(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::path(SIGN_PATH))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(ok_body("tok-r", "bot-r", 600)),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let mgr = SignManager::new(reqwest::Client::new());
|
||||
let e = mgr.refresh("ak", "sk", &server.uri(), "").await.unwrap();
|
||||
assert_eq!(e.token, "tok-r");
|
||||
assert_eq!(e.bot_id, "bot-r");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_token_surfaces_http_error_as_auth_failed() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.respond_with(wiremock::ResponseTemplate::new(401).set_body_string("Unauthorized"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let mgr = SignManager::new(reqwest::Client::new());
|
||||
let err = mgr
|
||||
.get_token("ak", "sk", &server.uri(), "")
|
||||
.await
|
||||
.unwrap_err();
|
||||
match err {
|
||||
YuanbaoError::AuthFailed(m) => assert!(m.contains("HTTP 401"), "got {m}"),
|
||||
other => panic!("expected AuthFailed, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_token_fails_on_non_zero_business_code() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(serde_json::json!({
|
||||
"code": 40001,
|
||||
"msg": "bad secret",
|
||||
})),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let mgr = SignManager::new(reqwest::Client::new());
|
||||
let err = mgr
|
||||
.get_token("ak", "sk", &server.uri(), "")
|
||||
.await
|
||||
.unwrap_err();
|
||||
match err {
|
||||
YuanbaoError::AuthFailed(m) => {
|
||||
assert!(m.contains("code=40001"), "got {m}");
|
||||
assert!(m.contains("bad secret"), "got {m}");
|
||||
}
|
||||
other => panic!("expected AuthFailed, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn force_refresh_evicts_cache_and_refetches() {
|
||||
let server = wiremock::MockServer::start().await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::path(SIGN_PATH))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(ok_body("tok-a", "bot", 600)),
|
||||
)
|
||||
.up_to_n_times(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||||
.and(wiremock::matchers::path(SIGN_PATH))
|
||||
.respond_with(
|
||||
wiremock::ResponseTemplate::new(200).set_body_json(ok_body("tok-b", "bot", 600)),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let mgr = SignManager::new(reqwest::Client::new());
|
||||
let first = mgr.get_token("ak", "sk", &server.uri(), "").await.unwrap();
|
||||
assert_eq!(first.token, "tok-a");
|
||||
let second = mgr
|
||||
.force_refresh("ak", "sk", &server.uri(), "to_env")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(second.token, "tok-b");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn clear_locks_drops_all_per_app_key_mutexes() {
|
||||
let mgr = SignManager::new(reqwest::Client::new());
|
||||
// Prime the locks map.
|
||||
let _ = mgr.get_refresh_lock("ak-1").await;
|
||||
let _ = mgr.get_refresh_lock("ak-2").await;
|
||||
assert_eq!(mgr.locks.lock().await.len(), 2);
|
||||
mgr.clear_locks().await;
|
||||
assert!(mgr.locks.lock().await.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -1,209 +0,0 @@
|
||||
//! Fence-aware Markdown splitter.
|
||||
//!
|
||||
//! When a long AI response is split into N chunks for the Yuanbao
|
||||
//! `max_message_length` cap, we must not break inside:
|
||||
//! - a fenced code block (``` … ``` or ~~~ … ~~~)
|
||||
//! - a Markdown table row (lines starting with `|`)
|
||||
//! - a list-continuation block
|
||||
//!
|
||||
//! Strategy: walk the input by line, tracking fence/table state, and
|
||||
//! emit a chunk every time adding the next line would push the buffer
|
||||
//! past the cap **and** the cap boundary is safe (not inside a fence,
|
||||
//! not in the middle of a table). If a single line is itself longer
|
||||
//! than the cap, hard-split at a char boundary.
|
||||
|
||||
/// Split `text` into chunks no larger than `cap_bytes` (utf-8 byte count),
|
||||
/// preserving fenced code blocks and table rows where possible.
|
||||
pub fn split_markdown(text: &str, cap_bytes: usize) -> Vec<String> {
|
||||
if text.len() <= cap_bytes {
|
||||
return vec![text.to_string()];
|
||||
}
|
||||
let cap = cap_bytes.max(1);
|
||||
// Reserve a small headroom so the trailing newline / final char fits
|
||||
// when we flush. For very small caps fall back to no margin so callers
|
||||
// testing tight bounds (cap=20) still get chunks under the cap.
|
||||
let safe_cap = if cap >= 32 {
|
||||
cap.saturating_sub(16)
|
||||
} else {
|
||||
cap
|
||||
};
|
||||
|
||||
let mut chunks: Vec<String> = Vec::new();
|
||||
let mut buf = String::with_capacity(cap);
|
||||
let mut in_fence = false;
|
||||
let mut fence_marker: Option<String> = None;
|
||||
|
||||
for line in text.split_inclusive('\n') {
|
||||
let trimmed = line.trim_end_matches('\n');
|
||||
let starts_fence = is_fence(trimmed);
|
||||
|
||||
// If this single line is wider than the cap, we must hard-split it.
|
||||
if line.len() > cap {
|
||||
flush(&mut chunks, &mut buf);
|
||||
for piece in hard_split(line, cap) {
|
||||
chunks.push(piece);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let candidate_len = buf.len() + line.len();
|
||||
if candidate_len > safe_cap && !buf.is_empty() && safe_to_break(in_fence) {
|
||||
flush(&mut chunks, &mut buf);
|
||||
}
|
||||
buf.push_str(line);
|
||||
|
||||
if let Some(marker) = starts_fence {
|
||||
if let Some(open) = &fence_marker {
|
||||
if marker == *open {
|
||||
in_fence = false;
|
||||
fence_marker = None;
|
||||
}
|
||||
} else {
|
||||
in_fence = true;
|
||||
fence_marker = Some(marker);
|
||||
}
|
||||
}
|
||||
}
|
||||
flush(&mut chunks, &mut buf);
|
||||
|
||||
// Drop empty trailing chunks (can happen if input ends on newline).
|
||||
chunks.retain(|c| !c.trim().is_empty());
|
||||
chunks
|
||||
}
|
||||
|
||||
fn flush(chunks: &mut Vec<String>, buf: &mut String) {
|
||||
if !buf.is_empty() {
|
||||
chunks.push(buf.trim_end().to_string());
|
||||
buf.clear();
|
||||
}
|
||||
}
|
||||
|
||||
fn safe_to_break(in_fence: bool) -> bool {
|
||||
!in_fence
|
||||
}
|
||||
|
||||
/// If `line` opens or closes a fenced code block, return the marker text
|
||||
/// (e.g. "```" or "~~~"). A line that contains a fence in the middle is
|
||||
/// NOT a fence marker; only lines that *start* with three or more
|
||||
/// backticks/tildes count.
|
||||
fn is_fence(line: &str) -> Option<String> {
|
||||
let trimmed = line.trim_start();
|
||||
if let Some(rest) = trimmed.strip_prefix("```") {
|
||||
// Allow optional language tag after the fence.
|
||||
let _ = rest;
|
||||
return Some("```".into());
|
||||
}
|
||||
if let Some(rest) = trimmed.strip_prefix("~~~") {
|
||||
let _ = rest;
|
||||
return Some("~~~".into());
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Last-resort splitter for a line that's wider than the cap.
|
||||
fn hard_split(line: &str, cap: usize) -> Vec<String> {
|
||||
let mut out = Vec::new();
|
||||
let mut remaining = line;
|
||||
while !remaining.is_empty() {
|
||||
if remaining.len() <= cap {
|
||||
out.push(remaining.to_string());
|
||||
break;
|
||||
}
|
||||
let mut idx = cap;
|
||||
while idx > 0 && !remaining.is_char_boundary(idx) {
|
||||
idx -= 1;
|
||||
}
|
||||
if idx == 0 {
|
||||
// pathological — emit one char at a time
|
||||
let take = remaining
|
||||
.char_indices()
|
||||
.nth(1)
|
||||
.map(|(i, _)| i)
|
||||
.unwrap_or(remaining.len());
|
||||
let (chunk, rest) = remaining.split_at(take);
|
||||
out.push(chunk.to_string());
|
||||
remaining = rest;
|
||||
} else {
|
||||
let (chunk, rest) = remaining.split_at(idx);
|
||||
out.push(chunk.to_string());
|
||||
remaining = rest;
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn short_input_returns_one_chunk() {
|
||||
let r = split_markdown("hello", 100);
|
||||
assert_eq!(r, vec!["hello"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn splits_on_newlines_respecting_cap() {
|
||||
let input = "a\n".repeat(100);
|
||||
let r = split_markdown(&input, 20);
|
||||
assert!(r.len() > 1);
|
||||
for c in &r {
|
||||
assert!(c.len() <= 20, "chunk too long: {c:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_fenced_code_block() {
|
||||
let input = "intro line\n\
|
||||
```rust\n\
|
||||
fn long_function_a() -> u32 { 42 }\n\
|
||||
fn long_function_b() -> u32 { 43 }\n\
|
||||
fn long_function_c() -> u32 { 44 }\n\
|
||||
```\n\
|
||||
trailing text";
|
||||
let chunks = split_markdown(input, 80);
|
||||
// Find the chunk(s) containing the fence — they must not split mid-fence.
|
||||
let mut open = 0;
|
||||
for c in &chunks {
|
||||
for line in c.lines() {
|
||||
if is_fence(line).is_some() {
|
||||
open += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
// The fence must appear as balanced pairs.
|
||||
assert_eq!(open % 2, 0, "unbalanced fences after split: {chunks:#?}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hard_split_very_long_line() {
|
||||
let line = "x".repeat(500);
|
||||
let r = split_markdown(&line, 100);
|
||||
for c in &r {
|
||||
assert!(c.len() <= 100, "chunk too long: {}", c.len());
|
||||
}
|
||||
assert_eq!(r.join("").len(), 500);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unicode_safe_hard_split() {
|
||||
let line = "中".repeat(200); // each char is 3 bytes → 600 total
|
||||
let r = split_markdown(&line, 50);
|
||||
for c in &r {
|
||||
assert!(c.len() <= 50, "chunk too long: {}", c.len());
|
||||
// verify it's valid utf-8 by reading it
|
||||
for ch in c.chars() {
|
||||
assert!(ch == '中');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_fence_detects_backticks() {
|
||||
assert_eq!(is_fence("```").as_deref(), Some("```"));
|
||||
assert_eq!(is_fence("```rust").as_deref(), Some("```"));
|
||||
assert_eq!(is_fence("~~~").as_deref(), Some("~~~"));
|
||||
assert_eq!(is_fence("text").as_deref(), None);
|
||||
assert_eq!(is_fence("``").as_deref(), None);
|
||||
}
|
||||
}
|
||||
@@ -1,441 +0,0 @@
|
||||
//! Shared domain types for the Yuanbao channel.
|
||||
//!
|
||||
//! Field naming follows the upstream Yuanbao protocol (`from_account`,
|
||||
//! `group_code`, `msg_id`, etc.) so that the protobuf decoders, the
|
||||
//! inbound pipeline, and the outbound encoders can all share the
|
||||
//! same `InboundMessage` / `MsgBodyElement` shapes without re-mapping.
|
||||
//!
|
||||
//! Source of truth: `gateway/platforms/yuanbao_proto.py` in
|
||||
//! hermes-agent (lines 415-705).
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Decoded ConnMsg envelope (head + payload).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ConnFrame {
|
||||
/// CmdType (`CMD_TYPE`): Request=0, Response=1, Push=2, PushAck=3.
|
||||
pub cmd_type: u32,
|
||||
/// Command word, e.g. `"auth-bind"`, `"ping"`, `"send_c2c_message"`.
|
||||
pub cmd: String,
|
||||
/// Module / service name, e.g. `"conn_access"` or `"yuanbao_openclaw_proxy"`.
|
||||
pub module: String,
|
||||
/// Per-message sequence number.
|
||||
pub seq_no: u32,
|
||||
/// Application-level message id.
|
||||
pub msg_id: String,
|
||||
/// Whether the server expects an ACK.
|
||||
pub need_ack: bool,
|
||||
/// Status code (head.status, field 10).
|
||||
pub status: u32,
|
||||
/// Biz payload bytes (ConnMsg.data, field 2).
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
/// One element of the TIM-style `msg_body` array (e.g. text, image, file).
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct MsgBodyElement {
|
||||
/// `"TIMTextElem"`, `"TIMImageElem"`, `"TIMFileElem"`, `"TIMSoundElem"`, …
|
||||
pub msg_type: String,
|
||||
pub msg_content: MsgContent,
|
||||
}
|
||||
|
||||
/// Generic union of all TIM `msg_content` shapes (text/image/file/sound).
|
||||
///
|
||||
/// Only the fields relevant to the active `msg_type` are populated; the
|
||||
/// rest stay at their `Default`.
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct MsgContent {
|
||||
/// Field 1 — text payload.
|
||||
pub text: Option<String>,
|
||||
/// Field 2 — file uuid (MD5 for images/files).
|
||||
pub uuid: Option<String>,
|
||||
/// Field 3 — image format code (1=JPEG, 2=GIF, 3=PNG, 4=BMP, 255=WEBP).
|
||||
pub image_format: Option<u32>,
|
||||
/// Field 4 — raw inline data (rarely used; usually `url` is set instead).
|
||||
pub data: Option<String>,
|
||||
/// Field 5 — element description.
|
||||
pub desc: Option<String>,
|
||||
/// Field 6 — extension JSON / blob.
|
||||
pub ext: Option<String>,
|
||||
/// Field 7 — voice payload identifier.
|
||||
pub sound: Option<String>,
|
||||
/// Field 8 — repeated `ImageInfo` for the image element.
|
||||
pub image_info_array: Vec<ImageInfo>,
|
||||
/// Field 9 — element index within a multi-image message.
|
||||
pub index: Option<u32>,
|
||||
/// Field 10 — resource URL.
|
||||
pub url: Option<String>,
|
||||
/// Field 11 — file size in bytes.
|
||||
pub file_size: Option<u32>,
|
||||
/// Field 12 — file name.
|
||||
pub file_name: Option<String>,
|
||||
}
|
||||
|
||||
/// Per-resolution image variant. `type` is 1=original, 2=large, 3=thumb.
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct ImageInfo {
|
||||
pub image_type: u32,
|
||||
pub size: u32,
|
||||
pub width: u32,
|
||||
pub height: u32,
|
||||
pub url: String,
|
||||
}
|
||||
|
||||
/// A single recall entry in `recall_msg_seq_list` (InboundMessagePush field 17).
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct ImMsgSeq {
|
||||
pub msg_seq: u32,
|
||||
pub msg_id: String,
|
||||
}
|
||||
|
||||
/// A decoded `InboundMessagePush` biz payload — what the yuanbao gateway
|
||||
/// pushes down to us for every incoming message.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct InboundMessage {
|
||||
pub callback_command: String,
|
||||
pub from_account: String,
|
||||
pub to_account: String,
|
||||
pub sender_nickname: String,
|
||||
/// Empty string for DMs, group ID for group messages.
|
||||
pub group_id: String,
|
||||
/// Empty string for DMs, group code (canonical group ref) for group messages.
|
||||
pub group_code: String,
|
||||
pub group_name: String,
|
||||
pub msg_seq: u32,
|
||||
pub msg_random: u32,
|
||||
/// Server-side message timestamp (seconds since epoch).
|
||||
pub msg_time: u32,
|
||||
pub msg_key: String,
|
||||
/// Stable application-level message ID.
|
||||
pub msg_id: String,
|
||||
pub msg_body: Vec<MsgBodyElement>,
|
||||
pub cloud_custom_data: String,
|
||||
pub event_time: u32,
|
||||
pub bot_owner_id: String,
|
||||
pub recall_msg_seq_list: Vec<ImMsgSeq>,
|
||||
pub claw_msg_type: u32,
|
||||
pub private_from_group_code: String,
|
||||
pub trace_id: String,
|
||||
}
|
||||
|
||||
impl InboundMessage {
|
||||
/// Whether this is a group message.
|
||||
pub fn is_group(&self) -> bool {
|
||||
!self.group_code.is_empty()
|
||||
}
|
||||
|
||||
/// Whether the message looks like a recall notification.
|
||||
pub fn is_recall(&self) -> bool {
|
||||
!self.recall_msg_seq_list.is_empty()
|
||||
}
|
||||
|
||||
/// Routing key — group_code for groups, sender uid for DMs.
|
||||
pub fn chat_id(&self) -> &str {
|
||||
if self.is_group() {
|
||||
&self.group_code
|
||||
} else {
|
||||
&self.from_account
|
||||
}
|
||||
}
|
||||
|
||||
/// Concatenated text content (joins all `TIMTextElem`s).
|
||||
pub fn extract_text(&self) -> String {
|
||||
let mut out = String::new();
|
||||
for el in &self.msg_body {
|
||||
if el.msg_type == "TIMTextElem" {
|
||||
if let Some(ref t) = el.msg_content.text {
|
||||
if !out.is_empty() {
|
||||
out.push('\n');
|
||||
}
|
||||
out.push_str(t);
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// All image URLs in the message (from `TIMImageElem` elements).
|
||||
pub fn extract_image_urls(&self) -> Vec<String> {
|
||||
let mut urls = Vec::new();
|
||||
for el in &self.msg_body {
|
||||
if el.msg_type == "TIMImageElem" {
|
||||
for info in &el.msg_content.image_info_array {
|
||||
if !info.url.is_empty() {
|
||||
urls.push(info.url.clone());
|
||||
}
|
||||
}
|
||||
if let Some(ref url) = el.msg_content.url {
|
||||
if !url.is_empty() && !urls.contains(url) {
|
||||
urls.push(url.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
urls
|
||||
}
|
||||
}
|
||||
|
||||
/// High-level classification produced by the inbound pipeline.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum MessageKind {
|
||||
#[default]
|
||||
Text,
|
||||
Image,
|
||||
File,
|
||||
Voice,
|
||||
Mixed,
|
||||
/// Recall notification — handled by `RecallGuard`, never dispatched.
|
||||
Recall,
|
||||
}
|
||||
|
||||
/// Where the message came from — used by the outbound side to address replies.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct Source {
|
||||
pub from_account: String,
|
||||
pub sender_nickname: String,
|
||||
pub group_code: String,
|
||||
/// `true` for group chats, `false` for DMs.
|
||||
pub is_group: bool,
|
||||
}
|
||||
|
||||
impl Source {
|
||||
/// Stable string for `ChannelMessage.sender` / `reply_target` —
|
||||
/// `g:<group_code>` for groups, raw uid for DMs. This format also
|
||||
/// round-trips through `parse_recipient` in `outbound.rs`.
|
||||
pub fn reply_target(&self) -> String {
|
||||
if self.is_group {
|
||||
format!("g:{}", self.group_code)
|
||||
} else {
|
||||
self.from_account.clone()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn text_elem(s: &str) -> MsgBodyElement {
|
||||
MsgBodyElement {
|
||||
msg_type: "TIMTextElem".into(),
|
||||
msg_content: MsgContent {
|
||||
text: Some(s.into()),
|
||||
..Default::default()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn image_elem(info_urls: &[&str], inline_url: Option<&str>) -> MsgBodyElement {
|
||||
MsgBodyElement {
|
||||
msg_type: "TIMImageElem".into(),
|
||||
msg_content: MsgContent {
|
||||
image_info_array: info_urls
|
||||
.iter()
|
||||
.map(|u| ImageInfo {
|
||||
image_type: 1,
|
||||
url: (*u).into(),
|
||||
..Default::default()
|
||||
})
|
||||
.collect(),
|
||||
url: inline_url.map(String::from),
|
||||
..Default::default()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dm_is_not_group() {
|
||||
let m = InboundMessage {
|
||||
from_account: "alice".into(),
|
||||
..Default::default()
|
||||
};
|
||||
assert!(!m.is_group());
|
||||
assert_eq!(m.chat_id(), "alice");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn group_is_group_and_chat_id_is_group_code() {
|
||||
let m = InboundMessage {
|
||||
group_code: "grp_42".into(),
|
||||
from_account: "alice".into(),
|
||||
..Default::default()
|
||||
};
|
||||
assert!(m.is_group());
|
||||
assert_eq!(m.chat_id(), "grp_42");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_recall_iff_recall_list_non_empty() {
|
||||
let mut m = InboundMessage::default();
|
||||
assert!(!m.is_recall());
|
||||
m.recall_msg_seq_list.push(ImMsgSeq {
|
||||
msg_seq: 7,
|
||||
msg_id: "x".into(),
|
||||
});
|
||||
assert!(m.is_recall());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_text_concatenates_text_elements() {
|
||||
let m = InboundMessage {
|
||||
msg_body: vec![
|
||||
text_elem("hello"),
|
||||
text_elem("world"),
|
||||
image_elem(&[], None),
|
||||
],
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(m.extract_text(), "hello\nworld");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_text_ignores_text_none_and_non_text() {
|
||||
let m = InboundMessage {
|
||||
msg_body: vec![
|
||||
MsgBodyElement {
|
||||
msg_type: "TIMTextElem".into(),
|
||||
msg_content: MsgContent::default(), // text: None
|
||||
},
|
||||
image_elem(&["https://x/y.png"], None),
|
||||
],
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(m.extract_text(), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_text_on_empty_msg_body_returns_empty() {
|
||||
let m = InboundMessage::default();
|
||||
assert_eq!(m.extract_text(), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_image_urls_from_image_info_array() {
|
||||
let m = InboundMessage {
|
||||
msg_body: vec![image_elem(&["https://a/1.png", "https://a/2.png"], None)],
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(
|
||||
m.extract_image_urls(),
|
||||
vec!["https://a/1.png".to_string(), "https://a/2.png".into()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_image_urls_falls_back_to_inline_url_field() {
|
||||
let m = InboundMessage {
|
||||
msg_body: vec![image_elem(&[], Some("https://a/inline.png"))],
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(
|
||||
m.extract_image_urls(),
|
||||
vec!["https://a/inline.png".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_image_urls_dedups_inline_when_already_in_info_array() {
|
||||
let dup = "https://a/dup.png";
|
||||
let m = InboundMessage {
|
||||
msg_body: vec![image_elem(&[dup], Some(dup))],
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(m.extract_image_urls(), vec![dup.to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_image_urls_skips_empty_url_in_info_array() {
|
||||
let m = InboundMessage {
|
||||
msg_body: vec![image_elem(&[""], None)],
|
||||
..Default::default()
|
||||
};
|
||||
assert!(m.extract_image_urls().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_image_urls_ignores_text_elements() {
|
||||
let m = InboundMessage {
|
||||
msg_body: vec![text_elem("hi"), image_elem(&["https://a/1.png"], None)],
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(m.extract_image_urls(), vec!["https://a/1.png".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn source_reply_target_dm_is_raw_uid() {
|
||||
let s = Source {
|
||||
from_account: "uid_alice".into(),
|
||||
is_group: false,
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(s.reply_target(), "uid_alice");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn source_reply_target_group_uses_g_prefix() {
|
||||
let s = Source {
|
||||
group_code: "grp_42".into(),
|
||||
is_group: true,
|
||||
..Default::default()
|
||||
};
|
||||
assert_eq!(s.reply_target(), "g:grp_42");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn message_kind_default_is_text() {
|
||||
assert_eq!(MessageKind::default(), MessageKind::Text);
|
||||
}
|
||||
}
|
||||
|
||||
/// Group metadata returned by `QueryGroupInfo`.
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct GroupInfo {
|
||||
pub code: i32,
|
||||
pub message: String,
|
||||
pub group_name: String,
|
||||
pub owner_id: String,
|
||||
pub owner_nickname: String,
|
||||
pub member_count: u32,
|
||||
}
|
||||
|
||||
/// One member returned by `GetGroupMemberList`.
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct GroupMember {
|
||||
pub user_id: String,
|
||||
pub nickname: String,
|
||||
/// 0=member, 1=admin, 2=owner.
|
||||
pub role: u32,
|
||||
pub join_time: u32,
|
||||
pub name_card: String,
|
||||
}
|
||||
|
||||
/// Paginated result of `GetGroupMemberList`.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct GroupMemberListPage {
|
||||
pub code: i32,
|
||||
pub message: String,
|
||||
pub members: Vec<GroupMember>,
|
||||
pub next_offset: u32,
|
||||
pub is_complete: bool,
|
||||
}
|
||||
|
||||
/// Cached account info — populated after `auth-bind` succeeds.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct Account {
|
||||
/// Bot user-id (used as `from_account` in outbound messages).
|
||||
pub uid: String,
|
||||
/// Display name (best-effort; may be empty until first inbound message).
|
||||
pub nickname: String,
|
||||
/// Server-assigned connection id (`AuthBindRsp.connect_id`, field 3).
|
||||
pub connect_id: String,
|
||||
}
|
||||
|
||||
/// Connection state machine (matches task list spec).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ConnectionState {
|
||||
Disconnected,
|
||||
Connecting,
|
||||
Authenticating,
|
||||
Connected,
|
||||
Reconnecting,
|
||||
}
|
||||
@@ -1,397 +0,0 @@
|
||||
//! Hand-rolled protobuf wire-format primitives.
|
||||
//!
|
||||
//! Only varints, length-delimited bytes, and the two fixed-width forms
|
||||
//! are supported — that's everything the yuanbao protocol uses. Kept
|
||||
//! separate from `proto.rs` so the latter stays under 500 lines and
|
||||
//! reads as a "schema" file.
|
||||
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
|
||||
use super::errors::YuanbaoError;
|
||||
|
||||
/// Global per-process sequence number for ConnMsg head.seq_no.
|
||||
static SEQ: AtomicU32 = AtomicU32::new(1);
|
||||
|
||||
pub fn next_seq_no() -> u32 {
|
||||
SEQ.fetch_add(1, Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub const WT_VARINT: u8 = 0;
|
||||
pub const WT_LEN: u8 = 2;
|
||||
|
||||
// ─── Varint ─────────────────────────────────────────────────────────
|
||||
|
||||
pub fn encode_varint(mut value: u64, buf: &mut Vec<u8>) {
|
||||
loop {
|
||||
let mut byte = (value & 0x7F) as u8;
|
||||
value >>= 7;
|
||||
if value != 0 {
|
||||
byte |= 0x80;
|
||||
}
|
||||
buf.push(byte);
|
||||
if value == 0 {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn decode_varint(data: &[u8], pos: usize) -> Result<(u64, usize), YuanbaoError> {
|
||||
let mut value: u64 = 0;
|
||||
let mut shift: u32 = 0;
|
||||
let mut i = pos;
|
||||
loop {
|
||||
if i >= data.len() {
|
||||
return Err(YuanbaoError::ProtoDecode("truncated varint".into()));
|
||||
}
|
||||
let byte = data[i];
|
||||
// On the 10th byte (shift == 63) a valid u64 varint can only have
|
||||
// the lowest bit set (values 0 or 1); anything higher overflows.
|
||||
if shift == 63 && byte > 1 {
|
||||
return Err(YuanbaoError::ProtoDecode(format!(
|
||||
"varint overflow: 10th byte is {byte:#04x}, expected 0x00 or 0x01"
|
||||
)));
|
||||
}
|
||||
value |= ((byte & 0x7F) as u64) << shift;
|
||||
i += 1;
|
||||
if byte & 0x80 == 0 {
|
||||
return Ok((value, i - pos));
|
||||
}
|
||||
shift += 7;
|
||||
if shift >= 64 {
|
||||
return Err(YuanbaoError::ProtoDecode("varint too long".into()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── Field encoders ────────────────────────────────────────────────
|
||||
|
||||
pub fn encode_field_varint(field: u32, value: u64, buf: &mut Vec<u8>) {
|
||||
encode_varint(((field as u64) << 3) | WT_VARINT as u64, buf);
|
||||
encode_varint(value, buf);
|
||||
}
|
||||
|
||||
pub fn encode_field_bytes(field: u32, data: &[u8], buf: &mut Vec<u8>) {
|
||||
encode_varint(((field as u64) << 3) | WT_LEN as u64, buf);
|
||||
encode_varint(data.len() as u64, buf);
|
||||
buf.extend_from_slice(data);
|
||||
}
|
||||
|
||||
pub fn encode_field_string(field: u32, s: &str, buf: &mut Vec<u8>) {
|
||||
encode_field_bytes(field, s.as_bytes(), buf);
|
||||
}
|
||||
|
||||
// ─── Field parsing ──────────────────────────────────────────────────
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum FieldValue {
|
||||
Varint(u64),
|
||||
Bytes(Vec<u8>),
|
||||
Fixed32(u32),
|
||||
Fixed64(u64),
|
||||
}
|
||||
|
||||
pub fn parse_fields(data: &[u8]) -> Result<Vec<(u32, FieldValue)>, YuanbaoError> {
|
||||
let mut out = Vec::new();
|
||||
let mut pos = 0;
|
||||
while pos < data.len() {
|
||||
let (tag, n) = decode_varint(data, pos)?;
|
||||
pos += n;
|
||||
let field = (tag >> 3) as u32;
|
||||
let wire = (tag & 0x07) as u8;
|
||||
match wire {
|
||||
WT_VARINT => {
|
||||
let (v, n) = decode_varint(data, pos)?;
|
||||
pos += n;
|
||||
out.push((field, FieldValue::Varint(v)));
|
||||
}
|
||||
WT_LEN => {
|
||||
let (len, n) = decode_varint(data, pos)?;
|
||||
pos += n;
|
||||
// Use checked conversions / arithmetic — a crafted oversize
|
||||
// varint length would otherwise overflow `usize` on 32-bit
|
||||
// targets and panic during slicing.
|
||||
let len_usize = usize::try_from(len).map_err(|_| {
|
||||
YuanbaoError::ProtoDecode(format!(
|
||||
"len field {field} too large for platform: {len}"
|
||||
))
|
||||
})?;
|
||||
let end = pos.checked_add(len_usize).ok_or_else(|| {
|
||||
YuanbaoError::ProtoDecode(format!(
|
||||
"len field {field} overflows position: pos={pos} len={len}"
|
||||
))
|
||||
})?;
|
||||
if end > data.len() {
|
||||
return Err(YuanbaoError::ProtoDecode(format!(
|
||||
"truncated len field {field}: need {len} have {}",
|
||||
data.len() - pos
|
||||
)));
|
||||
}
|
||||
out.push((field, FieldValue::Bytes(data[pos..end].to_vec())));
|
||||
pos = end;
|
||||
}
|
||||
1 => {
|
||||
if pos + 8 > data.len() {
|
||||
return Err(YuanbaoError::ProtoDecode("truncated fixed64".into()));
|
||||
}
|
||||
let v = u64::from_le_bytes(data[pos..pos + 8].try_into().unwrap());
|
||||
pos += 8;
|
||||
out.push((field, FieldValue::Fixed64(v)));
|
||||
}
|
||||
5 => {
|
||||
if pos + 4 > data.len() {
|
||||
return Err(YuanbaoError::ProtoDecode("truncated fixed32".into()));
|
||||
}
|
||||
let v = u32::from_le_bytes(data[pos..pos + 4].try_into().unwrap());
|
||||
pos += 4;
|
||||
out.push((field, FieldValue::Fixed32(v)));
|
||||
}
|
||||
other => {
|
||||
return Err(YuanbaoError::ProtoDecode(format!(
|
||||
"unsupported wire type {other} at field {field}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
pub fn get_string(fields: &[(u32, FieldValue)], num: u32) -> String {
|
||||
for (n, v) in fields {
|
||||
if *n == num {
|
||||
if let FieldValue::Bytes(b) = v {
|
||||
return String::from_utf8_lossy(b).into_owned();
|
||||
}
|
||||
}
|
||||
}
|
||||
String::new()
|
||||
}
|
||||
|
||||
pub fn get_varint(fields: &[(u32, FieldValue)], num: u32) -> u64 {
|
||||
for (n, v) in fields {
|
||||
if *n == num {
|
||||
if let FieldValue::Varint(x) = v {
|
||||
return *x;
|
||||
}
|
||||
}
|
||||
}
|
||||
0
|
||||
}
|
||||
|
||||
pub fn get_bytes(fields: &[(u32, FieldValue)], num: u32) -> Vec<u8> {
|
||||
for (n, v) in fields {
|
||||
if *n == num {
|
||||
if let FieldValue::Bytes(b) = v {
|
||||
return b.clone();
|
||||
}
|
||||
}
|
||||
}
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
pub fn get_repeated_bytes(fields: &[(u32, FieldValue)], num: u32) -> Vec<Vec<u8>> {
|
||||
fields
|
||||
.iter()
|
||||
.filter_map(|(n, v)| match v {
|
||||
FieldValue::Bytes(b) if *n == num => Some(b.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn varint_roundtrip() {
|
||||
for &v in &[0u64, 1, 127, 128, 300, 16384, u32::MAX as u64, u64::MAX] {
|
||||
let mut buf = Vec::new();
|
||||
encode_varint(v, &mut buf);
|
||||
let (got, n) = decode_varint(&buf, 0).unwrap();
|
||||
assert_eq!(got, v, "varint roundtrip failed for {v}");
|
||||
assert_eq!(n, buf.len());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn varint_truncated_errors() {
|
||||
let buf = vec![0x80, 0x80]; // continuation bit set but no end
|
||||
assert!(decode_varint(&buf, 0).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn field_roundtrip() {
|
||||
let mut buf = Vec::new();
|
||||
encode_field_varint(1, 42, &mut buf);
|
||||
encode_field_string(2, "hello", &mut buf);
|
||||
encode_field_bytes(3, b"\x00\x01\x02", &mut buf);
|
||||
|
||||
let fields = parse_fields(&buf).unwrap();
|
||||
assert_eq!(get_varint(&fields, 1), 42);
|
||||
assert_eq!(get_string(&fields, 2), "hello");
|
||||
assert_eq!(get_bytes(&fields, 3), vec![0, 1, 2]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_field_skipped_gracefully() {
|
||||
let mut buf = Vec::new();
|
||||
encode_field_varint(99, 123, &mut buf);
|
||||
encode_field_string(1, "wanted", &mut buf);
|
||||
let fields = parse_fields(&buf).unwrap();
|
||||
assert_eq!(get_string(&fields, 1), "wanted");
|
||||
assert_eq!(get_string(&fields, 2), ""); // missing field returns default
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seq_numbers_are_monotonic() {
|
||||
let a = next_seq_no();
|
||||
let b = next_seq_no();
|
||||
assert!(b > a);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn varint_too_long_errors() {
|
||||
// 11 continuation bytes — the 10th byte (0x80) triggers the overflow
|
||||
// guard before the loop even reaches the 11th.
|
||||
let buf = vec![0x80; 11];
|
||||
match decode_varint(&buf, 0).unwrap_err() {
|
||||
YuanbaoError::ProtoDecode(m) => {
|
||||
assert!(m.contains("too long") || m.contains("overflow"), "got {m}");
|
||||
}
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fields_truncated_bytes_field_errors() {
|
||||
// Field 1 (wire type 2) declaring length 5 but only 1 byte of payload.
|
||||
let buf = vec![
|
||||
(1 << 3) | 2, // tag: field=1, wire=2
|
||||
5, // claimed len
|
||||
b'a',
|
||||
];
|
||||
match parse_fields(&buf).unwrap_err() {
|
||||
YuanbaoError::ProtoDecode(m) => assert!(m.contains("truncated"), "got {m}"),
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fields_oversize_len_field_errors_without_panic() {
|
||||
// Field 1 (wire type 2) with a varint length encoding `u64::MAX` —
|
||||
// previously this would attempt `pos + len as usize`, overflowing
|
||||
// on 32-bit and slicing past the buffer on 64-bit. Now it must
|
||||
// return a structured decode error.
|
||||
let mut buf = Vec::new();
|
||||
buf.push((1 << 3) | 2); // tag: field=1, wire=2
|
||||
encode_varint(u64::MAX, &mut buf); // adversarial length
|
||||
match parse_fields(&buf).unwrap_err() {
|
||||
YuanbaoError::ProtoDecode(m) => {
|
||||
assert!(
|
||||
m.contains("too large") || m.contains("overflows") || m.contains("truncated"),
|
||||
"expected overflow/truncation error, got {m}"
|
||||
);
|
||||
}
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fields_reads_fixed64() {
|
||||
let mut buf = Vec::new();
|
||||
buf.push((1 << 3) | 1); // tag: field=1, wire=1 (fixed64)
|
||||
buf.extend_from_slice(&0x1122_3344_5566_7788u64.to_le_bytes());
|
||||
let f = parse_fields(&buf).unwrap();
|
||||
match f[0].1 {
|
||||
FieldValue::Fixed64(v) => assert_eq!(v, 0x1122_3344_5566_7788),
|
||||
ref other => panic!("expected Fixed64 got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fields_truncated_fixed64_errors() {
|
||||
let mut buf = Vec::new();
|
||||
buf.push((1 << 3) | 1);
|
||||
buf.extend_from_slice(&[0u8; 4]); // only 4/8 bytes
|
||||
match parse_fields(&buf).unwrap_err() {
|
||||
YuanbaoError::ProtoDecode(m) => assert!(m.contains("fixed64"), "got {m}"),
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fields_reads_fixed32() {
|
||||
let mut buf = Vec::new();
|
||||
buf.push((1 << 3) | 5); // tag: field=1, wire=5 (fixed32)
|
||||
buf.extend_from_slice(&0xCAFEBABEu32.to_le_bytes());
|
||||
let f = parse_fields(&buf).unwrap();
|
||||
match f[0].1 {
|
||||
FieldValue::Fixed32(v) => assert_eq!(v, 0xCAFEBABE),
|
||||
ref other => panic!("expected Fixed32 got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fields_truncated_fixed32_errors() {
|
||||
let mut buf = Vec::new();
|
||||
buf.push((1 << 3) | 5);
|
||||
buf.extend_from_slice(&[0u8; 2]); // only 2/4 bytes
|
||||
match parse_fields(&buf).unwrap_err() {
|
||||
YuanbaoError::ProtoDecode(m) => assert!(m.contains("fixed32"), "got {m}"),
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fields_unsupported_wire_type_errors() {
|
||||
// wire type 3 (start group) is not supported.
|
||||
let buf = vec![(1 << 3) | 3];
|
||||
match parse_fields(&buf).unwrap_err() {
|
||||
YuanbaoError::ProtoDecode(m) => {
|
||||
assert!(m.contains("unsupported wire type 3"), "got {m}")
|
||||
}
|
||||
other => panic!("unexpected {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_string_returns_empty_when_field_is_varint() {
|
||||
// Field 1 exists but encoded as varint, not bytes — get_string must
|
||||
// skip past it and return the default.
|
||||
let mut buf = Vec::new();
|
||||
encode_field_varint(1, 7, &mut buf);
|
||||
let fields = parse_fields(&buf).unwrap();
|
||||
assert_eq!(get_string(&fields, 1), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_varint_returns_zero_when_field_is_bytes() {
|
||||
let mut buf = Vec::new();
|
||||
encode_field_string(1, "not a varint", &mut buf);
|
||||
let fields = parse_fields(&buf).unwrap();
|
||||
assert_eq!(get_varint(&fields, 1), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_bytes_returns_empty_when_field_is_varint() {
|
||||
let mut buf = Vec::new();
|
||||
encode_field_varint(1, 7, &mut buf);
|
||||
let fields = parse_fields(&buf).unwrap();
|
||||
assert!(get_bytes(&fields, 1).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_repeated_bytes_collects_multiple_same_field() {
|
||||
let mut buf = Vec::new();
|
||||
encode_field_string(1, "a", &mut buf);
|
||||
encode_field_string(1, "bb", &mut buf);
|
||||
encode_field_string(2, "c", &mut buf); // different field — should be skipped
|
||||
encode_field_string(1, "ddd", &mut buf);
|
||||
let fields = parse_fields(&buf).unwrap();
|
||||
let got = get_repeated_bytes(&fields, 1);
|
||||
assert_eq!(got.len(), 3);
|
||||
assert_eq!(got[0], b"a");
|
||||
assert_eq!(got[1], b"bb");
|
||||
assert_eq!(got[2], b"ddd");
|
||||
}
|
||||
}
|
||||
@@ -532,6 +532,12 @@ pub async fn start_channels(mut config: Config) -> Result<()> {
|
||||
);
|
||||
}
|
||||
|
||||
// Assemble the ChannelHost capability surface (shutdown, STT/TTS, reaction
|
||||
// gate, approvals, conversation store, event sink). Ported rich providers
|
||||
// reach host capabilities through this instead of calling core internals.
|
||||
let channel_host =
|
||||
crate::openhuman::channels::host::build_channel_host(Arc::new(config.clone()));
|
||||
|
||||
// Collect active channels
|
||||
let mut channels: Vec<Arc<dyn Channel>> = Vec::new();
|
||||
|
||||
@@ -544,19 +550,32 @@ pub async fn start_channels(mut config: Config) -> Result<()> {
|
||||
draft_update_interval_ms = tg.draft_update_interval_ms,
|
||||
"[channels] telegram enabled in core config (bot token not logged)"
|
||||
);
|
||||
channels.push(Arc::new(
|
||||
TelegramChannel::new(
|
||||
tg.bot_token.clone(),
|
||||
tg.allowed_users.clone(),
|
||||
tg.mention_only,
|
||||
)
|
||||
.with_streaming(
|
||||
tg.stream_mode,
|
||||
tg.draft_update_interval_ms,
|
||||
tg.silent_streaming,
|
||||
)
|
||||
.with_chat_id(tg.chat_id.clone()),
|
||||
let mut telegram = TelegramChannel::new(
|
||||
tg.bot_token.clone(),
|
||||
tg.allowed_users.clone(),
|
||||
tg.mention_only,
|
||||
)
|
||||
.with_streaming(
|
||||
tg.stream_mode,
|
||||
tg.draft_update_interval_ms,
|
||||
tg.silent_streaming,
|
||||
)
|
||||
.with_chat_id(tg.chat_id.clone())
|
||||
.with_http_client(crate::openhuman::config::build_runtime_proxy_client(
|
||||
"channel.telegram",
|
||||
));
|
||||
// Inject host capabilities: voice STT, persisted allowlist, reaction
|
||||
// event fan-out. Each is optional — telegram degrades gracefully.
|
||||
if let Some(transcriber) = channel_host.transcriber() {
|
||||
telegram = telegram.with_transcriber(transcriber);
|
||||
}
|
||||
if let Some(allowlist) = channel_host.allowlist() {
|
||||
telegram = telegram.with_allowlist(allowlist);
|
||||
}
|
||||
if let Some(events) = channel_host.events() {
|
||||
telegram = telegram.with_events(events);
|
||||
}
|
||||
channels.push(Arc::new(telegram));
|
||||
} else {
|
||||
tracing::info!(
|
||||
"[channels] telegram not configured (no channels_config.telegram in saved config)"
|
||||
@@ -564,13 +583,14 @@ pub async fn start_channels(mut config: Config) -> Result<()> {
|
||||
}
|
||||
|
||||
if let Some(ref dc) = config.channels_config.discord {
|
||||
channels.push(Arc::new(DiscordChannel::new(
|
||||
channels.push(Arc::new(DiscordChannel::with_http_client(
|
||||
dc.bot_token.clone(),
|
||||
dc.guild_id.clone(),
|
||||
dc.channel_id.clone(),
|
||||
dc.allowed_users.clone(),
|
||||
dc.listen_to_bots,
|
||||
dc.mention_only,
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.discord"),
|
||||
)));
|
||||
}
|
||||
|
||||
@@ -588,13 +608,14 @@ pub async fn start_channels(mut config: Config) -> Result<()> {
|
||||
}
|
||||
|
||||
if let Some(ref mm) = config.channels_config.mattermost {
|
||||
channels.push(Arc::new(MattermostChannel::new(
|
||||
channels.push(Arc::new(MattermostChannel::with_http_client(
|
||||
mm.url.clone(),
|
||||
mm.bot_token.clone(),
|
||||
mm.channel_id.clone(),
|
||||
mm.allowed_users.clone(),
|
||||
mm.thread_replies.unwrap_or(true),
|
||||
mm.mention_only.unwrap_or(false),
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.mattermost"),
|
||||
)));
|
||||
}
|
||||
|
||||
@@ -646,12 +667,16 @@ pub async fn start_channels(mut config: Config) -> Result<()> {
|
||||
// Web mode: requires session_path
|
||||
#[cfg(feature = "whatsapp-web")]
|
||||
if wa.is_web_config() {
|
||||
channels.push(Arc::new(WhatsAppWebChannel::new(
|
||||
let mut wa_channel = WhatsAppWebChannel::new(
|
||||
wa.session_path.clone().unwrap_or_default(),
|
||||
wa.pair_phone.clone(),
|
||||
wa.pair_code.clone(),
|
||||
wa.allowed_numbers.clone(),
|
||||
)));
|
||||
);
|
||||
if let Some(lifecycle) = channel_host.lifecycle() {
|
||||
wa_channel = wa_channel.with_lifecycle(lifecycle);
|
||||
}
|
||||
channels.push(Arc::new(wa_channel));
|
||||
} else {
|
||||
tracing::warn!("WhatsApp Web configured but session_path not set");
|
||||
}
|
||||
@@ -699,18 +724,20 @@ pub async fn start_channels(mut config: Config) -> Result<()> {
|
||||
}
|
||||
|
||||
if let Some(ref dt) = config.channels_config.dingtalk {
|
||||
channels.push(Arc::new(DingTalkChannel::new(
|
||||
channels.push(Arc::new(DingTalkChannel::with_http_client(
|
||||
dt.client_id.clone(),
|
||||
dt.client_secret.clone(),
|
||||
dt.allowed_users.clone(),
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.dingtalk"),
|
||||
)));
|
||||
}
|
||||
|
||||
if let Some(ref qq) = config.channels_config.qq {
|
||||
channels.push(Arc::new(QQChannel::new(
|
||||
channels.push(Arc::new(QQChannel::with_http_client(
|
||||
qq.app_id.clone(),
|
||||
qq.app_secret.clone(),
|
||||
qq.allowed_users.clone(),
|
||||
crate::openhuman::config::build_runtime_proxy_client("channel.qq"),
|
||||
)));
|
||||
}
|
||||
|
||||
|
||||
@@ -4,295 +4,23 @@
|
||||
// a bearer token. Tokens can be persisted in config so restarts don't require
|
||||
// re-pairing.
|
||||
|
||||
use parking_lot::Mutex;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::collections::HashSet;
|
||||
use std::io::Write as _;
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
#[cfg(unix)]
|
||||
use std::os::unix::fs::{OpenOptionsExt as _, PermissionsExt as _};
|
||||
use std::os::unix::fs::OpenOptionsExt as _;
|
||||
|
||||
/// `PairingGuard`, `constant_time_eq`, and the pairing lockout constants now
|
||||
/// live in tinychannels (portable, self-contained); re-export them so existing
|
||||
/// callers (channels, `core::auth`, tests) keep their paths.
|
||||
pub use tinychannels::security::{
|
||||
constant_time_eq, generate_code, generate_token, hash_token, is_token_hash, PairingGuard,
|
||||
MAX_PAIR_ATTEMPTS, PAIR_LOCKOUT_SECS,
|
||||
};
|
||||
|
||||
/// Environment variable for the core JSON-RPC bearer token (see `crate::core::auth`).
|
||||
pub const CORE_TOKEN_ENV_VAR: &str = "OPENHUMAN_CORE_TOKEN";
|
||||
|
||||
/// Maximum failed pairing attempts before lockout.
|
||||
const MAX_PAIR_ATTEMPTS: u32 = 5;
|
||||
/// Lockout duration after too many failed pairing attempts.
|
||||
const PAIR_LOCKOUT_SECS: u64 = 300; // 5 minutes
|
||||
|
||||
/// Manages pairing state for channels that use bearer-token auth after pairing.
|
||||
///
|
||||
/// Bearer tokens are stored as SHA-256 hashes to prevent plaintext exposure
|
||||
/// in config files. When a new token is generated, the plaintext is returned
|
||||
/// to the client once, and only the hash is retained.
|
||||
// TODO: I've just made this work with parking_lot but it should use either flume or tokio's async mutexes
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PairingGuard {
|
||||
/// Whether pairing is required at all.
|
||||
require_pairing: bool,
|
||||
/// One-time pairing code (generated on startup, consumed on first pair).
|
||||
pairing_code: Arc<Mutex<Option<String>>>,
|
||||
/// Set of SHA-256 hashed bearer tokens (persisted across restarts).
|
||||
paired_tokens: Arc<Mutex<HashSet<String>>>,
|
||||
/// Brute-force protection: failed attempt counter + lockout time.
|
||||
failed_attempts: Arc<Mutex<(u32, Option<Instant>)>>,
|
||||
}
|
||||
|
||||
impl PairingGuard {
|
||||
/// Create a new pairing guard.
|
||||
///
|
||||
/// If `require_pairing` is true and no tokens exist yet, a fresh
|
||||
/// pairing code is generated and returned via `pairing_code()`.
|
||||
///
|
||||
/// Existing tokens are accepted in both forms:
|
||||
/// - Plaintext (`zc_...`): hashed on load for backward compatibility
|
||||
/// - Already hashed (64-char hex): stored as-is
|
||||
pub fn new(require_pairing: bool, existing_tokens: &[String]) -> (Self, Option<String>) {
|
||||
let tokens: HashSet<String> = existing_tokens
|
||||
.iter()
|
||||
.map(|t| {
|
||||
if is_token_hash(t) {
|
||||
t.clone()
|
||||
} else {
|
||||
hash_token(t)
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let code = if require_pairing && tokens.is_empty() {
|
||||
Some(generate_code())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
log::info!(
|
||||
"[openhuman:pairing] Guard created: require_pairing={}, existing_tokens={}, code_generated={}",
|
||||
require_pairing,
|
||||
tokens.len(),
|
||||
code.is_some()
|
||||
);
|
||||
let guard = Self {
|
||||
require_pairing,
|
||||
pairing_code: Arc::new(Mutex::new(code.clone())),
|
||||
paired_tokens: Arc::new(Mutex::new(tokens)),
|
||||
failed_attempts: Arc::new(Mutex::new((0, None))),
|
||||
};
|
||||
(guard, code)
|
||||
}
|
||||
|
||||
/// The one-time pairing code (only set when no tokens exist yet).
|
||||
pub fn pairing_code(&self) -> Option<String> {
|
||||
self.pairing_code.lock().clone()
|
||||
}
|
||||
|
||||
/// Invalidate the one-time pairing code without pairing a token.
|
||||
///
|
||||
/// `try_pair` consumes the code as part of a successful `/bind <code>`, but
|
||||
/// some onboarding flows bind the operator through a different signal (e.g.
|
||||
/// the Telegram self-token `/start` path allowlists the first sender
|
||||
/// directly). Those paths must still finish the one-time flow so the code
|
||||
/// can't later be replayed by another sender who obtains it — call this once
|
||||
/// the operator is bound. Idempotent; a no-op if the code is already gone.
|
||||
pub fn invalidate_code(&self) {
|
||||
let mut code = self.pairing_code.lock();
|
||||
if code.is_some() {
|
||||
*code = None;
|
||||
log::info!(
|
||||
"[openhuman:pairing] one-time pairing code invalidated (operator onboarded)"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether pairing is required at all.
|
||||
pub fn require_pairing(&self) -> bool {
|
||||
self.require_pairing
|
||||
}
|
||||
|
||||
fn try_pair_blocking(&self, code: &str) -> Result<Option<String>, u64> {
|
||||
// Check brute force lockout
|
||||
{
|
||||
let attempts = self.failed_attempts.lock();
|
||||
if let (count, Some(locked_at)) = &*attempts {
|
||||
if *count >= MAX_PAIR_ATTEMPTS {
|
||||
let elapsed = locked_at.elapsed().as_secs();
|
||||
if elapsed < PAIR_LOCKOUT_SECS {
|
||||
log::warn!(
|
||||
"[openhuman:pairing] Pairing locked out: {} failed attempts, {}s remaining",
|
||||
count,
|
||||
PAIR_LOCKOUT_SECS - elapsed
|
||||
);
|
||||
return Err(PAIR_LOCKOUT_SECS - elapsed);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
let mut pairing_code = self.pairing_code.lock();
|
||||
if let Some(ref expected) = *pairing_code {
|
||||
if constant_time_eq(code.trim(), expected.trim()) {
|
||||
// Reset failed attempts on success
|
||||
{
|
||||
let mut attempts = self.failed_attempts.lock();
|
||||
*attempts = (0, None);
|
||||
}
|
||||
let token = generate_token();
|
||||
let mut tokens = self.paired_tokens.lock();
|
||||
tokens.insert(hash_token(&token));
|
||||
|
||||
// Consume the pairing code so it cannot be reused
|
||||
*pairing_code = None;
|
||||
|
||||
log::info!("[openhuman:pairing] Pairing successful, token issued");
|
||||
return Ok(Some(token));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Increment failed attempts
|
||||
{
|
||||
let mut attempts = self.failed_attempts.lock();
|
||||
attempts.0 += 1;
|
||||
log::warn!(
|
||||
"[openhuman:pairing] Pairing attempt failed ({}/{})",
|
||||
attempts.0,
|
||||
MAX_PAIR_ATTEMPTS
|
||||
);
|
||||
if attempts.0 >= MAX_PAIR_ATTEMPTS {
|
||||
attempts.1 = Some(Instant::now());
|
||||
log::warn!("[openhuman:pairing] Max attempts reached, lockout activated");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
/// Attempt to pair with the given code. Returns a bearer token on success.
|
||||
/// Returns `Err(lockout_seconds)` if locked out due to brute force.
|
||||
pub async fn try_pair(&self, code: &str) -> Result<Option<String>, u64> {
|
||||
let this = self.clone();
|
||||
let code = code.to_string();
|
||||
// TODO: make this function the main one without spawning a task
|
||||
let handle = tokio::task::spawn_blocking(move || this.try_pair_blocking(&code));
|
||||
|
||||
handle
|
||||
.await
|
||||
.expect("failed to spawn blocking task this should not happen")
|
||||
}
|
||||
|
||||
/// Check if a bearer token is valid (compares against stored hashes).
|
||||
///
|
||||
/// Always fails closed on empty/whitespace tokens. When pairing is not required,
|
||||
/// configured tokens are still honored if present; with no tokens configured,
|
||||
/// every request is rejected.
|
||||
pub fn is_authenticated(&self, token: &str) -> bool {
|
||||
if token.trim().is_empty() {
|
||||
log::debug!("[openhuman:pairing] is_authenticated: rejected empty bearer token");
|
||||
return false;
|
||||
}
|
||||
|
||||
let tokens = self.paired_tokens.lock();
|
||||
if tokens.is_empty() {
|
||||
log::debug!(
|
||||
"[openhuman:pairing] is_authenticated: no paired tokens configured (require_pairing={})",
|
||||
self.require_pairing
|
||||
);
|
||||
return false;
|
||||
}
|
||||
|
||||
let hashed = hash_token(token);
|
||||
let ok = tokens.contains(&hashed);
|
||||
if !ok {
|
||||
log::debug!("[openhuman:pairing] is_authenticated: bearer token not in paired set");
|
||||
}
|
||||
ok
|
||||
}
|
||||
|
||||
/// Returns true if pairing is satisfied (has at least one token).
|
||||
pub fn is_paired(&self) -> bool {
|
||||
let tokens = self.paired_tokens.lock();
|
||||
!tokens.is_empty()
|
||||
}
|
||||
|
||||
/// Get all paired token hashes (for persisting to config).
|
||||
pub fn tokens(&self) -> Vec<String> {
|
||||
let tokens = self.paired_tokens.lock();
|
||||
tokens.iter().cloned().collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate a 6-digit numeric pairing code using cryptographically secure randomness.
|
||||
fn generate_code() -> String {
|
||||
// UUID v4 uses getrandom (backed by /dev/urandom on Linux, BCryptGenRandom
|
||||
// on Windows) — a CSPRNG. We extract 4 bytes from it for a uniform random
|
||||
// number in [0, 1_000_000).
|
||||
//
|
||||
// Rejection sampling eliminates modulo bias: values above the largest
|
||||
// multiple of 1_000_000 that fits in u32 are discarded and re-drawn.
|
||||
// The rejection probability is ~0.02%, so this loop almost always exits
|
||||
// on the first iteration.
|
||||
const UPPER_BOUND: u32 = 1_000_000;
|
||||
const REJECT_THRESHOLD: u32 = (u32::MAX / UPPER_BOUND) * UPPER_BOUND;
|
||||
|
||||
loop {
|
||||
let uuid = uuid::Uuid::new_v4();
|
||||
let bytes = uuid.as_bytes();
|
||||
let raw = u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
|
||||
|
||||
if raw < REJECT_THRESHOLD {
|
||||
return format!("{:06}", raw % UPPER_BOUND);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Generate a cryptographically-adequate bearer token with 256-bit entropy.
|
||||
///
|
||||
/// Uses `rand::rng()` which is backed by the OS CSPRNG
|
||||
/// (/dev/urandom on Linux, BCryptGenRandom on Windows, SecRandomCopyBytes
|
||||
/// on macOS). The 32 random bytes (256 bits) are hex-encoded for a
|
||||
/// 64-character token, providing 256 bits of entropy.
|
||||
fn generate_token() -> String {
|
||||
use rand::RngExt as _;
|
||||
let mut bytes = [0u8; 32];
|
||||
rand::rng().fill(&mut bytes);
|
||||
format!("zc_{}", hex::encode(bytes))
|
||||
}
|
||||
|
||||
/// SHA-256 hash a bearer token for storage. Returns lowercase hex.
|
||||
fn hash_token(token: &str) -> String {
|
||||
format!("{:x}", Sha256::digest(token.as_bytes()))
|
||||
}
|
||||
|
||||
/// Check if a stored value looks like a SHA-256 hash (64 hex chars)
|
||||
/// rather than a plaintext token.
|
||||
fn is_token_hash(value: &str) -> bool {
|
||||
value.len() == 64 && value.chars().all(|c| c.is_ascii_hexdigit())
|
||||
}
|
||||
|
||||
/// Constant-time string comparison to prevent timing attacks.
|
||||
///
|
||||
/// Does not short-circuit on length mismatch — always iterates over the
|
||||
/// longer input to avoid leaking length information via timing.
|
||||
pub fn constant_time_eq(a: &str, b: &str) -> bool {
|
||||
let a = a.as_bytes();
|
||||
let b = b.as_bytes();
|
||||
|
||||
// Track length mismatch as a usize (non-zero = different lengths)
|
||||
let len_diff = a.len() ^ b.len();
|
||||
|
||||
// XOR each byte, padding the shorter input with zeros.
|
||||
// Iterates over max(a.len(), b.len()) to avoid timing differences.
|
||||
let max_len = a.len().max(b.len());
|
||||
let mut byte_diff = 0u8;
|
||||
for i in 0..max_len {
|
||||
let x = *a.get(i).unwrap_or(&0);
|
||||
let y = *b.get(i).unwrap_or(&0);
|
||||
byte_diff |= x ^ y;
|
||||
}
|
||||
(len_diff == 0) & (byte_diff == 0)
|
||||
}
|
||||
|
||||
/// Check if a host string represents a non-localhost bind address.
|
||||
pub fn is_public_bind(host: &str) -> bool {
|
||||
!matches!(
|
||||
|
||||
@@ -8,7 +8,7 @@ use std::time::Duration;
|
||||
use openhuman_core::core::event_bus::{DomainEvent, EventHandler};
|
||||
use openhuman_core::openhuman::agent_memory::memory_loader::MemoryCitation;
|
||||
use openhuman_core::openhuman::channels::bus::ChannelInboundSubscriber;
|
||||
use openhuman_core::openhuman::channels::providers::presentation::test_support as presentation_test_support;
|
||||
use openhuman_core::openhuman::channels::providers::web::presentation::test_support as presentation_test_support;
|
||||
use openhuman_core::openhuman::channels::providers::web::{
|
||||
subscribe_web_channel_events, test_support as web_test_support,
|
||||
};
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
//! End-to-end coverage for the ChannelHost capability boundary.
|
||||
//!
|
||||
//! Builds the *real* OpenHuman `ChannelHost` via `build_channel_host` and
|
||||
//! drives each capability through the `dyn ChannelHost` trait objects exactly
|
||||
//! as a ported channel provider would — proving the assembled host works
|
||||
//! across the portable boundary, not just the concrete adapters in isolation.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use openhuman_core::openhuman::channels::host::build_channel_host;
|
||||
use openhuman_core::openhuman::config::Config;
|
||||
use tinychannels::host::{ApprovalDecision, ConversationMessage, ReactionQuery};
|
||||
|
||||
/// A config pinned to a throwaway workspace with the local runtime disabled,
|
||||
/// so filesystem-backed capabilities are isolated and inference is short-circuited.
|
||||
fn test_config(workspace: &std::path::Path) -> Config {
|
||||
let mut config = Config::default();
|
||||
config.workspace_dir = workspace.to_path_buf();
|
||||
config.local_ai.runtime_enabled = false;
|
||||
config
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn host_advertises_the_wired_capability_set() {
|
||||
let host = build_channel_host(Arc::new(Config::default()));
|
||||
let caps = host.capabilities();
|
||||
assert!(caps.lifecycle);
|
||||
assert!(caps.stt);
|
||||
assert!(caps.tts);
|
||||
assert!(caps.reaction_gate);
|
||||
assert!(caps.approvals);
|
||||
assert!(caps.conversation_store);
|
||||
assert!(caps.event_sink);
|
||||
assert!(!caps.is_lean());
|
||||
// Not yet backed portably — a provider needing these degrades gracefully.
|
||||
assert!(!caps.turn_dispatch);
|
||||
assert!(!caps.run_ledger);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn conversation_store_roundtrips_through_the_boundary() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let host = build_channel_host(Arc::new(test_config(dir.path())));
|
||||
let store = host.conversations().expect("conversation store present");
|
||||
|
||||
let thread = "e2e-thread";
|
||||
store
|
||||
.append(
|
||||
thread,
|
||||
ConversationMessage {
|
||||
role: "user".into(),
|
||||
content: "ping".into(),
|
||||
timestamp: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("append user");
|
||||
store
|
||||
.append(
|
||||
thread,
|
||||
ConversationMessage {
|
||||
role: "assistant".into(),
|
||||
content: "pong".into(),
|
||||
timestamp: None,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("append assistant");
|
||||
|
||||
let history = store.history(thread, 10).await.expect("history");
|
||||
assert_eq!(history.len(), 2);
|
||||
assert_eq!(history[0].content, "ping");
|
||||
assert_eq!(history[1].content, "pong");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reaction_gate_short_circuits_when_runtime_disabled() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let host = build_channel_host(Arc::new(test_config(dir.path())));
|
||||
let gate = host.reactions().expect("reaction gate present");
|
||||
|
||||
let decision = gate
|
||||
.should_react(ReactionQuery {
|
||||
message: "hello there".into(),
|
||||
channel_type: "web".into(),
|
||||
})
|
||||
.await
|
||||
.expect("should_react ok");
|
||||
assert!(!decision.should_react);
|
||||
assert!(decision.emoji.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn approval_gate_parses_replies_through_the_boundary() {
|
||||
let host = build_channel_host(Arc::new(Config::default()));
|
||||
let gate = host.approvals().expect("approval gate present");
|
||||
assert_eq!(gate.parse_reply("yes"), Some(ApprovalDecision::Approve));
|
||||
assert_eq!(gate.parse_reply("no"), Some(ApprovalDecision::Deny));
|
||||
assert_eq!(gate.parse_reply("not-a-decision"), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn event_sink_publishes_web_channel_events() {
|
||||
let host = build_channel_host(Arc::new(Config::default()));
|
||||
let events = host.events().expect("event sink present");
|
||||
let ok = events
|
||||
.publish(
|
||||
"web",
|
||||
"chat_segment",
|
||||
serde_json::json!({
|
||||
"event": "chat_segment",
|
||||
"client_id": "c",
|
||||
"thread_id": "t",
|
||||
"request_id": "r",
|
||||
"message": "hi"
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
assert!(ok.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifecycle_registry_accepts_shutdown_hooks() {
|
||||
let host = build_channel_host(Arc::new(Config::default()));
|
||||
let lifecycle = host.lifecycle().expect("lifecycle present");
|
||||
// Registration must not panic; the hook runs only at real process shutdown.
|
||||
lifecycle.register_shutdown(
|
||||
"e2e-test-hook",
|
||||
Box::new(|| Box::pin(async move { /* no-op */ })),
|
||||
);
|
||||
}
|
||||
@@ -15,7 +15,9 @@ use openhuman_core::openhuman::channels::test_support::{
|
||||
use openhuman_core::openhuman::channels::LarkChannel;
|
||||
use reqwest::StatusCode as ReqwestStatusCode;
|
||||
use serde_json::json;
|
||||
use tokio_tungstenite::tungstenite::Message as WsMsg;
|
||||
// Lark's WS seam lives in tinychannels (tungstenite 0.29); use its re-export so
|
||||
// the message type matches the function signature across the version boundary.
|
||||
use tinychannels::tokio_tungstenite::tungstenite::Message as WsMsg;
|
||||
|
||||
async fn spawn_mock(app: Router) -> String {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! Integration tests for the TokenJuice module.
|
||||
//!
|
||||
//! Iterates vendored TinyJuice `*.fixture.json` files under
|
||||
//! `vendor/tinyjuice/src/tests/fixtures/` and asserts that
|
||||
//! `vendor/tinyjuice/tests/fixtures/` and asserts that
|
||||
//! `reduce_execution_with_rules` produces the expected output.
|
||||
|
||||
use openhuman_core::openhuman::tokenjuice::{
|
||||
@@ -17,7 +17,7 @@ const KNOWN_DRIFT_FIXTURES: &[&str] = &[
|
||||
|
||||
fn fixtures_dir() -> std::path::PathBuf {
|
||||
let manifest = std::env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not set");
|
||||
std::path::PathBuf::from(manifest).join("vendor/tinyjuice/src/tests/fixtures")
|
||||
std::path::PathBuf::from(manifest).join("vendor/tinyjuice/tests/fixtures")
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Vendored
+1
-1
Submodule vendor/tinychannels updated: 985c8f87b4...94fa9f5cb0
Reference in New Issue
Block a user