From 8fd81bc29a2beb01721e149c0a90b2e3ed825886 Mon Sep 17 00:00:00 2001
From: sashabaranov <677093+sashabaranov@users.noreply.github.com>
Date: Sun, 12 Feb 2023 16:21:43 +0400
Subject: [PATCH] Break infinite loops in case [DONE] is missing (#67)

* Break infinite loops in case [DONE] is missing

* lint
---
 stream.go | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/stream.go b/stream.go
index 8c87e4c..04a32fc 100644
--- a/stream.go
+++ b/stream.go
@@ -11,12 +11,19 @@ import (
 	"net/http"
 )
 
+var (
+	emptyMessagesLimit            = 100
+	ErrTooManyEmptyStreamMessages = errors.New("stream has sent too many empty messages")
+)
+
 type CompletionStream struct {
 	reader   *bufio.Reader
 	response *http.Response
 }
 
 func (stream *CompletionStream) Recv() (response CompletionResponse, err error) {
+	emptyMessagesCount := 0
+
 waitForData:
 	line, err := stream.reader.ReadBytes('\n')
 	if err != nil {
@@ -28,6 +35,12 @@ waitForData:
 	var headerData = []byte("data: ")
 	line = bytes.TrimSpace(line)
 	if !bytes.HasPrefix(line, headerData) {
+		emptyMessagesCount++
+		if emptyMessagesCount > emptyMessagesLimit {
+			err = ErrTooManyEmptyStreamMessages
+			return
+		}
+
 		goto waitForData
 	}