mirror of
https://github.com/gravitational/teleport.git
synced 2026-11-03 16:14:19 +08:00
Zac Bergquist
4e0968c0a5
newTestListener would call (*testing.T).Log when the Accept call fails. This resulted in racy behavior due to the log sometimes happening after the test completes (it's normal for accept to return an error after the listener is closed, which we do on test cleanup).
361 lines
9.8 KiB
Go
361 lines
9.8 KiB
Go
/*
|
|
* Teleport
|
|
* Copyright (C) 2023 Gravitational, Inc.
|
|
*
|
|
* This program is free software: you can redistribute it and/or modify
|
|
* it under the terms of the GNU Affero General Public License as published by
|
|
* the Free Software Foundation, either version 3 of the License, or
|
|
* (at your option) any later version.
|
|
*
|
|
* This program is distributed in the hope that it will be useful,
|
|
* but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
* GNU Affero General Public License for more details.
|
|
*
|
|
* You should have received a copy of the GNU Affero General Public License
|
|
* along with this program. If not, see <http://www.gnu.org/licenses/>.
|
|
*/
|
|
|
|
package client
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"io"
|
|
"net"
|
|
"os"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
"github.com/gravitational/trace"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"golang.org/x/crypto/ssh"
|
|
|
|
tracessh "github.com/gravitational/teleport/api/observability/tracing/ssh"
|
|
"github.com/gravitational/teleport/lib/observability/tracing"
|
|
"github.com/gravitational/teleport/lib/sshutils"
|
|
)
|
|
|
|
func TestHelperFunctions(t *testing.T) {
|
|
assert.Equal(t, "one", nodeName(targetNode{addr: "one"}))
|
|
assert.Equal(t, "one", nodeName(targetNode{addr: "one:22"}))
|
|
assert.Equal(t, "example.com", nodeName(targetNode{addr: "one", hostname: "example.com"}))
|
|
}
|
|
|
|
func TestNewSession(t *testing.T) {
|
|
nc := &NodeClient{
|
|
Namespace: "blue",
|
|
Tracer: tracing.NoopProvider().Tracer("test"),
|
|
}
|
|
|
|
ctx := context.Background()
|
|
// defaults:
|
|
ses, err := newSession(ctx, nc, nil, nil, nil, nil, nil, true)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, ses)
|
|
require.Equal(t, nc, ses.NodeClient())
|
|
require.Equal(t, nc.Namespace, ses.namespace)
|
|
require.NotNil(t, ses.env)
|
|
require.Equal(t, os.Stderr, ses.terminal.Stderr())
|
|
require.Equal(t, os.Stdout, ses.terminal.Stdout())
|
|
require.Equal(t, os.Stdin, ses.terminal.Stdin())
|
|
|
|
// pass environ map
|
|
env := map[string]string{
|
|
sshutils.SessionEnvVar: "session-id",
|
|
}
|
|
ses, err = newSession(ctx, nc, nil, env, nil, nil, nil, true)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, ses)
|
|
require.Empty(t, cmp.Diff(ses.env, env))
|
|
// the session ID must be taken from tne environ map, if passed:
|
|
require.Equal(t, "session-id", string(ses.id))
|
|
}
|
|
|
|
// TestProxyConnection verifies that client or server-side disconnect
|
|
// propagates all the way to the opposite side.
|
|
func TestProxyConnection(t *testing.T) {
|
|
// remoteSrv mocks a remote listener, accepting port-forwarded connections
|
|
// over SSH.
|
|
remoteConCh := make(chan net.Conn)
|
|
remoteErrCh := make(chan error, 3)
|
|
remoteSrv := newTestListener(t, func(con net.Conn) {
|
|
defer con.Close()
|
|
|
|
remoteConCh <- con
|
|
|
|
// Echo any data back to the sender.
|
|
_, err := io.Copy(con, con)
|
|
if err != nil && strings.Contains(err.Error(), "use of closed network connection") {
|
|
err = nil
|
|
}
|
|
remoteErrCh <- err
|
|
})
|
|
defer remoteSrv.Close()
|
|
|
|
// localSrv mocks a local tsh listener, accepting local connections for
|
|
// port-forwarding to remote SSH node.
|
|
proxyErrCh := make(chan error, 3)
|
|
localSrv := newTestListener(t, func(con net.Conn) {
|
|
defer con.Close()
|
|
|
|
proxyErrCh <- proxyConnection(context.Background(), con, remoteSrv.Addr().String(), new(net.Dialer))
|
|
})
|
|
defer localSrv.Close()
|
|
|
|
// Dial localSrv. This should trigger proxyConnection and a dial to
|
|
// remoteSrv.
|
|
localCon, err := net.Dial("tcp", localSrv.Addr().String())
|
|
require.NoError(t, err)
|
|
clientErrCh := make(chan error, 3)
|
|
go func(con net.Conn) {
|
|
_, err := io.Copy(io.Discard, con)
|
|
if err != nil && strings.Contains(err.Error(), "use of closed network connection") {
|
|
err = nil
|
|
}
|
|
clientErrCh <- err
|
|
}(localCon)
|
|
|
|
// Discard remoteCon to unblock the remote handler.
|
|
<-remoteConCh
|
|
|
|
// Simulate a client-side disconnect. All other parties (tsh proxy and
|
|
// remove listener) should disconnect as well.
|
|
t.Log("simulate client-side disconnect")
|
|
err = localCon.Close()
|
|
require.NoError(t, err)
|
|
|
|
for i := 0; i < 3; i++ {
|
|
select {
|
|
case err := <-proxyErrCh:
|
|
require.NoError(t, err)
|
|
case err := <-remoteErrCh:
|
|
require.NoError(t, err)
|
|
case err := <-clientErrCh:
|
|
require.NoError(t, err)
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("proxyConnection, client and server didn't disconnect within 5s after client connection was closed")
|
|
}
|
|
}
|
|
|
|
// Dial localSrv again. This should trigger proxyConnection and a dial to
|
|
// remoteSrv.
|
|
localCon, err = net.Dial("tcp", localSrv.Addr().String())
|
|
require.NoError(t, err)
|
|
go func(con net.Conn) {
|
|
_, err := io.Copy(io.Discard, con)
|
|
if err != nil && strings.Contains(err.Error(), "use of closed network connection") {
|
|
err = nil
|
|
}
|
|
clientErrCh <- err
|
|
}(localCon)
|
|
|
|
// Simulate a server-side disconnect. All other parties (tsh proxy and
|
|
// local client) should disconnect as well.
|
|
t.Log("simulate server-side disconnect")
|
|
remoteCon := <-remoteConCh
|
|
err = remoteCon.Close()
|
|
require.NoError(t, err)
|
|
|
|
for i := 0; i < 3; i++ {
|
|
select {
|
|
case err := <-proxyErrCh:
|
|
require.NoError(t, err)
|
|
case err := <-remoteErrCh:
|
|
require.NoError(t, err)
|
|
case err := <-clientErrCh:
|
|
require.NoError(t, err)
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("proxyConnection, client and server didn't disconnect within 5s after remote connection was closed")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestListenAndForwardCancel(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
testFun func(client *NodeClient, ctx context.Context, listener *wrappedListener)
|
|
}{
|
|
{
|
|
name: "listenAndForward",
|
|
testFun: func(client *NodeClient, ctx context.Context, listener *wrappedListener) {
|
|
client.listenAndForward(ctx, listener, "localAddr", "remoteAddr")
|
|
},
|
|
},
|
|
{
|
|
name: "dynamicListenAndForward",
|
|
testFun: func(client *NodeClient, ctx context.Context, listener *wrappedListener) {
|
|
client.dynamicListenAndForward(ctx, listener, "localAddr")
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
client := &NodeClient{
|
|
Client: &tracessh.Client{
|
|
Client: &ssh.Client{
|
|
Conn: &fakeSSHConn{},
|
|
},
|
|
},
|
|
Tracer: tracing.NoopProvider().Tracer("test"),
|
|
}
|
|
|
|
// Create two anchors. An "accept" anchor that unblocks once the listener has
|
|
// accepted a connection and an "unblock" anchor that unblocks when Accept
|
|
// unblocks.
|
|
acceptCh := make(chan struct{})
|
|
unblockCh := make(chan struct{})
|
|
|
|
// Create a new cancelable listener.
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
ln, err := newWrappedListener(acceptCh)
|
|
require.NoError(t, err)
|
|
|
|
// Start testFun (listenAndForward or dynamicListenAndForward)
|
|
// and close the unblock channel once "Accept" has unblocked.
|
|
go func() {
|
|
tt.testFun(client, ctx, ln)
|
|
close(unblockCh)
|
|
}()
|
|
|
|
// Block until "Accept" has been called. After this it is safe to assume the
|
|
// listener is accepting.
|
|
select {
|
|
case <-acceptCh:
|
|
case <-time.After(1 * time.Minute):
|
|
t.Fatal("Timed out waiting for Accept to be called.")
|
|
}
|
|
|
|
// At this point, "Accept" should still be blocking.
|
|
select {
|
|
case <-unblockCh:
|
|
t.Fatalf("Failed because Accept was unblocked.")
|
|
default:
|
|
}
|
|
|
|
// Cancel "Accept" to unblock it.
|
|
cancel()
|
|
|
|
// Verify that "Accept" has unblocked.
|
|
select {
|
|
case <-unblockCh:
|
|
case <-time.After(1 * time.Minute):
|
|
t.Fatal("Timed out waiting for Accept to unblock.")
|
|
}
|
|
})
|
|
}
|
|
|
|
}
|
|
|
|
func newTestListener(t *testing.T, handle func(net.Conn)) net.Listener {
|
|
l, err := net.Listen("tcp", "localhost:0")
|
|
require.NoError(t, err)
|
|
|
|
go func() {
|
|
for {
|
|
con, err := l.Accept()
|
|
if err != nil {
|
|
return
|
|
}
|
|
go handle(con)
|
|
}
|
|
}()
|
|
|
|
return l
|
|
}
|
|
|
|
// fakeSSHConn is a NOP connection that implements the ssh.Conn interface.
|
|
// Only used in tests.
|
|
type fakeSSHConn struct {
|
|
ssh.Conn
|
|
}
|
|
|
|
func (c *fakeSSHConn) Close() error {
|
|
return nil
|
|
}
|
|
|
|
// wrappedListener is a listener that uses a channel to notify the caller
|
|
// when "Accept" has been called.
|
|
type wrappedListener struct {
|
|
net.Listener
|
|
acceptCh chan struct{}
|
|
}
|
|
|
|
func newWrappedListener(acceptCh chan struct{}) (*wrappedListener, error) {
|
|
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
|
if err != nil {
|
|
return nil, trace.Wrap(err)
|
|
}
|
|
|
|
return &wrappedListener{
|
|
Listener: ln,
|
|
acceptCh: acceptCh,
|
|
}, nil
|
|
}
|
|
|
|
func (l wrappedListener) Accept() (net.Conn, error) {
|
|
close(l.acceptCh)
|
|
return l.Listener.Accept()
|
|
}
|
|
|
|
func TestLineLabeledWriter(t *testing.T) {
|
|
t.Parallel()
|
|
tests := []struct {
|
|
name string
|
|
inputs []string
|
|
expected string
|
|
}{
|
|
{
|
|
name: "typical input",
|
|
inputs: []string{"this is\nsome test\ninput"},
|
|
expected: "[label] this is\n[label] some test\n[label] input",
|
|
},
|
|
{
|
|
name: "don't add empty line at end",
|
|
inputs: []string{"dangling newline\n"},
|
|
expected: "[label] dangling newline\n",
|
|
},
|
|
{
|
|
name: "blank lines in middle",
|
|
inputs: []string{"this\n\nis\n\nsome input"},
|
|
expected: "[label] this\n[label] \n[label] is\n[label] \n[label] some input",
|
|
},
|
|
{
|
|
name: "line break between writes",
|
|
inputs: []string{"line 1\n", "line 2\n"},
|
|
expected: "[label] line 1\n[label] line 2\n",
|
|
},
|
|
{
|
|
name: "line break immediately on second write",
|
|
inputs: []string{"line 1", "\nline 2"},
|
|
expected: "[label] line 1\n[label] line 2",
|
|
},
|
|
{
|
|
name: "line continues between writes",
|
|
inputs: []string{"this is all ", "one continuous line"},
|
|
expected: "[label] this is all one continuous line",
|
|
},
|
|
}
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
var buf bytes.Buffer
|
|
w := newLineLabeledWriter(&buf, "label")
|
|
|
|
totalBytes := 0
|
|
expectedBytes := 0
|
|
for _, line := range tc.inputs {
|
|
n, err := w.Write([]byte(line))
|
|
assert.NoError(t, err)
|
|
totalBytes += n
|
|
expectedBytes += len(line)
|
|
}
|
|
assert.Equal(t, expectedBytes, totalBytes)
|
|
assert.Equal(t, tc.expected, buf.String())
|
|
})
|
|
}
|
|
}
|