Skip to content

Commit 8fd81bc

Browse files
authored
Break infinite loops in case [DONE] is missing (#67)
* Break infinite loops in case [DONE] is missing * lint
1 parent 5191ea6 commit 8fd81bc

File tree

1 file changed

+13
-0
lines changed

1 file changed

+13
-0
lines changed

stream.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,19 @@ import (
1111
"net/http"
1212
)
1313

14+
var (
15+
emptyMessagesLimit = 100
16+
ErrTooManyEmptyStreamMessages = errors.New("stream has sent too many empty messages")
17+
)
18+
1419
type CompletionStream struct {
1520
reader *bufio.Reader
1621
response *http.Response
1722
}
1823

1924
func (stream *CompletionStream) Recv() (response CompletionResponse, err error) {
25+
emptyMessagesCount := 0
26+
2027
waitForData:
2128
line, err := stream.reader.ReadBytes('\n')
2229
if err != nil {
@@ -28,6 +35,12 @@ waitForData:
2835
var headerData = []byte("data: ")
2936
line = bytes.TrimSpace(line)
3037
if !bytes.HasPrefix(line, headerData) {
38+
emptyMessagesCount++
39+
if emptyMessagesCount > emptyMessagesLimit {
40+
err = ErrTooManyEmptyStreamMessages
41+
return
42+
}
43+
3144
goto waitForData
3245
}
3346

0 commit comments

Comments
 (0)