From 5af7e51e39cb0defd23449207bd55082acba0837 Mon Sep 17 00:00:00 2001
From: kkast <kkastsevich@gmail.com>
Date: Wed, 13 Mar 2024 01:30:20 +0300
Subject: [PATCH] oak-audit: issue 1

---
 light-clients/ics10-grandpa/src/client_def.rs   | 7 +++++++
 light-clients/ics10-grandpa/src/client_state.rs | 8 ++++++++
 2 files changed, 15 insertions(+)

diff --git a/light-clients/ics10-grandpa/src/client_def.rs b/light-clients/ics10-grandpa/src/client_def.rs
index 9fdf0289a..a9d594162 100644
--- a/light-clients/ics10-grandpa/src/client_def.rs
+++ b/light-clients/ics10-grandpa/src/client_def.rs
@@ -92,6 +92,13 @@ where
 					))
 					.into())
 				}
+				if header.finality_proof.unknown_headers.len() >= client_state.session_length() {
+					return Err(Error::Custom(format!(
+						"Number of unknown headers submitted exceeds max number of blocks in a session: max {}, got {}",
+						client_state.session_length(), header.finality_proof.unknown_headers.len()
+					))
+					.into())
+				}
 				let headers_with_finality_proof = ParachainHeadersWithFinalityProof {
 					finality_proof: header.finality_proof,
 					parachain_header: header.parachain_header,
diff --git a/light-clients/ics10-grandpa/src/client_state.rs b/light-clients/ics10-grandpa/src/client_state.rs
index 941f7eb89..857126b67 100644
--- a/light-clients/ics10-grandpa/src/client_state.rs
+++ b/light-clients/ics10-grandpa/src/client_state.rs
@@ -146,6 +146,14 @@ impl<H> ClientState<H> {
 		elapsed > self.relay_chain.trusting_period()
 	}
 
+	pub fn session_length(&self) -> usize {
+		match self.relay_chain {
+			RelayChain::Polkadot => 2400,
+			RelayChain::Kusama => 600,
+			RelayChain::Rococo => 600,
+		}
+	}
+
 	pub fn with_frozen_height(self, h: Height) -> Result<Self, Error> {
 		if h == Height::zero() {
 			return Err(Error::Custom(