Break infinite loops in case [DONE] is missing (#67)
* Break infinite loops in case [DONE] is missing * lint
This commit is contained in:
13
stream.go
13
stream.go
@@ -11,12 +11,19 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
emptyMessagesLimit = 100
|
||||||
|
ErrTooManyEmptyStreamMessages = errors.New("stream has sent too many empty messages")
|
||||||
|
)
|
||||||
|
|
||||||
type CompletionStream struct {
|
type CompletionStream struct {
|
||||||
reader *bufio.Reader
|
reader *bufio.Reader
|
||||||
response *http.Response
|
response *http.Response
|
||||||
}
|
}
|
||||||
|
|
||||||
func (stream *CompletionStream) Recv() (response CompletionResponse, err error) {
|
func (stream *CompletionStream) Recv() (response CompletionResponse, err error) {
|
||||||
|
emptyMessagesCount := 0
|
||||||
|
|
||||||
waitForData:
|
waitForData:
|
||||||
line, err := stream.reader.ReadBytes('\n')
|
line, err := stream.reader.ReadBytes('\n')
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -28,6 +35,12 @@ waitForData:
|
|||||||
var headerData = []byte("data: ")
|
var headerData = []byte("data: ")
|
||||||
line = bytes.TrimSpace(line)
|
line = bytes.TrimSpace(line)
|
||||||
if !bytes.HasPrefix(line, headerData) {
|
if !bytes.HasPrefix(line, headerData) {
|
||||||
|
emptyMessagesCount++
|
||||||
|
if emptyMessagesCount > emptyMessagesLimit {
|
||||||
|
err = ErrTooManyEmptyStreamMessages
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
goto waitForData
|
goto waitForData
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user