Add additional IgnoreHandshakeMessageAsync restrictions

Metalnem · May 17, 2018 · e97a3ed · e97a3ed
1 parent e3a3689
commit e97a3ed
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/NoiseSocket/NoiseSocket.cs b/NoiseSocket/NoiseSocket.cs
@@ -589,6 +589,17 @@ public async Task IgnoreHandshakeMessageAsync(CancellationToken cancellationToke
 				throw new InvalidOperationException(error);
 			}
 
+			if (client && (savedMessages == null || savedMessages.Count != 3))
+			{
+				string error = $"Client can ignore only the initial handshake message from the server in the Retry case.";
+				throw new InvalidOperationException(error);
+			}
+
+			if (!client && (savedMessages == null || savedMessages.Count != 1))
+			{
+				throw new InvalidOperationException($"Server can ignore only the initial handshake message from the client.");
+			}
+
 			var noiseMessage = await ReadPacketAsync(stream, cancellationToken).ConfigureAwait(false);
 			ProcessMessage(HandshakeOperation.ReadHandshakeMessage, noiseMessage, false);
 		}