-
Notifications
You must be signed in to change notification settings - Fork 9
/
stream.go
67 lines (56 loc) · 1.35 KB
/
stream.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
package libp2p
import (
"errors"
"io"
"time"
"github.com/gauss-project/aurorafs/pkg/p2p"
"github.com/libp2p/go-libp2p-core/network"
)
var (
closeDeadline = 30 * time.Second
errExpectedEof = errors.New("read: expected eof")
)
var _ p2p.Stream = (*stream)(nil)
type stream struct {
network.Stream
headers map[string][]byte
responseHeaders map[string][]byte
}
func NewStream(s network.Stream) p2p.Stream {
return &stream{Stream: s}
}
func newStream(s network.Stream) *stream {
return &stream{Stream: s}
}
func (s *stream) Headers() p2p.Headers {
return s.headers
}
func (s *stream) ResponseHeaders() p2p.Headers {
return s.responseHeaders
}
func (s *stream) FullClose() error {
// close the stream to make sure it is gc'd
defer s.Close()
if err := s.CloseWrite(); err != nil {
_ = s.Reset()
return err
}
// So we don't wait forever
_ = s.SetDeadline(time.Now().Add(closeDeadline))
// We *have* to observe the EOF. Otherwise, we leak the stream.
// Now, technically, we should do this *before*
// returning from SendMessage as the message
// hasn't really been sent yet until we see the
// EOF but we don't actually *know* what
// protocol the other side is speaking.
n, err := s.Read([]byte{0})
if n > 0 || err == nil {
_ = s.Reset()
return errExpectedEof
}
if err != io.EOF {
_ = s.Reset()
return err
}
return nil
}