/
stream.go
52 lines (42 loc) · 1.34 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
package p2pnet
import (
"context"
"io"
"log"
host "github.com/libp2p/go-libp2p-host"
inet "github.com/libp2p/go-libp2p-net"
peer "github.com/libp2p/go-libp2p-peer"
protocol "github.com/libp2p/go-libp2p-protocol"
)
//StreamMgr handles connections of network
type StreamMgr struct {
host host.Host
}
// Stream represents a bidirectional channel between two agents
type Stream interface {
io.ReadWriter
// Close closes the stream for writing. Reading will still work (that
// is, the remote side can still write).
io.Closer
// Reset closes both ends of the stream. Use this to tell the remote
// side to hang up and go away.
Reset() error
}
// StreamHandler is the type of function used to listen for
// streams opened by the remote side.
type StreamHandler func(Stream)
// SetHandler sets the protocol handler on host.
func (mgr *StreamMgr) SetHandler(pid protocol.ID, handler StreamHandler) {
mgr.host.SetStreamHandler(pid, func(s inet.Stream) {
log.Printf("Stream opened for protocol %s", pid)
handler(s)
})
}
// RemoveHandler removes the protocol handler
func (mgr *StreamMgr) RemoveHandler(pid protocol.ID) {
mgr.host.RemoveStreamHandler(pid)
}
// NewStream opens a new stream to given peer
func (mgr *StreamMgr) NewStream(ctx context.Context, peer peer.ID, pid protocol.ID) (Stream, error) {
return mgr.host.NewStream(ctx, peer, pid)
}