From 8449d2c91a7e14a2d5bef60d1152cd43c3c5ebd7 Mon Sep 17 00:00:00 2001
From: Jan Bujak <jan@parity.io>
Date: Mon, 26 Jul 2021 07:45:56 +0000
Subject: [PATCH] Do not leak active head data in statement distribution

---
 .../network/statement-distribution/src/lib.rs | 48 +++++++++++++------
 1 file changed, 34 insertions(+), 14 deletions(-)

diff --git a/node/network/statement-distribution/src/lib.rs b/node/network/statement-distribution/src/lib.rs
index 8d0cecdd63b4..311d2eacf1d9 100644
--- a/node/network/statement-distribution/src/lib.rs
+++ b/node/network/statement-distribution/src/lib.rs
@@ -1856,24 +1856,44 @@ impl StatementDistribution {
 						"New active leaf",
 					);
 
-					let session_index = runtime.get_session_index(ctx.sender(), relay_parent).await?;
-					let info = runtime.get_session_info_by_index(ctx.sender(), relay_parent, session_index).await?;
-					let session_info = &info.session_info;
-
-					active_heads.entry(relay_parent)
-						.or_insert(ActiveHeadData::new(session_info.validators.clone(), session_index, span));
-
-					active_heads.retain(|h, _| {
-						let live = !deactivated.contains(h);
-						if !live {
+					let session_index = match runtime.get_session_index(ctx.sender(), relay_parent).await {
+						Ok(session_index) => session_index,
+						Err(error) => {
+							tracing::trace!(
+								target: LOG_TARGET,
+								hash = ?relay_parent,
+								"Couldn't activate leaf: no session index found: {}",
+								error
+							);
+							continue;
+						}
+					};
+					let info = match runtime.get_session_info_by_index(ctx.sender(), relay_parent, session_index).await {
+						Ok(info) => info,
+						Err(error) => {
 							tracing::trace!(
 								target: LOG_TARGET,
-								hash = ?h,
-								"Deactivating leaf",
+								hash = ?relay_parent,
+								"Couldn't activate leaf: no session info found: {}",
+								error
 							);
+							continue;
 						}
-						live
-					});
+					};
+
+					let session_info = &info.session_info;
+					active_heads.entry(relay_parent)
+						.or_insert(ActiveHeadData::new(session_info.validators.clone(), session_index, span));
+				}
+
+				for deactivated in deactivated {
+					if active_heads.remove(&deactivated).is_some() {
+						tracing::trace!(
+							target: LOG_TARGET,
+							hash = ?deactivated,
+							"Deactivating leaf",
+						);
+					}
 				}
 			}
 			FromOverseer::Signal(OverseerSignal::BlockFinalized(..)) => {