/
mux.go
82 lines (68 loc) · 2.33 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
package sshd
import (
"context"
"sync"
"golang.org/x/crypto/ssh"
)
// DefaultServeMux is the default ServeMux used by Serve.
var DefaultServeMux = NewServeMux()
// Handler 在建立 channel 时调用 ServeChannel.
type Handler interface {
// ServeChannel 对 ssh.NewChannel 处理, ctx 来源于 ConnContext,
// 如果返回的 error 不为 nil 将会在 Server 中输出日志.
ServeChannel(ctx context.Context, conn *ssh.ServerConn, newChannel ssh.NewChannel) error
}
// HandlerFunc 函数类型的 Handler
type HandlerFunc func(ctx context.Context, conn *ssh.ServerConn, newChannel ssh.NewChannel) error
// ServeChannel implements Handler.ServeChannel
func (h HandlerFunc) ServeChannel(ctx context.Context, conn *ssh.ServerConn, newChannel ssh.NewChannel) error {
return h(ctx, conn, newChannel)
}
// ServeMux is an SSH request multiplexer.
type ServeMux struct {
mut sync.RWMutex
handlers map[string]Handler
}
// NewServeMux allocates and returns a new ServeMux.
func NewServeMux() *ServeMux {
return &ServeMux{
handlers: make(map[string]Handler),
}
}
// Handle registers the handler for the given channel type.
// Panics If a handler already existed for channel type.
func (mux *ServeMux) Handle(channelType string, handler Handler) {
mux.mut.Lock()
defer mux.mut.Unlock()
if len(channelType) == 0 {
panic("mux: invalid channel type")
}
if handler == nil {
panic("mux: nil handler")
}
if mux.handlers == nil {
mux.handlers = make(map[string]Handler)
}
if _, existed := mux.handlers[channelType]; existed {
panic("mux: multiple registrations for " + channelType)
}
mux.handlers[channelType] = handler
}
// HandleFunc registers the handler function for the given channel type.
func (mux *ServeMux) HandleFunc(channelType string, handler func(context.Context, *ssh.ServerConn, ssh.NewChannel) error) {
if handler == nil {
panic("mux: nil handler")
}
mux.Handle(channelType, HandlerFunc(handler))
}
// ServeChannel implements Handler
func (mux *ServeMux) ServeChannel(ctx context.Context, conn *ssh.ServerConn, newChannel ssh.NewChannel) error {
channelType := newChannel.ChannelType()
mux.mut.RLock()
handler, ok := mux.handlers[channelType]
mux.mut.RUnlock()
if ok && handler != nil {
return handler.ServeChannel(ctx, conn, newChannel)
}
return newChannel.Reject(ssh.UnknownChannelType, "unsupported channel type")
}