/
mux.go
142 lines (121 loc) · 3.54 KB
/
mux.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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
package protocol
import (
"fmt"
"io"
"sync"
context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
inet "github.com/ipfs/go-ipfs/p2p/net"
eventlog "github.com/ipfs/go-ipfs/thirdparty/eventlog"
lgbl "github.com/ipfs/go-ipfs/util/eventlog/loggables"
)
var log = eventlog.Logger("net/mux")
type streamHandlerMap map[ID]inet.StreamHandler
// Mux provides simple stream multixplexing.
// It helps you precisely when:
// * You have many streams
// * You have function handlers
//
// It contains the handlers for each protocol accepted.
// It dispatches handlers for streams opened by remote peers.
type Mux struct {
lock sync.RWMutex
handlers streamHandlerMap
defaultHandler inet.StreamHandler
}
func NewMux() *Mux {
return &Mux{
handlers: streamHandlerMap{},
}
}
// Protocols returns the list of protocols this muxer has handlers for
func (m *Mux) Protocols() []ID {
m.lock.RLock()
l := make([]ID, 0, len(m.handlers))
for p := range m.handlers {
l = append(l, p)
}
m.lock.RUnlock()
return l
}
// ReadHeader reads the stream and returns the next Handler function
// according to the muxer encoding.
func (m *Mux) ReadHeader(s io.Reader) (ID, inet.StreamHandler, error) {
p, err := ReadHeader(s)
if err != nil {
return "", nil, err
}
m.lock.RLock()
defer m.lock.RUnlock()
h, found := m.handlers[p]
switch {
case !found && m.defaultHandler != nil:
return p, m.defaultHandler, nil
case !found && m.defaultHandler == nil:
return p, nil, fmt.Errorf("%s no handler with name: %s (%d)", m, p, len(p))
default:
return p, h, nil
}
}
// String returns the muxer's printing representation
func (m *Mux) String() string {
m.lock.RLock()
defer m.lock.RUnlock()
return fmt.Sprintf("<Muxer %p %d>", m, len(m.handlers))
}
func (m *Mux) SetDefaultHandler(h inet.StreamHandler) {
m.lock.Lock()
m.defaultHandler = h
m.lock.Unlock()
}
// SetHandler sets the protocol handler on the Network's Muxer.
// This operation is threadsafe.
func (m *Mux) SetHandler(p ID, h inet.StreamHandler) {
log.Debugf("%s setting handler for protocol: %s (%d)", m, p, len(p))
m.lock.Lock()
m.handlers[p] = h
m.lock.Unlock()
}
// RemoveHandler removes the protocol handler on the Network's Muxer.
// This operation is threadsafe.
func (m *Mux) RemoveHandler(p ID) {
log.Debugf("%s removing handler for protocol: %s (%d)", m, p, len(p))
m.lock.Lock()
delete(m.handlers, p)
m.lock.Unlock()
}
// Handle reads the next name off the Stream, and calls a handler function
// This is done in its own goroutine, to avoid blocking the caller.
func (m *Mux) Handle(s inet.Stream) {
go m.HandleSync(s)
}
// HandleSync reads the next name off the Stream, and calls a handler function
// This is done synchronously. The handler function will return before
// HandleSync returns.
func (m *Mux) HandleSync(s inet.Stream) {
ctx := context.Background()
name, handler, err := m.ReadHeader(s)
if err != nil {
err = fmt.Errorf("protocol mux error: %s", err)
log.Event(ctx, "muxError", lgbl.Error(err))
s.Close()
return
}
log.Debugf("muxer handle protocol %s: %s", s.Conn().RemotePeer(), name)
handler(s)
}
// ReadLengthPrefix reads the name from Reader with a length-byte-prefix.
func ReadLengthPrefix(r io.Reader) (string, error) {
// c-string identifier
// the first byte is our length
l := make([]byte, 1)
if _, err := io.ReadFull(r, l); err != nil {
return "", err
}
length := int(l[0])
// the next are our identifier
name := make([]byte, length)
if _, err := io.ReadFull(r, name); err != nil {
return "", err
}
return string(name), nil
}