server/ssh/http/connection_status_test.go
Rama master · réplica de lectura del repositorio original.
package http
import (
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/labstack/echo/v5"
"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/services"
servicemocks "github.com/shellhub-io/shellhub/server/api/services/mocks"
"github.com/shellhub-io/shellhub/server/api/store"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
)
func newConnectionTestServer(t *testing.T, service services.Service, cfg *Config) *echo.Echo {
t.Helper()
e := echo.New()
e.Binder = handlers.NewBinder()
e.Validator = handlers.NewValidator()
e.HTTPErrorHandler = handlers.NewErrors(nil)
h := &Handlers{Config: cfg, Dialer: nil, Service: service}
e.GET(HandleConnectionV1Path, h.HandleConnectionV1)
e.GET(HandleConnectionV2Path, h.HandleConnectionV2)
return e
}
const testDeviceUID = "0000000000000000000000000000000000000000000000000000000000000000"
func TestHandleConnectionV1StatusCodes(t *testing.T) {
tests := []struct {
description string
config *Config
tenant string
setupMock func(*servicemocks.MockService)
expectedStatus int
expectedBody string
}{
{
description: "an unknown device is refused with a server error",
config: &Config{RequireAcceptedTunnel: true}, //nolint:exhaustruct // the test sets only the field it exercises
tenant: "00000000-0000-4000-0000-000000000000",
setupMock: func(m *servicemocks.MockService) {
m.EXPECT().
GetDevice(mock.Anything, mock.Anything, models.UID(testDeviceUID)).
Return(nil, services.NewErrDeviceNotFound(models.UID(testDeviceUID), store.ErrNoDocuments)).
Once()
},
expectedStatus: http.StatusInternalServerError,
},
{
description: "a device that is not accepted is refused",
config: &Config{RequireAcceptedTunnel: true}, //nolint:exhaustruct // the test sets only the field it exercises
tenant: "00000000-0000-4000-0000-000000000000",
setupMock: func(m *servicemocks.MockService) {
m.EXPECT().
GetDevice(mock.Anything, mock.Anything, models.UID(testDeviceUID)).
Return(&models.Device{UID: testDeviceUID, Status: models.DeviceStatusPending}, nil). //nolint:exhaustruct // the test sets only the field it exercises
Once()
},
expectedStatus: http.StatusForbidden,
},
{
description: "the pre-0.15 fallback refuses an unknown device with a server error",
config: &Config{RequireAcceptedTunnel: false}, //nolint:exhaustruct // the test sets only the field it exercises
tenant: "",
setupMock: func(m *servicemocks.MockService) {
m.EXPECT().
GetDevice(mock.Anything, mock.Anything, models.UID(testDeviceUID)).
Return(nil, services.NewErrDeviceNotFound(models.UID(testDeviceUID), store.ErrNoDocuments)).
Once()
},
expectedStatus: http.StatusInternalServerError,
},
{
description: "a request that is not a websocket upgrade gets only the upgrader's response",
config: &Config{RequireAcceptedTunnel: false}, //nolint:exhaustruct // only the tunnel gate matters here
tenant: "00000000-0000-4000-0000-000000000000",
setupMock: func(*servicemocks.MockService) {},
expectedStatus: http.StatusBadRequest,
expectedBody: "Bad Request\n",
},
}
for _, tt := range tests {
t.Run(tt.description, func(t *testing.T) {
serviceMock := servicemocks.NewMockService(t)
tt.setupMock(serviceMock)
e := newConnectionTestServer(t, serviceMock, tt.config)
req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, HandleConnectionV1Path, nil)
req.Header.Set("X-Request-ID", "request-id")
req = identitytest.As(req, identity.Identity{DeviceUID: testDeviceUID, TenantID: tt.tenant})
rec := httptest.NewRecorder()
e.ServeHTTP(rec, req)
assert.Equal(t, tt.expectedStatus, rec.Code)
if tt.expectedBody != "" {
assert.Equal(t, tt.expectedBody, rec.Body.String())
}
})
}
}
// A device that does not exist is asked for once; anything else is retried,
// because the loopback client used to absorb a database blip inside the agent's
// request and the fleet reconnects in lockstep without it.
func TestHandleConnectionV1RetriesOnlyInfrastructureFailures(t *testing.T) {
tests := []struct {
description string
err error
expectedAttempts int
}{
{
description: "a missing device is not retried",
err: services.NewErrDeviceNotFound(models.UID(testDeviceUID), store.ErrNoDocuments),
expectedAttempts: 1,
},
{
description: "an infrastructure failure is retried",
err: services.NewErrDeviceNotFound(models.UID(testDeviceUID), errors.New("connection refused")),
expectedAttempts: deviceResolveAttempts,
},
}
for _, tt := range tests {
t.Run(tt.description, func(t *testing.T) {
serviceMock := servicemocks.NewMockService(t)
serviceMock.EXPECT().
GetDevice(mock.Anything, mock.Anything, models.UID(testDeviceUID)).
Return(nil, tt.err).
Times(tt.expectedAttempts)
e := newConnectionTestServer(t, serviceMock, &Config{RequireAcceptedTunnel: true}) //nolint:exhaustruct // the test sets only the field it exercises
req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, HandleConnectionV1Path, nil)
req.Header.Set("X-Request-ID", "request-id")
req = identitytest.As(req, identity.Identity{DeviceUID: testDeviceUID, TenantID: "00000000-0000-4000-0000-000000000000"})
rec := httptest.NewRecorder()
e.ServeHTTP(rec, req)
assert.Equal(t, http.StatusInternalServerError, rec.Code)
serviceMock.AssertExpectations(t)
})
}
}
func TestHandleConnectionV2StatusCodes(t *testing.T) {
tests := []struct {
description string
headers map[string]string
caller identity.Identity
expectedStatus int
expectedBody string
}{
{
description: "a malformed tenant is rejected before anything else",
headers: map[string]string{"X-Request-ID": "request-id"},
caller: identity.Identity{DeviceUID: testDeviceUID, TenantID: "not-a-uuid"},
expectedStatus: http.StatusBadRequest,
},
{
description: "a short device uid is rejected",
headers: map[string]string{"X-Request-ID": "request-id"},
caller: identity.Identity{DeviceUID: "too-short", TenantID: "00000000-0000-4000-0000-000000000000"},
expectedStatus: http.StatusBadRequest,
},
{
description: "a missing request id is rejected",
caller: identity.Identity{DeviceUID: testDeviceUID, TenantID: "00000000-0000-4000-0000-000000000000"},
expectedStatus: http.StatusBadRequest,
},
{
description: "a request that is not a websocket upgrade gets only the upgrader's response",
headers: map[string]string{"X-Request-ID": "request-id"},
caller: identity.Identity{DeviceUID: testDeviceUID, TenantID: "00000000-0000-4000-0000-000000000000"},
expectedStatus: http.StatusBadRequest,
expectedBody: "Bad Request\n",
},
}
for _, tt := range tests {
t.Run(tt.description, func(t *testing.T) {
e := newConnectionTestServer(t, nil, &Config{}) //nolint:exhaustruct // the test sets only the field it exercises
req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, HandleConnectionV2Path, nil)
for name, value := range tt.headers {
req.Header.Set(name, value)
}
req = identitytest.As(req, tt.caller)
rec := httptest.NewRecorder()
e.ServeHTTP(rec, req)
require.Equal(t, tt.expectedStatus, rec.Code)
if tt.expectedBody != "" {
assert.Equal(t, tt.expectedBody, rec.Body.String())
}
})
}
}