fix(optimism): Remove loop in poll_next to prevent infinite cycles

Author: RomanHodulak

crates/optimism/flashblocks/src/ws/stream.rs

@@ -72,53 +72,51 @@ where
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = self.get_mut();
 
-        loop {
-            if this.state == State::Initial {
-                this.connect();
-            }
+        if this.state == State::Initial {
+            this.connect();
+        }
 
-            if this.state == State::Connect {
-                match ready!(this.connect.poll_unpin(cx)) {
-                    Ok((sink, stream)) => this.stream(sink, stream),
-                    Err(err) => {
-                        this.connect();
+        if this.state == State::Connect {
+            match ready!(this.connect.poll_unpin(cx)) {
+                Ok((sink, stream)) => this.stream(sink, stream),
+                Err(err) => {
+                    this.connect();
 
-                        return Poll::Ready(Some(Err(err)));
-                    }
+                    return Poll::Ready(Some(Err(err)));
                 }
             }
+        }
 
-            while let State::Stream(pong) = &mut this.state {
-                if pong.is_some() {
-                    let mut sink = Pin::new(this.sink.as_mut().unwrap());
-                    let _ = ready!(sink.as_mut().poll_ready(cx));
-                    if let Some(pong) = pong.take() {
-                        let _ = sink.as_mut().start_send(pong);
-                    }
-                    let _ = ready!(sink.as_mut().poll_flush(cx));
+        while let State::Stream(pong) = &mut this.state {
+            if pong.is_some() {
+                let mut sink = Pin::new(this.sink.as_mut().unwrap());
+                let _ = ready!(sink.as_mut().poll_ready(cx));
+                if let Some(pong) = pong.take() {
+                    let _ = sink.as_mut().start_send(pong);
                 }
+                let _ = ready!(sink.as_mut().poll_flush(cx));
+            }
 
-                let Some(msg) = ready!(this
-                    .stream
-                    .as_mut()
-                    .expect("Stream state should be unreachable without stream")
-                    .poll_next_unpin(cx))
-                else {
-                    this.connect();
+            let Some(msg) = ready!(this
+                .stream
+                .as_mut()
+                .expect("Stream state should be unreachable without stream")
+                .poll_next_unpin(cx))
+            else {
+                this.connect();
 
-                    break;
-                };
+                break;
+            };
 
-                match msg {
-                    Ok(Message::Binary(bytes)) => {
-                        return Poll::Ready(Some(FlashBlock::decode(bytes)))
-                    }
-                    Ok(Message::Ping(bytes)) => this.ping(bytes),
-                    Ok(msg) => debug!("Received unexpected message: {:?}", msg),
-                    Err(err) => return Poll::Ready(Some(Err(err.into()))),
-                }
+            match msg {
+                Ok(Message::Binary(bytes)) => return Poll::Ready(Some(FlashBlock::decode(bytes))),
+                Ok(Message::Ping(bytes)) => this.ping(bytes),
+                Ok(msg) => debug!("Received unexpected message: {:?}", msg),
+                Err(err) => return Poll::Ready(Some(Err(err.into()))),
             }
         }
+
+        Poll::Pending
     }
 }
 
@@ -223,6 +221,7 @@ mod tests {
     use crate::ExecutionPayloadBaseV1;
     use alloy_primitives::bytes::Bytes;
     use brotli::enc::BrotliEncoderParams;
+    use futures_util::poll;
     use std::{future, iter};
     use tokio_tungstenite::tungstenite::{Error, Utf8Bytes};
 
@@ -510,7 +509,7 @@ mod tests {
         let ws_url = "http://localhost".parse().unwrap();
         let mut stream = WsFlashBlockStream::with_connector(ws_url, connector);
 
-        let _ = stream.next().await;
+        let _ = poll!(stream.next());
 
         let FakeSink(actual_buffered_messages, actual_sent_messages) = stream.sink.unwrap();