0.7.17
This commit is contained in:
429
kb_lib/src/ws_hybrid_observation.rs
Normal file
429
kb_lib/src/ws_hybrid_observation.rs
Normal file
@@ -0,0 +1,429 @@
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user