server/ssh/web/web_test.go
Rama master · réplica de lectura del repositorio original.
package web
import (
"crypto/ed25519"
"crypto/rand"
"crypto/rsa"
"encoding/json"
"encoding/pem"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
"github.com/labstack/echo/v5"
"github.com/shellhub-io/shellhub/pkg/api/authorizer"
"github.com/shellhub-io/shellhub/pkg/api/identity"
"github.com/shellhub-io/shellhub/pkg/api/identity/identitytest"
"github.com/shellhub-io/shellhub/pkg/models"
"github.com/shellhub-io/shellhub/server/api/pkg/echo/handlers"
"github.com/shellhub-io/shellhub/server/api/pkg/gateway"
"github.com/shellhub-io/shellhub/server/api/pkg/origin"
servicemocks "github.com/shellhub-io/shellhub/server/api/services/mocks"
"github.com/shellhub-io/shellhub/server/api/store"
"github.com/shellhub-io/shellhub/server/ssh/pkg/webhandoff"
"github.com/sirupsen/logrus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"golang.org/x/crypto/ssh"
"golang.org/x/net/websocket"
)
func writeHostKey(t *testing.T) string {
t.Helper()
_, key, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
block, err := ssh.MarshalPrivateKey(key, "")
require.NoError(t, err)
path := filepath.Join(t.TempDir(), "ssh.key")
require.NoError(t, os.WriteFile(path, pem.EncodeToMemory(block), 0o600))
return path
}
// TestExitLogLevel verifies that exitLogLevel returns logrus.WarnLevel for
// expected/banner-derived errors and logrus.ErrorLevel for genuine server faults.
func TestExitLogLevel(t *testing.T) {
tests := []struct {
description string
err error
expected logrus.Level
}{
{
description: "ErrConnect uses Warn",
err: ErrConnect,
expected: logrus.WarnLevel,
},
{
description: "ErrAccessDenied uses Warn",
err: ErrAccessDenied,
expected: logrus.WarnLevel,
},
{
description: "ErrInvalidSSHID uses Warn",
err: ErrInvalidSSHID,
expected: logrus.WarnLevel,
},
{
description: "ErrAuthentication uses Warn",
err: ErrAuthentication,
expected: logrus.WarnLevel,
},
{
description: "ErrGetAuth uses Warn",
err: ErrGetAuth,
expected: logrus.WarnLevel,
},
{
description: "ErrFindDevice uses Warn",
err: ErrFindDevice,
expected: logrus.WarnLevel,
},
{
description: "ErrForbiddenPublicKey uses Warn",
err: ErrForbiddenPublicKey,
expected: logrus.WarnLevel,
},
{
description: "ErrBridgeCredentialsNotFound uses Warn",
err: ErrBridgeCredentialsNotFound,
expected: logrus.WarnLevel,
},
{
description: "ErrWebSocketGetToken uses Warn",
err: ErrWebSocketGetToken,
expected: logrus.WarnLevel,
},
{
description: "ErrWebSocketGetDimensions uses Warn",
err: ErrWebSocketGetDimensions,
expected: logrus.WarnLevel,
},
{
description: "ErrWebSocketGetIP uses Warn",
err: ErrWebSocketGetIP,
expected: logrus.WarnLevel,
},
{
description: "ErrSession uses Error",
err: ErrSession,
expected: logrus.ErrorLevel,
},
{
description: "ErrPty uses Error",
err: ErrPty,
expected: logrus.ErrorLevel,
},
{
description: "ErrShell uses Error",
err: ErrShell,
expected: logrus.ErrorLevel,
},
{
description: "unrecognised error uses Error",
err: ErrPublicKey,
expected: logrus.ErrorLevel,
},
}
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
assert.Equal(t, test.expected, exitLogLevel(test.err))
})
}
}
func TestNewSSHServerBridge_CredentialsNotFound(t *testing.T) {
e := echo.New()
require.NoError(t, NewSSHServerBridge(e, nil, nil, webhandoff.NewStore(), &Config{HostKeyFile: writeHostKey(t)}))
server := httptest.NewServer(e)
defer server.Close()
assert.NotPanics(t, func() {
assert.Contains(t, openBridge(t, server.URL, "nonexistent"), ErrBridgeCredentialsNotFound.Error())
}, "handler must not panic when credentials are not found")
}
const tenantID = "00000000-0000-4000-0000-000000000000"
func deviceInNamespace(inScope bool) (*models.Device, error) {
if !inScope {
return nil, store.ErrNoDocuments
}
return &models.Device{UID: "device-uid", TenantID: tenantID}, nil //nolint:exhaustruct // the handler only checks that the device resolves
}
func TestWebSessionRouteRequiresTheConnectPermission(t *testing.T) {
tests := []struct {
description string
role string
deviceInScope bool
expectedCode int
}{
{
description: "observer is refused",
role: authorizer.RoleObserver.String(),
expectedCode: http.StatusForbidden,
},
{
description: "an identity carrying no role is refused",
role: "",
expectedCode: http.StatusForbidden,
},
{
description: "operator is served a token for a device in their namespace",
role: authorizer.RoleOperator.String(),
deviceInScope: true,
expectedCode: http.StatusOK,
},
{
description: "a device outside the caller's namespace is refused",
role: authorizer.RoleOperator.String(),
expectedCode: http.StatusForbidden,
},
}
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
service := servicemocks.NewMockService(t)
if test.role == authorizer.RoleOperator.String() {
service.On("GetDevice", mock.Anything, mock.Anything, models.UID("device-uid")).
Return(deviceInNamespace(test.deviceInScope)).Once()
}
e := echo.New()
e.HTTPErrorHandler = handlers.NewErrors(nil)
e.Use(gateway.WithContext())
require.NoError(t, NewSSHServerBridge(e, nil, service, webhandoff.NewStore(), &Config{HostKeyFile: writeHostKey(t)}))
server := httptest.NewServer(identitytest.Serve(e, identity.Identity{
TenantID: tenantID,
Role: authorizer.RoleFromString(test.role),
}))
defer server.Close()
body := strings.NewReader(`{"device":"device-uid","username":"root","password":"secret"}`)
req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, server.URL+WebSessionRoute, body)
require.NoError(t, err)
req.Header.Set("Content-Type", "application/json")
res, err := http.DefaultClient.Do(req)
require.NoError(t, err)
defer res.Body.Close() //nolint:errcheck
assert.Equal(t, test.expectedCode, res.StatusCode)
if test.expectedCode != http.StatusOK {
var refusal struct {
Message string `json:"message"`
Error string `json:"error"`
}
require.NoError(t, json.NewDecoder(res.Body).Decode(&refusal))
assert.True(t, refusal.Message != "" || refusal.Error != "",
"a refusal must name what failed, through the error handler's message or this route's own error")
return
}
var success struct {
Token string `json:"token"`
}
require.NoError(t, json.NewDecoder(res.Body).Decode(&success))
assert.NotEmpty(t, success.Token)
})
}
}
func TestNewSSHServerBridge_TokenOpensOneBridge(t *testing.T) {
service := servicemocks.NewMockService(t)
service.On("GetDevice", mock.Anything, mock.Anything, models.UID("device-uid")).
Return(deviceInNamespace(true)).Once()
service.On("GetDevice", mock.Anything, mock.Anything, models.UID("device-uid")).
Return(nil, store.ErrNoDocuments).Once()
server := startBridge(t, service)
status, reply := postSession(t, server.URL, `{"device":"device-uid","username":"root","fingerprint":"fingerprint"}`)
require.Equal(t, http.StatusOK, status)
assert.Contains(t, openBridge(t, server.URL, reply.Token), ErrGetAuth.Error())
assert.Contains(t, openBridge(t, server.URL, reply.Token), ErrBridgeCredentialsNotFound.Error())
}
func TestNewSSHServerBridge_FailedOpenSpendsTheToken(t *testing.T) {
tests := []struct {
description string
query string
resolved bool
expected error
}{
{
description: "an invalid terminal size",
query: "cols=wide&rows=24",
resolved: true,
expected: ErrWebSocketGetDimensions,
},
{
description: "a client origin the router did not resolve",
query: "cols=80&rows=24",
resolved: false,
expected: ErrWebSocketGetIP,
},
}
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
service := servicemocks.NewMockService(t)
service.On("GetDevice", mock.Anything, mock.Anything, models.UID("device-uid")).
Return(deviceInNamespace(true)).Once()
server := startBridgeResolving(t, service, test.resolved)
status, reply := postSession(t, server.URL, `{"device":"device-uid","username":"root","fingerprint":"fingerprint"}`)
require.Equal(t, http.StatusOK, status)
assert.Contains(t, openBridgeWith(t, server.URL, "token="+reply.Token+"&"+test.query), test.expected.Error())
assert.Contains(t, openBridge(t, server.URL, reply.Token), ErrBridgeCredentialsNotFound.Error())
})
}
}
func TestWebSessionRouteRefusesAPasswordItCannotEncrypt(t *testing.T) {
service := servicemocks.NewMockService(t)
service.On("GetDevice", mock.Anything, mock.Anything, models.UID("device-uid")).
Return(deviceInNamespace(true)).Once()
server := startBridge(t, service)
status, reply := postSession(t, server.URL, `{"device":"device-uid","username":"root","password":"`+strings.Repeat("a", 191)+`"}`)
assert.Equal(t, http.StatusBadRequest, status)
assert.Contains(t, reply.Error, ErrCredentialsEncryptPassword.Error())
assert.Contains(t, reply.Error, rsa.ErrMessageTooLong.Error())
}
func startBridge(t *testing.T, service *servicemocks.MockService) *httptest.Server {
t.Helper()
return startBridgeResolving(t, service, true)
}
func startBridgeResolving(t *testing.T, service *servicemocks.MockService, resolved bool) *httptest.Server {
t.Helper()
e := echo.New()
e.HTTPErrorHandler = handlers.NewErrors(nil)
e.Use(gateway.WithContext())
if resolved {
e.Pre(func(next echo.HandlerFunc) echo.HandlerFunc {
return func(c *echo.Context) error {
req := c.Request()
c.SetRequest(req.WithContext(origin.NewContext(req.Context(), origin.Loopback().Resolve(req))))
return next(c)
}
})
}
require.NoError(t, NewSSHServerBridge(e, nil, service, webhandoff.NewStore(), &Config{HostKeyFile: writeHostKey(t)}))
server := httptest.NewServer(identitytest.Serve(e, identity.Identity{TenantID: tenantID, Role: authorizer.RoleOperator}))
t.Cleanup(server.Close)
return server
}
type sessionReply struct {
Token string `json:"token"`
Error string `json:"error"`
}
func postSession(t *testing.T, serverURL, body string) (int, sessionReply) {
t.Helper()
req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, serverURL+WebSessionRoute, strings.NewReader(body))
require.NoError(t, err)
req.Header.Set("Content-Type", "application/json")
res, err := http.DefaultClient.Do(req)
require.NoError(t, err)
defer res.Body.Close() //nolint:errcheck // the body is decoded before the close runs, and a failed close cannot change it
var reply sessionReply
require.NoError(t, json.NewDecoder(res.Body).Decode(&reply))
return res.StatusCode, reply
}
func openBridge(t *testing.T, serverURL, token string) string {
t.Helper()
return openBridgeWith(t, serverURL, "token="+token+"&cols=80&rows=24")
}
func openBridgeWith(t *testing.T, serverURL, query string) string {
t.Helper()
cfg, err := websocket.NewConfig("ws"+strings.TrimPrefix(serverURL, "http")+WebsocketSSHBridgeRoute+"?"+query, serverURL)
require.NoError(t, err)
conn, err := websocket.DialConfig(cfg)
require.NoError(t, err)
defer conn.Close() //nolint:errcheck // the bridge has already answered, and a failed close does not change it
var raw []byte
require.NoError(t, websocket.Message.Receive(conn, &raw))
var msg Message
require.NoError(t, json.Unmarshal(raw, &msg))
require.Equal(t, messageKindError, msg.Kind)
data, ok := msg.Data.(string)
require.True(t, ok)
return data
}
func TestNewSSHServerBridge_MissingHostKey(t *testing.T) {
err := NewSSHServerBridge(echo.New(), nil, nil, webhandoff.NewStore(), &Config{HostKeyFile: filepath.Join(t.TempDir(), "absent.key")})
assert.ErrorIs(t, err, ErrBridgeReadHostKey)
}