-
Notifications
You must be signed in to change notification settings - Fork 1
/
test_server.go
159 lines (129 loc) · 4.13 KB
/
test_server.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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
package testserver
import (
"bufio"
"encoding/pem"
"fmt"
"net"
"crypto/rand"
"crypto/rsa"
"crypto/x509"
assert "github.com/stretchr/testify/require"
"golang.org/x/crypto/ssh"
)
// SSHServer represents a test SSH Server
type SSHServer struct {
listener net.Listener
}
// SSHHandler is the interface that is implemented to handle an SSH channel.
type SSHHandler interface {
// Handler is a function that handles i/o to/from an SSH channel
Handle(t assert.TestingT, ch ssh.Channel)
}
// HandlerFactory is a test function that will deliver an SSHHandler.
type HandlerFactory func(t assert.TestingT) SSHHandler
// NewSSHServer deflivers a new test SSH Server, with a Handler that simply echoes lines received.
// The server implements password authentication with the given credentials.
func NewSSHServer(t assert.TestingT, uname, password string) *SSHServer {
return NewSSHServerHandler(t, uname, password, func(t assert.TestingT) SSHHandler { return &echoer{} })
}
// NewSSHServerHandler deflivers a new test SSH Server, with a custom channel handler.
// The server implements password authentication with the given credentials.
func NewSSHServerHandler(t assert.TestingT, uname, password string, factory HandlerFactory) *SSHServer {
listener, err := net.Listen("tcp", "localhost:0")
assert.NoError(t, err, "Listen failed")
go acceptConnections(t, listener, newSSHServerConfig(t, uname, password), factory)
return &SSHServer{listener: listener}
}
// Port delivers the tcp port number on which the server is listening.
func (ts *SSHServer) Port() int {
return ts.listener.Addr().(*net.TCPAddr).Port
}
// Close closes any resources used by the server.
func (ts *SSHServer) Close() {
// nolint: gosec, errcheck
ts.listener.Close()
}
func acceptConnections(t assert.TestingT, listener net.Listener, config *ssh.ServerConfig, factory HandlerFactory) {
// nolint: gosec, errcheck
for {
nConn, err := listener.Accept()
if err != nil {
return
}
_, chch, reqch, err := ssh.NewServerConn(nConn, config)
if err != nil {
return
}
go ssh.DiscardRequests(reqch)
// Service the incoming Channel channel.
for newChannel := range chch {
dataChan, requests, err := newChannel.Accept()
assert.NoError(t, err, "Failed to accept new channel")
// Handle the "subsystem" request.
go func(in <-chan *ssh.Request) {
for req := range in {
assert.NoError(t, req.Reply(req.Type == "subsystem", nil), "Request reply failed")
}
}(requests)
go func() {
defer dataChan.Close()
factory(t).Handle(t, dataChan)
}()
}
}
}
func newSSHServerConfig(t assert.TestingT, uname, password string) *ssh.ServerConfig {
config := &ssh.ServerConfig{
PasswordCallback: func(c ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
if c.User() == uname && string(pass) == password {
return nil, nil
}
return nil, fmt.Errorf("password rejected for %q", c.User())
},
}
config.AddHostKey(generateHostKey(t))
return config
}
func generateHostKey(t assert.TestingT) (hostkey ssh.Signer) { // nolint: interfacer
reader := rand.Reader
bitSize := 2048
var err error
var key *rsa.PrivateKey
if key, err = rsa.GenerateKey(reader, bitSize); err == nil {
privateBytes := encodePrivateKeyToPEM(key)
if hostkey, err = ssh.ParsePrivateKey(privateBytes); err == nil {
return
}
}
t.Errorf("Failed to generate host key %v", err)
return
}
func encodePrivateKeyToPEM(privateKey *rsa.PrivateKey) []byte {
// Get ASN.1 DER format
privDER := x509.MarshalPKCS1PrivateKey(privateKey)
// pem.Block
privBlock := pem.Block{
Type: "RSA PRIVATE KEY",
Headers: nil,
Bytes: privDER,
}
// Private key in PEM format
privatePEM := pem.EncodeToMemory(&privBlock)
return privatePEM
}
type echoer struct{}
// Simple Handler implementation that echoes lines.
func (e *echoer) Handle(t assert.TestingT, ch ssh.Channel) {
chReader := bufio.NewReader(ch)
chWriter := bufio.NewWriter(ch)
for {
input, err := chReader.ReadString('\n')
if err != nil {
return
}
_, err = chWriter.WriteString(fmt.Sprintf("GOT:%s", input))
assert.NoError(t, err, "Write failed")
err = chWriter.Flush()
assert.NoError(t, err, "Flush failed")
}
}