blob: f8cf30a1e66af08e267f52488e6ac9c1a0119f7c [file] [log] [blame]
khenaidooffe076b2019-01-15 16:08:08 -05001// Copyright 2016 The CMux Authors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12// implied. See the License for the specific language governing
13// permissions and limitations under the License.
14
15package cmux
16
17import (
18 "bytes"
19 "io"
20)
21
22// bufferedReader is an optimized implementation of io.Reader that behaves like
23// ```
24// io.MultiReader(bytes.NewReader(buffer.Bytes()), io.TeeReader(source, buffer))
25// ```
26// without allocating.
27type bufferedReader struct {
28 source io.Reader
29 buffer bytes.Buffer
30 bufferRead int
31 bufferSize int
32 sniffing bool
33 lastErr error
34}
35
36func (s *bufferedReader) Read(p []byte) (int, error) {
37 if s.bufferSize > s.bufferRead {
38 // If we have already read something from the buffer before, we return the
39 // same data and the last error if any. We need to immediately return,
40 // otherwise we may block for ever, if we try to be smart and call
41 // source.Read() seeking a little bit of more data.
42 bn := copy(p, s.buffer.Bytes()[s.bufferRead:s.bufferSize])
43 s.bufferRead += bn
44 return bn, s.lastErr
45 } else if !s.sniffing && s.buffer.Cap() != 0 {
46 // We don't need the buffer anymore.
47 // Reset it to release the internal slice.
48 s.buffer = bytes.Buffer{}
49 }
50
51 // If there is nothing more to return in the sniffed buffer, read from the
52 // source.
53 sn, sErr := s.source.Read(p)
54 if sn > 0 && s.sniffing {
55 s.lastErr = sErr
56 if wn, wErr := s.buffer.Write(p[:sn]); wErr != nil {
57 return wn, wErr
58 }
59 }
60 return sn, sErr
61}
62
63func (s *bufferedReader) reset(snif bool) {
64 s.sniffing = snif
65 s.bufferRead = 0
66 s.bufferSize = s.buffer.Len()
67}