pkg/revdial/revdial_test.go
Rama master · réplica de lectura del repositorio original.
package revdial_test
import (
"bufio"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"testing/synctest"
"time"
"github.com/gorilla/websocket"
"github.com/shellhub-io/shellhub/pkg/clock"
"github.com/shellhub-io/shellhub/pkg/revdial"
"github.com/shellhub-io/shellhub/pkg/wsconnadapter"
log "github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
const (
connPath = "/ssh/revdial"
controlPath = "/control"
)
type controlMessage struct {
Command string `json:"command,omitempty"`
ConnPath string `json:"connPath,omitempty"`
Err string `json:"err,omitempty"`
}
type agent struct {
t *testing.T
ctl net.Conn
wsURL string
ready chan string
}
type control func(t *testing.T, mux *http.ServeMux, wsURL string) (serverConn net.Conn, agentConn net.Conn)
func websocketControl(t *testing.T, mux *http.ServeMux, wsURL string) (net.Conn, net.Conn) {
t.Helper()
accepted := make(chan net.Conn, 1)
mux.HandleFunc(controlPath, func(w http.ResponseWriter, r *http.Request) {
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
if err != nil {
t.Errorf("the control connection could not be upgraded: %v", err)
return
}
accepted <- wsconnadapter.New(conn)
})
peer, _, err := websocket.DefaultDialer.DialContext(t.Context(), wsURL+controlPath, nil) //nolint:bodyclose // gorilla/websocket documents that the handshake response body need not be closed.
require.NoError(t, err)
return <-accepted, wsconnadapter.New(peer)
}
func pipeControl(*testing.T, *http.ServeMux, string) (net.Conn, net.Conn) {
return net.Pipe()
}
func newDialer(t *testing.T) (*revdial.Dialer, *agent) {
t.Helper()
return newDialerOver(t, websocketControl)
}
func newDialerOver(t *testing.T, ctl control) (*revdial.Dialer, *agent) {
t.Helper()
mux := http.NewServeMux()
mux.Handle(connPath, revdial.ConnHandler(websocket.Upgrader{}))
srv := httptest.NewServer(mux)
t.Cleanup(srv.Close)
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
serverCtl, agentCtl := ctl(t, mux, wsURL)
t.Cleanup(func() {
_ = serverCtl.Close()
_ = agentCtl.Close()
})
logger := log.New()
logger.SetOutput(io.Discard)
d := revdial.NewDialer(log.NewEntry(logger), serverCtl, connPath, clock.New())
t.Cleanup(func() { _ = d.Close() })
a := &agent{
t: t,
ctl: agentCtl,
wsURL: wsURL,
ready: make(chan string, 1024),
}
go forwardControlCommand(a.ctl, "conn-ready", func(msg controlMessage) { a.ready <- msg.ConnPath })
return d, a
}
func forwardControlCommand(conn net.Conn, command string, deliver func(controlMessage)) {
reader := bufio.NewReader(conn)
for {
line, err := reader.ReadBytes('\n')
if err != nil {
return
}
var msg controlMessage
if err := json.Unmarshal(line, &msg); err != nil {
return
}
if msg.Command == command {
deliver(msg)
}
}
}
func (a *agent) awaitConnReady() (string, bool) {
select {
case path := <-a.ready:
return path, true
case <-time.After(5 * time.Second):
return "", assert.Fail(a.t, "the server never announced a connection to pick up")
}
}
func (a *agent) dialBack(ctx context.Context, path string) (net.Conn, bool) {
conn, _, err := websocket.DefaultDialer.DialContext(ctx, a.wsURL+path, nil) //nolint:bodyclose // gorilla/websocket documents that the handshake response body need not be closed.
if err != nil {
a.t.Errorf("the agent could not dial back %s: %v", path, err)
return nil, false
}
adapter := wsconnadapter.New(conn)
a.t.Cleanup(func() { _ = adapter.Close() })
return adapter, true
}
func (a *agent) failPickup(path string, reason string) bool {
line, err := json.Marshal(controlMessage{Command: "pickup-failed", ConnPath: path, Err: reason})
if err == nil {
_, err = a.ctl.Write(append(line, '\n'))
}
if err != nil {
a.t.Errorf("the agent could not report a failed pick-up: %v", err)
return false
}
return true
}
func dialWithin(t *testing.T, d *revdial.Dialer, timeout time.Duration) (net.Conn, error) {
t.Helper()
ctx, cancel := context.WithTimeout(t.Context(), timeout)
defer cancel()
return d.Dial(ctx)
}
func TestDialServesEveryConcurrentDialItsOwnConnection(t *testing.T) {
t.Parallel()
d, a := newDialer(t)
const dials = 16
go func() {
for range dials {
path, ok := a.awaitConnReady()
if !ok {
return
}
conn, ok := a.dialBack(t.Context(), path)
if !ok {
return
}
go func() {
_, _ = conn.Write([]byte(path))
}()
}
}()
var wg sync.WaitGroup
received := make(chan string, dials)
failures := make(chan error, dials)
for range dials {
wg.Go(func() {
conn, err := dialWithin(t, d, 5*time.Second)
if err != nil {
failures <- err
return
}
defer func() { _ = conn.Close() }()
if err := conn.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil { //nolint:forbidigo // a deadline needs the wall clock
failures <- err
return
}
buf := make([]byte, 256)
n, err := conn.Read(buf)
if err != nil {
failures <- err
return
}
received <- string(buf[:n])
})
}
wg.Wait()
close(failures)
close(received)
for err := range failures {
require.NoError(t, err, "every dial must get a connection or its own error")
}
paths := make(map[string]bool)
for path := range received {
assert.False(t, paths[path], "two dials received the same connection: %s", path)
paths[path] = true
}
assert.Len(t, paths, dials, "every dial must end up holding a distinct connection")
}
func TestDialReportsItsOwnPickupFailure(t *testing.T) {
t.Parallel()
d, a := newDialer(t)
go func() {
if path, ok := a.awaitConnReady(); ok {
a.failPickup(path, "the agent could not reach the server")
}
}()
conn, err := dialWithin(t, d, 5*time.Second)
require.ErrorContains(t, err, "the agent could not reach the server")
assert.Nil(t, conn)
}
func TestDialIgnoresAPickupFailureForADialNoLongerWaiting(t *testing.T) {
t.Parallel()
d, a := newDialerOver(t, pipeControl)
gaveUp, cancel := context.WithCancel(t.Context())
abandoned := make(chan error, 1)
go func() {
_, err := d.Dial(gaveUp)
abandoned <- err
}()
stalePath, ok := a.awaitConnReady()
require.True(t, ok)
cancel()
require.ErrorIs(t, <-abandoned, context.Canceled)
agentDone := make(chan struct{})
go func() {
defer close(agentDone)
path, ok := a.awaitConnReady()
if !ok || !a.failPickup(stalePath, "stale") {
return
}
a.dialBack(t.Context(), path)
}()
conn, err := dialWithin(t, d, 5*time.Second)
<-agentDone
require.NoError(t, err, "a failure for a dial that gave up must not fail the one still waiting")
require.NotNil(t, conn)
_ = conn.Close()
}
func TestDialClosesTheConnectionOfADialThatGaveUp(t *testing.T) {
t.Parallel()
d, a := newDialer(t)
ctx, cancel := context.WithCancel(t.Context())
dialed := make(chan error, 1)
go func() {
_, err := d.Dial(ctx)
dialed <- err
}()
path, ok := a.awaitConnReady()
require.True(t, ok)
cancel()
select {
case err := <-dialed:
require.ErrorIs(t, err, context.Canceled)
case <-time.After(5 * time.Second):
require.FailNow(t, "a cancelled dial must return")
}
late, ok := a.dialBack(t.Context(), path)
require.True(t, ok)
requireClosedByPeer(t, late)
}
func TestCloseReturnsWhenNobodyIsWatchingDone(t *testing.T) {
t.Parallel()
serverCtl, agentCtl := net.Pipe()
t.Cleanup(func() { _ = agentCtl.Close() })
logger := log.New()
logger.SetOutput(io.Discard)
d := revdial.NewDialer(log.NewEntry(logger), serverCtl, connPath, clock.New())
closed := make(chan struct{})
go func() {
_ = d.Close()
close(closed)
}()
select {
case <-closed:
case <-time.After(5 * time.Second):
require.FailNow(t, "Close must return even when nothing is receiving on Done")
}
select {
case <-d.Done():
default:
require.FailNow(t, "Done must be closed once Close returned")
}
}
func TestConnHandlerAnswersANonUpgradePickupWithOnlyTheUpgradersResponse(t *testing.T) {
t.Parallel()
d, a := newDialer(t)
dialed := make(chan error, 1)
go func() {
_, err := dialWithin(t, d, 5*time.Second)
dialed <- err
}()
path, ok := a.awaitConnReady()
require.True(t, ok)
httpURL := "http" + strings.TrimPrefix(a.wsURL, "ws")
req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, httpURL+path, nil)
require.NoError(t, err)
res, err := http.DefaultClient.Do(req)
require.NoError(t, err)
defer func() { _ = res.Body.Close() }()
body, err := io.ReadAll(res.Body)
require.NoError(t, err)
assert.Equal(t, http.StatusBadRequest, res.StatusCode)
assert.Equal(t, "Bad Request\n", string(body))
require.True(t, a.failPickup(path, "the agent gave up"))
require.Error(t, <-dialed)
}
type dialFunc func(context.Context, string) (*websocket.Conn, *http.Response, error)
type pickup struct {
path string
conn net.Conn
}
type pickupServer struct {
url string
accepted chan pickup
}
func newPickupServer(t *testing.T) *pickupServer {
t.Helper()
accepted := make(chan pickup, 64)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
if err != nil {
t.Errorf("the pick-up could not be upgraded: %v", err)
return
}
adapter := wsconnadapter.New(conn)
t.Cleanup(func() { _ = adapter.Close() })
accepted <- pickup{path: r.URL.Path, conn: adapter}
}))
t.Cleanup(srv.Close)
return &pickupServer{url: "ws" + strings.TrimPrefix(srv.URL, "http"), accepted: accepted}
}
func (s *pickupServer) dial(ctx context.Context, path string) (*websocket.Conn, *http.Response, error) {
return websocket.DefaultDialer.DialContext(ctx, s.url+path, nil)
}
type controlPeer struct {
t *testing.T
ctl net.Conn
failures chan controlMessage
}
func newListener(t *testing.T, dial dialFunc) (*revdial.Listener, *controlPeer) {
t.Helper()
serverCtl, agentCtl := net.Pipe()
t.Cleanup(func() {
_ = serverCtl.Close()
_ = agentCtl.Close()
})
s := &controlPeer{t: t, ctl: serverCtl, failures: make(chan controlMessage, 64)}
go forwardControlCommand(s.ctl, "pickup-failed", func(msg controlMessage) { s.failures <- msg })
ln := revdial.NewListener(agentCtl, dial)
t.Cleanup(func() { _ = ln.Close() })
return ln, s
}
func (s *controlPeer) announce(path string) {
line, err := json.Marshal(controlMessage{Command: "conn-ready", ConnPath: path})
require.NoError(s.t, err)
_, err = s.ctl.Write(append(line, '\n'))
require.NoError(s.t, err)
}
func receiveWithin[T any](t *testing.T, ch <-chan T, what string) T {
t.Helper()
select {
case v := <-ch:
return v
case <-time.After(5 * time.Second):
require.FailNow(t, what)
var zero T
return zero
}
}
func requireClosedByPeer(t *testing.T, conn net.Conn) {
t.Helper()
require.NoError(t, conn.SetReadDeadline(time.Now().Add(5*time.Second))) //nolint:forbidigo // a deadline needs the wall clock
_, err := conn.Read(make([]byte, 1))
require.Error(t, err)
if netErr, ok := errors.AsType[net.Error](err); ok {
require.False(t, netErr.Timeout(), "a connection nobody will claim must be closed, not left open")
}
}
func TestListenerClosesAPickupThatFinishesAfterClose(t *testing.T) {
t.Parallel()
for range 50 {
pickups := newPickupServer(t)
dialing := make(chan struct{})
release := make(chan struct{})
ln, peer := newListener(t, func(ctx context.Context, path string) (*websocket.Conn, *http.Response, error) {
close(dialing)
<-release
return pickups.dial(ctx, path)
})
peer.announce("/pickup")
receiveWithin(t, dialing, "the listener never started the pick-up")
require.NoError(t, ln.Close())
close(release)
requireClosedByPeer(t, receiveWithin(t, pickups.accepted, "the pick-up never reached the server").conn)
_, err := ln.Accept()
require.ErrorIs(t, err, revdial.ErrListenerClosed)
}
}
func overflowAcceptQueue(t *testing.T, pickups *pickupServer, peer *controlPeer) (map[string]net.Conn, controlMessage) {
t.Helper()
announced := revdial.AcceptQueueSize + 1
for i := range announced {
peer.announce(fmt.Sprintf("/pickup/%d", i))
}
byPath := make(map[string]net.Conn, announced)
for range announced {
p := receiveWithin(t, pickups.accepted, "a pick-up never reached the server")
byPath[p.path] = p.conn
}
failure := receiveWithin(t, peer.failures, "the pick-up beyond the accept queue was never reported as failed")
require.Contains(t, byPath, failure.ConnPath)
return byPath, failure
}
func TestListenerCloseClosesThePickupsWaitingToBeAccepted(t *testing.T) {
t.Parallel()
pickups := newPickupServer(t)
ln, peer := newListener(t, pickups.dial)
queued, failure := overflowAcceptQueue(t, pickups, peer)
delete(queued, failure.ConnPath)
require.NoError(t, ln.Close())
for _, conn := range queued {
requireClosedByPeer(t, conn)
}
}
func TestListenerFailsAPickupThatFindsTheAcceptQueueFull(t *testing.T) {
t.Parallel()
pickups := newPickupServer(t)
_, peer := newListener(t, pickups.dial)
byPath, failure := overflowAcceptQueue(t, pickups, peer)
requireClosedByPeer(t, byPath[failure.ConnPath])
select {
case extra := <-peer.failures:
require.FailNow(t, "only the pick-up beyond the accept queue may fail", "also failed: %s", extra.ConnPath)
case <-time.After(100 * time.Millisecond):
}
}
func TestListenerLeavesNoGoroutineBehindAfterClose(t *testing.T) {
t.Parallel()
synctest.Test(t, func(t *testing.T) {
pending := revdial.ControlWriteQueueSize + 4
dialing := make(chan struct{}, pending)
release := make(chan struct{})
ln, peer := newListener(t, func(context.Context, string) (*websocket.Conn, *http.Response, error) {
dialing <- struct{}{}
<-release
return nil, nil, errors.New("the server is unreachable")
})
for i := range pending {
peer.announce(fmt.Sprintf("/pickup/%d", i))
}
for range pending {
receiveWithin(t, dialing, "a pick-up never started")
}
require.NoError(t, ln.Close())
close(release)
})
}