2
0
mirror of https://github.com/soheilhy/cmux.git synced 2024-09-19 18:45:48 +08:00
cmux/buffer.go
Soheil Hassas Yeganeh d5924ef0b4 Fix a blocking issue in buffer reader
After sniffing and buffering data, if we try to read from
the socket again, bufio.Reader may block. This breaks HTTP handlers
in go1.5.2+ if one tries on browsers or with curl. Go's HTTP client,
however, is not broken. This issue is also there with TeeReader.

Return immediately with the data in the sniffed buffer.
2016-04-24 12:55:13 -04:00

50 lines
1.2 KiB
Go

package cmux
import (
"bytes"
"io"
)
// bufferedReader is an optimized implementation of io.Reader that behaves like
// ```
// io.MultiReader(bytes.NewReader(buffer.Bytes()), io.TeeReader(source, buffer))
// ```
// without allocating.
type bufferedReader struct {
source io.Reader
buffer bytes.Buffer
bufferRead int
bufferSize int
sniffing bool
lastErr error
}
func (s *bufferedReader) Read(p []byte) (int, error) {
if s.bufferSize > s.bufferRead {
// If we have already read something from the buffer before, we return the
// same data and the last error if any. We need to immediately return,
// otherwise we may block for ever, if we try to be smart and call
// source.Read() seeking a little bit of more data.
bn := copy(p, s.buffer.Bytes()[s.bufferRead:s.bufferSize])
s.bufferRead += bn
return bn, s.lastErr
}
// If there is nothing more to return in the sniffed buffer, read from the
// source.
sn, sErr := s.source.Read(p)
if sn > 0 && s.sniffing {
s.lastErr = sErr
if wn, wErr := s.buffer.Write(p[:sn]); wErr != nil {
return wn, wErr
}
}
return sn, sErr
}
func (s *bufferedReader) reset(snif bool) {
s.sniffing = snif
s.bufferRead = 0
s.bufferSize = s.buffer.Len()
}