diff --git a/core/src/cluster_info.rs b/core/src/cluster_info.rs index 320a02e29d..09db308b38 100644 --- a/core/src/cluster_info.rs +++ b/core/src/cluster_info.rs @@ -313,7 +313,7 @@ impl Default for ClusterInfo { } } -#[derive(Debug, Default, Deserialize, Serialize, AbiExample)] +#[derive(Clone, Debug, Default, Deserialize, Serialize, AbiExample)] pub struct PruneData { /// Pubkey of the node that sent this prune data pub pubkey: Pubkey, @@ -413,6 +413,8 @@ enum Protocol { PullRequest(CrdsFilter, CrdsValue), PullResponse(Pubkey, Vec), PushMessage(Pubkey, Vec), + // TODO: Remove the redundant outer pubkey here, + // and use the inner PruneData.pubkey instead. PruneMessage(Pubkey, PruneData), PingMessage(Ping), PongMessage(Pong), @@ -496,7 +498,13 @@ impl Sanitize for Protocol { } Protocol::PullResponse(_, val) => val.sanitize(), Protocol::PushMessage(_, val) => val.sanitize(), - Protocol::PruneMessage(_, val) => val.sanitize(), + Protocol::PruneMessage(from, val) => { + if *from != val.pubkey { + Err(SanitizeError::InvalidValue) + } else { + val.sanitize() + } + } Protocol::PingMessage(ping) => ping.sanitize(), Protocol::PongMessage(pong) => pong.sanitize(), } @@ -4169,6 +4177,23 @@ mod tests { assert_eq!(msg.sanitize(), Err(SanitizeError::ValueOutOfBounds)); } + #[test] + fn test_protocol_prune_message_sanitize() { + let keypair = Keypair::new(); + let mut prune_data = PruneData { + pubkey: keypair.pubkey(), + prunes: vec![], + signature: Signature::default(), + destination: Pubkey::new_unique(), + wallclock: timestamp(), + }; + prune_data.sign(&keypair); + let prune_message = Protocol::PruneMessage(keypair.pubkey(), prune_data.clone()); + assert_eq!(prune_message.sanitize(), Ok(())); + let prune_message = Protocol::PruneMessage(Pubkey::new_unique(), prune_data); + assert_eq!(prune_message.sanitize(), Err(SanitizeError::InvalidValue)); + } + // computes the maximum size for pull request blooms fn max_bloom_size() -> usize { let filter_size = serialized_size(&CrdsFilter::default())