430 lines
15 KiB
Rust
430 lines
15 KiB
Rust
// file: kb_lib/src/ws_hybrid_observation.rs
|
|
|
|
//! Hybrid WebSocket technical observation service.
|
|
|
|
use std::hash::Hasher;
|
|
|
|
/// One hybrid WebSocket observation result.
|
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
|
pub struct KbWsHybridObservationResult {
|
|
/// Stable observation name.
|
|
pub observation_name: std::string::String,
|
|
/// Stable deduplication key.
|
|
pub dedupe_key: std::string::String,
|
|
/// Optional watched address.
|
|
pub watched_address: std::option::Option<std::string::String>,
|
|
/// Optional observed slot.
|
|
pub slot: std::option::Option<u64>,
|
|
/// Whether this observation was newly recorded during this service lifetime.
|
|
pub created_observation: bool,
|
|
}
|
|
|
|
/// Hybrid WebSocket technical observation service.
|
|
#[derive(Debug, Clone)]
|
|
pub struct KbWsHybridObservationService {
|
|
persistence: crate::KbDetectionPersistenceService,
|
|
seen_dedupe_keys:
|
|
std::sync::Arc<tokio::sync::Mutex<std::collections::HashSet<std::string::String>>>,
|
|
}
|
|
|
|
impl KbWsHybridObservationService {
|
|
/// Creates a new hybrid WebSocket technical observation service.
|
|
pub fn new(database: std::sync::Arc<crate::KbDatabase>) -> Self {
|
|
let persistence = crate::KbDetectionPersistenceService::new(database.clone());
|
|
Self {
|
|
persistence,
|
|
seen_dedupe_keys: std::sync::Arc::new(tokio::sync::Mutex::new(
|
|
std::collections::HashSet::<std::string::String>::new(),
|
|
)),
|
|
}
|
|
}
|
|
|
|
/// Records one `logsNotification` payload.
|
|
pub async fn record_logs_notification(
|
|
&self,
|
|
endpoint_name: std::option::Option<std::string::String>,
|
|
payload: &serde_json::Value,
|
|
) -> Result<crate::KbWsHybridObservationResult, crate::KbError> {
|
|
let signature = kb_extract_string_by_candidate_keys(payload, &["signature"]);
|
|
self.record_observation_inner(
|
|
"ws.hybrid.logs_notification".to_string(),
|
|
"signal.ws.hybrid.logs_notification".to_string(),
|
|
endpoint_name,
|
|
signature,
|
|
payload,
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Records one `programNotification` payload for one watched program id.
|
|
pub async fn record_program_notification(
|
|
&self,
|
|
endpoint_name: std::option::Option<std::string::String>,
|
|
watched_program_id: std::string::String,
|
|
payload: &serde_json::Value,
|
|
) -> Result<crate::KbWsHybridObservationResult, crate::KbError> {
|
|
let pubkey = kb_extract_string_by_candidate_keys(payload, &["pubkey"]);
|
|
let watched_address = match pubkey {
|
|
Some(pubkey) => Some(pubkey),
|
|
None => Some(watched_program_id),
|
|
};
|
|
self.record_observation_inner(
|
|
"ws.hybrid.program_notification".to_string(),
|
|
"signal.ws.hybrid.program_notification".to_string(),
|
|
endpoint_name,
|
|
watched_address,
|
|
payload,
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Records one `accountNotification` payload for one watched account address.
|
|
pub async fn record_account_notification(
|
|
&self,
|
|
endpoint_name: std::option::Option<std::string::String>,
|
|
watched_account: std::string::String,
|
|
payload: &serde_json::Value,
|
|
) -> Result<crate::KbWsHybridObservationResult, crate::KbError> {
|
|
self.record_observation_inner(
|
|
"ws.hybrid.account_notification".to_string(),
|
|
"signal.ws.hybrid.account_notification".to_string(),
|
|
endpoint_name,
|
|
Some(watched_account),
|
|
payload,
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn record_observation_inner(
|
|
&self,
|
|
observation_name: std::string::String,
|
|
signal_name: std::string::String,
|
|
endpoint_name: std::option::Option<std::string::String>,
|
|
watched_address: std::option::Option<std::string::String>,
|
|
payload: &serde_json::Value,
|
|
) -> Result<crate::KbWsHybridObservationResult, crate::KbError> {
|
|
let slot = kb_extract_slot(payload);
|
|
let payload_hash = kb_hash_payload(payload);
|
|
let dedupe_key = kb_build_ws_observation_dedupe_key(
|
|
observation_name.as_str(),
|
|
endpoint_name.as_deref(),
|
|
watched_address.as_deref(),
|
|
slot,
|
|
payload_hash.as_str(),
|
|
);
|
|
let mut seen_guard = self.seen_dedupe_keys.lock().await;
|
|
let already_seen = seen_guard.contains(&dedupe_key);
|
|
if !already_seen {
|
|
seen_guard.insert(dedupe_key.clone());
|
|
}
|
|
drop(seen_guard);
|
|
if already_seen {
|
|
return Ok(crate::KbWsHybridObservationResult {
|
|
observation_name,
|
|
dedupe_key,
|
|
watched_address,
|
|
slot,
|
|
created_observation: false,
|
|
});
|
|
}
|
|
let observation_payload = payload.clone();
|
|
let observation_result = self
|
|
.persistence
|
|
.record_observation(&crate::KbDetectionObservationInput::new(
|
|
observation_name.clone(),
|
|
crate::KbObservationSourceKind::WsRpc,
|
|
endpoint_name.clone(),
|
|
dedupe_key.clone(),
|
|
slot,
|
|
observation_payload.clone(),
|
|
))
|
|
.await;
|
|
if let Err(error) = observation_result {
|
|
return Err(error);
|
|
}
|
|
let signal_result = self
|
|
.persistence
|
|
.record_signal(&crate::KbDetectionSignalInput::new(
|
|
signal_name,
|
|
crate::KbAnalysisSignalSeverity::Low,
|
|
dedupe_key.clone(),
|
|
None,
|
|
None,
|
|
observation_payload,
|
|
))
|
|
.await;
|
|
if let Err(error) = signal_result {
|
|
return Err(error);
|
|
}
|
|
Ok(crate::KbWsHybridObservationResult {
|
|
observation_name,
|
|
dedupe_key,
|
|
watched_address,
|
|
slot,
|
|
created_observation: true,
|
|
})
|
|
}
|
|
}
|
|
|
|
fn kb_extract_slot(value: &serde_json::Value) -> std::option::Option<u64> {
|
|
kb_extract_u64_by_candidate_keys(value, &["slot"])
|
|
}
|
|
|
|
fn kb_extract_string_by_candidate_keys(
|
|
value: &serde_json::Value,
|
|
candidate_keys: &[&str],
|
|
) -> std::option::Option<std::string::String> {
|
|
if let Some(object) = value.as_object() {
|
|
for candidate_key in candidate_keys {
|
|
let direct_option = object.get(*candidate_key);
|
|
if let Some(direct) = direct_option {
|
|
let text_option = direct.as_str();
|
|
if let Some(text) = text_option {
|
|
return Some(text.to_string());
|
|
}
|
|
}
|
|
}
|
|
for nested_value in object.values() {
|
|
let nested_result = kb_extract_string_by_candidate_keys(nested_value, candidate_keys);
|
|
if nested_result.is_some() {
|
|
return nested_result;
|
|
}
|
|
}
|
|
return None;
|
|
}
|
|
if let Some(array) = value.as_array() {
|
|
for nested_value in array {
|
|
let nested_result = kb_extract_string_by_candidate_keys(nested_value, candidate_keys);
|
|
if nested_result.is_some() {
|
|
return nested_result;
|
|
}
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
fn kb_extract_u64_by_candidate_keys(
|
|
value: &serde_json::Value,
|
|
candidate_keys: &[&str],
|
|
) -> std::option::Option<u64> {
|
|
if let Some(object) = value.as_object() {
|
|
for candidate_key in candidate_keys {
|
|
let direct_option = object.get(*candidate_key);
|
|
if let Some(direct) = direct_option {
|
|
let number_option = direct.as_u64();
|
|
if let Some(number) = number_option {
|
|
return Some(number);
|
|
}
|
|
}
|
|
}
|
|
for nested_value in object.values() {
|
|
let nested_result = kb_extract_u64_by_candidate_keys(nested_value, candidate_keys);
|
|
if nested_result.is_some() {
|
|
return nested_result;
|
|
}
|
|
}
|
|
return None;
|
|
}
|
|
if let Some(array) = value.as_array() {
|
|
for nested_value in array {
|
|
let nested_result = kb_extract_u64_by_candidate_keys(nested_value, candidate_keys);
|
|
if nested_result.is_some() {
|
|
return nested_result;
|
|
}
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
fn kb_hash_payload(payload: &serde_json::Value) -> std::string::String {
|
|
let payload_text_result = serde_json::to_string(payload);
|
|
let payload_text = match payload_text_result {
|
|
Ok(payload_text) => payload_text,
|
|
Err(_) => return "serde_error".to_string(),
|
|
};
|
|
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
|
std::hash::Hash::hash(&payload_text, &mut hasher);
|
|
hasher.finish().to_string()
|
|
}
|
|
|
|
fn kb_build_ws_observation_dedupe_key(
|
|
source_method: &str,
|
|
endpoint_name: std::option::Option<&str>,
|
|
address: std::option::Option<&str>,
|
|
slot: std::option::Option<u64>,
|
|
payload_hash: &str,
|
|
) -> std::string::String {
|
|
format!(
|
|
"{}:{}:{}:{}:{}",
|
|
source_method,
|
|
endpoint_name.unwrap_or_default(),
|
|
address.unwrap_or_default(),
|
|
slot.unwrap_or_default(),
|
|
payload_hash
|
|
)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
async fn make_database() -> std::sync::Arc<crate::KbDatabase> {
|
|
let tempdir_result = tempfile::tempdir();
|
|
let tempdir = match tempdir_result {
|
|
Ok(tempdir) => tempdir,
|
|
Err(error) => panic!("tempdir must succeed: {}", error),
|
|
};
|
|
let database_path = tempdir.path().join("ws_hybrid_observation.sqlite3");
|
|
let config = crate::KbDatabaseConfig {
|
|
enabled: true,
|
|
backend: crate::KbDatabaseBackend::Sqlite,
|
|
sqlite: crate::KbSqliteDatabaseConfig {
|
|
path: database_path.to_string_lossy().to_string(),
|
|
create_if_missing: true,
|
|
busy_timeout_ms: 5000,
|
|
max_connections: 1,
|
|
auto_initialize_schema: true,
|
|
use_wal: true,
|
|
},
|
|
};
|
|
let database_result = crate::KbDatabase::connect_and_initialize(&config).await;
|
|
let database = match database_result {
|
|
Ok(database) => database,
|
|
Err(error) => panic!("database init must succeed: {}", error),
|
|
};
|
|
std::sync::Arc::new(database)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn record_program_notification_is_idempotent_for_same_payload() {
|
|
let database = make_database().await;
|
|
let service = crate::KbWsHybridObservationService::new(database);
|
|
let payload = serde_json::json!({
|
|
"method": "programNotification",
|
|
"params": {
|
|
"result": {
|
|
"context": {
|
|
"slot": 111
|
|
},
|
|
"value": {
|
|
"pubkey": "ProgramOwnedAccount111",
|
|
"account": {
|
|
"lamports": 1,
|
|
"owner": "Program111"
|
|
}
|
|
}
|
|
}
|
|
}
|
|
});
|
|
let first_result = service
|
|
.record_program_notification(
|
|
Some("helius_primary_ws_programs".to_string()),
|
|
"Program111".to_string(),
|
|
&payload,
|
|
)
|
|
.await;
|
|
let first = match first_result {
|
|
Ok(first) => first,
|
|
Err(error) => panic!("first program notification must succeed: {}", error),
|
|
};
|
|
assert!(first.created_observation);
|
|
let second_result = service
|
|
.record_program_notification(
|
|
Some("helius_primary_ws_programs".to_string()),
|
|
"Program111".to_string(),
|
|
&payload,
|
|
)
|
|
.await;
|
|
let second = match second_result {
|
|
Ok(second) => second,
|
|
Err(error) => panic!("second program notification must succeed: {}", error),
|
|
};
|
|
assert!(!second.created_observation);
|
|
assert_eq!(first.dedupe_key, second.dedupe_key);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn record_account_notification_is_idempotent_for_same_payload() {
|
|
let database = make_database().await;
|
|
let service = crate::KbWsHybridObservationService::new(database);
|
|
let payload = serde_json::json!({
|
|
"method": "accountNotification",
|
|
"params": {
|
|
"result": {
|
|
"context": {
|
|
"slot": 222
|
|
},
|
|
"value": {
|
|
"lamports": 10,
|
|
"owner": "Program111"
|
|
}
|
|
}
|
|
}
|
|
});
|
|
let first_result = service
|
|
.record_account_notification(
|
|
Some("mainnet_public_ws_accounts".to_string()),
|
|
"PoolAccount111".to_string(),
|
|
&payload,
|
|
)
|
|
.await;
|
|
let first = match first_result {
|
|
Ok(first) => first,
|
|
Err(error) => panic!("first account notification must succeed: {}", error),
|
|
};
|
|
assert!(first.created_observation);
|
|
let second_result = service
|
|
.record_account_notification(
|
|
Some("mainnet_public_ws_accounts".to_string()),
|
|
"PoolAccount111".to_string(),
|
|
&payload,
|
|
)
|
|
.await;
|
|
let second = match second_result {
|
|
Ok(second) => second,
|
|
Err(error) => panic!("second account notification must succeed: {}", error),
|
|
};
|
|
assert!(!second.created_observation);
|
|
assert_eq!(first.dedupe_key, second.dedupe_key);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn record_logs_notification_uses_signature_for_dedupe() {
|
|
let database = make_database().await;
|
|
let service = crate::KbWsHybridObservationService::new(database);
|
|
let payload = serde_json::json!({
|
|
"method": "logsNotification",
|
|
"params": {
|
|
"result": {
|
|
"context": {
|
|
"slot": 333
|
|
},
|
|
"value": {
|
|
"signature": "LogsSignature111",
|
|
"err": null,
|
|
"logs": [
|
|
"Program log: Instruction: InitializePool"
|
|
]
|
|
}
|
|
}
|
|
}
|
|
});
|
|
let first_result = service
|
|
.record_logs_notification(Some("mainnet_public_ws_logs".to_string()), &payload)
|
|
.await;
|
|
let first = match first_result {
|
|
Ok(first) => first,
|
|
Err(error) => panic!("first logs notification must succeed: {}", error),
|
|
};
|
|
assert!(first.created_observation);
|
|
assert_eq!(first.watched_address, Some("LogsSignature111".to_string()));
|
|
let second_result = service
|
|
.record_logs_notification(Some("mainnet_public_ws_logs".to_string()), &payload)
|
|
.await;
|
|
let second = match second_result {
|
|
Ok(second) => second,
|
|
Err(error) => panic!("second logs notification must succeed: {}", error),
|
|
};
|
|
assert!(!second.created_observation);
|
|
assert_eq!(first.dedupe_key, second.dedupe_key);
|
|
}
|
|
}
|