Skip to content

Commit

Permalink
Break infinite loops in case [DONE] is missing (#67)
Browse files Browse the repository at this point in the history
* Break infinite loops in case [DONE] is missing

* lint
  • Loading branch information
sashabaranov authored Feb 12, 2023
1 parent 5191ea6 commit 8fd81bc
Showing 1 changed file with 13 additions and 0 deletions.
13 changes: 13 additions & 0 deletions stream.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
}

Expand Down

0 comments on commit 8fd81bc

Please sign in to comment.