Files
poseidon/internal/api/websocket_test.go
Maximilian Paß ab938bfc22 Refactor MemoryLeakTestSuite
as we identified two issues where the goroutine count from before differs from after the test.

1) It seemed like a Go runtime specific Goroutine appeared in rare cases before the test. To avoid this, we introduced a short timeout before looking up the Goroutines.
Another solution might be to do the lookup twice and check if the count matches.

2) A Goroutine that periodically monitors some storage unexpectedly got closed in rare cases. As we could not identify the cause for this, we removed the leaking Goroutines by properly cleaning up.
2024-02-28 11:52:51 +01:00

495 lines
17 KiB
Go

package api
import (
"bufio"
"context"
"crypto/tls"
"fmt"
"github.com/gorilla/mux"
"github.com/gorilla/websocket"
"github.com/openHPI/poseidon/internal/environment"
"github.com/openHPI/poseidon/internal/nomad"
"github.com/openHPI/poseidon/internal/runner"
"github.com/openHPI/poseidon/pkg/dto"
"github.com/openHPI/poseidon/tests"
"github.com/openHPI/poseidon/tests/helpers"
"github.com/sirupsen/logrus"
"github.com/sirupsen/logrus/hooks/test"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"io"
"net/http"
"net/http/httptest"
"net/url"
"testing"
"time"
)
func TestWebSocketTestSuite(t *testing.T) {
suite.Run(t, new(WebSocketTestSuite))
}
type WebSocketTestSuite struct {
tests.MemoryLeakTestSuite
router *mux.Router
executionID string
runner runner.Runner
apiMock *nomad.ExecutorAPIMock
server *httptest.Server
}
func (s *WebSocketTestSuite) SetupTest() {
s.MemoryLeakTestSuite.SetupTest()
runnerID := "runner-id"
s.runner, s.apiMock = newNomadAllocationWithMockedAPIClient(runnerID)
// default execution
s.executionID = tests.DefaultExecutionID
s.runner.StoreExecution(s.executionID, &executionRequestLs)
mockAPIExecuteLs(s.apiMock)
runnerManager := &runner.ManagerMock{}
runnerManager.On("Get", s.runner.ID()).Return(s.runner, nil)
s.router = NewRouter(runnerManager, nil)
s.server = httptest.NewServer(s.router)
}
func (s *WebSocketTestSuite) TearDownTest() {
defer s.MemoryLeakTestSuite.TearDownTest()
s.server.Close()
err := s.runner.Destroy(nil)
s.Require().NoError(err)
}
func (s *WebSocketTestSuite) TestWebsocketConnectionCanBeEstablished() {
wsURL, err := s.webSocketURL("ws", s.runner.ID(), s.executionID)
s.Require().NoError(err)
conn, _, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
s.Require().NoError(err)
<-time.After(tests.ShortTimeout)
err = conn.Close()
s.NoError(err)
}
func (s *WebSocketTestSuite) TestWebsocketReturns404IfExecutionDoesNotExist() {
wsURL, err := s.webSocketURL("ws", s.runner.ID(), "invalid-execution-id")
s.Require().NoError(err)
_, response, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
s.Require().ErrorIs(err, websocket.ErrBadHandshake)
s.Equal(http.StatusNotFound, response.StatusCode)
}
func (s *WebSocketTestSuite) TestWebsocketReturns400IfRequestedViaHttp() {
wsURL, err := s.webSocketURL("http", s.runner.ID(), s.executionID)
s.Require().NoError(err)
response, err := http.Get(wsURL.String())
s.Require().NoError(err)
// This functionality is implemented by the WebSocket library.
s.Equal(http.StatusBadRequest, response.StatusCode)
_, err = io.ReadAll(response.Body)
s.NoError(err)
}
func (s *WebSocketTestSuite) TestWebsocketConnection() {
s.runner.StoreExecution(s.executionID, &executionRequestHead)
mockAPIExecuteHead(s.apiMock)
wsURL, err := s.webSocketURL("ws", s.runner.ID(), s.executionID)
s.Require().NoError(err)
connection, _, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
s.Require().NoError(err)
err = connection.SetReadDeadline(time.Now().Add(5 * time.Second))
s.Require().NoError(err)
s.Run("Receives start message", func() {
message, err := helpers.ReceiveNextWebSocketMessage(connection)
s.Require().NoError(err)
s.Equal(dto.WebSocketMetaStart, message.Type)
})
s.Run("Executes the request in the runner", func() {
<-time.After(tests.ShortTimeout)
s.apiMock.AssertCalled(s.T(), "ExecuteCommand",
mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.AnythingOfType("bool"),
mock.Anything, mock.Anything, mock.Anything)
})
s.Run("Can send input", func() {
err = connection.WriteMessage(websocket.TextMessage, []byte("Hello World\n"))
s.Require().NoError(err)
})
messages, err := helpers.ReceiveAllWebSocketMessages(connection)
s.Require().Error(err)
s.True(websocket.IsCloseError(err, websocket.CloseNormalClosure))
s.Run("Receives output message", func() {
stdout, _, _ := helpers.WebSocketOutputMessages(messages)
s.Equal("Hello World", stdout)
})
s.Run("Receives exit message", func() {
controlMessages := helpers.WebSocketControlMessages(messages)
s.Require().Equal(1, len(controlMessages))
s.Equal(dto.WebSocketExit, controlMessages[0].Type)
})
}
func (s *WebSocketTestSuite) TestCancelWebSocketConnection() {
executionID := "sleeping-execution"
s.runner.StoreExecution(executionID, &executionRequestSleep)
canceled := mockAPIExecuteSleep(s.apiMock)
wsURL, err := webSocketURL("ws", s.server, s.router, s.runner.ID(), executionID)
s.Require().NoError(err)
connection, _, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
s.Require().NoError(err)
message, err := helpers.ReceiveNextWebSocketMessage(connection)
s.Require().NoError(err)
s.Equal(dto.WebSocketMetaStart, message.Type)
select {
case <-canceled:
s.Fail("ExecuteInteractively canceled unexpected")
default:
}
err = connection.WriteControl(websocket.CloseMessage,
websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
s.Require().NoError(err)
select {
case <-canceled:
case <-time.After(time.Second):
s.Fail("ExecuteInteractively not canceled")
}
}
func (s *WebSocketTestSuite) TestWebSocketConnectionTimeout() {
executionID := "time-out-execution"
limitExecution := executionRequestSleep
limitExecution.TimeLimit = 2
s.runner.StoreExecution(executionID, &limitExecution)
canceled := mockAPIExecuteSleep(s.apiMock)
wsURL, err := webSocketURL("ws", s.server, s.router, s.runner.ID(), executionID)
s.Require().NoError(err)
connection, _, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
s.Require().NoError(err)
message, err := helpers.ReceiveNextWebSocketMessage(connection)
s.Require().NoError(err)
s.Equal(dto.WebSocketMetaStart, message.Type)
select {
case <-canceled:
s.Fail("ExecuteInteractively canceled unexpected")
case <-time.After(time.Duration(limitExecution.TimeLimit-1) * time.Second):
<-time.After(time.Second)
}
select {
case <-canceled:
case <-time.After(time.Second):
s.Fail("ExecuteInteractively not canceled")
}
message, err = helpers.ReceiveNextWebSocketMessage(connection)
s.Require().NoError(err)
s.Equal(dto.WebSocketMetaTimeout, message.Type)
}
func (s *WebSocketTestSuite) TestWebsocketStdoutAndStderr() {
executionID := "ls-execution"
s.runner.StoreExecution(executionID, &executionRequestLs)
mockAPIExecuteLs(s.apiMock)
wsURL, err := webSocketURL("ws", s.server, s.router, s.runner.ID(), executionID)
s.Require().NoError(err)
connection, _, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
s.Require().NoError(err)
messages, err := helpers.ReceiveAllWebSocketMessages(connection)
s.Require().Error(err)
s.True(websocket.IsCloseError(err, websocket.CloseNormalClosure))
stdout, stderr, _ := helpers.WebSocketOutputMessages(messages)
s.Contains(stdout, "existing-file")
s.Contains(stderr, "non-existing-file")
}
func (s *WebSocketTestSuite) TestWebsocketError() {
executionID := "error-execution"
s.runner.StoreExecution(executionID, &executionRequestError)
mockAPIExecuteError(s.apiMock)
wsURL, err := webSocketURL("ws", s.server, s.router, s.runner.ID(), executionID)
s.Require().NoError(err)
connection, _, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
s.Require().NoError(err)
messages, err := helpers.ReceiveAllWebSocketMessages(connection)
s.Require().Error(err)
s.True(websocket.IsCloseError(err, websocket.CloseNormalClosure))
_, _, errMessages := helpers.WebSocketOutputMessages(messages)
s.Require().Equal(1, len(errMessages))
s.Equal("Error executing the request", errMessages[0])
}
func (s *WebSocketTestSuite) TestWebsocketNonZeroExit() {
executionID := "exit-execution"
s.runner.StoreExecution(executionID, &executionRequestExitNonZero)
mockAPIExecuteExitNonZero(s.apiMock)
wsURL, err := webSocketURL("ws", s.server, s.router, s.runner.ID(), executionID)
s.Require().NoError(err)
connection, _, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
s.Require().NoError(err)
messages, err := helpers.ReceiveAllWebSocketMessages(connection)
s.Require().Error(err)
s.True(websocket.IsCloseError(err, websocket.CloseNormalClosure))
controlMessages := helpers.WebSocketControlMessages(messages)
s.Equal(2, len(controlMessages))
s.Equal(&dto.WebSocketMessage{Type: dto.WebSocketExit, ExitCode: 42}, controlMessages[1])
}
func (s *MainTestSuite) TestWebsocketTLS() {
runnerID := "runner-id"
r, apiMock := newNomadAllocationWithMockedAPIClient(runnerID)
executionID := tests.DefaultExecutionID
r.StoreExecution(executionID, &executionRequestLs)
mockAPIExecuteLs(apiMock)
runnerManager := &runner.ManagerMock{}
runnerManager.On("Get", r.ID()).Return(r, nil)
router := NewRouter(runnerManager, nil)
server, err := helpers.StartTLSServer(s.T(), router)
s.Require().NoError(err)
defer server.Close()
wsURL, err := webSocketURL("wss", server, router, runnerID, executionID)
s.Require().NoError(err)
config := &tls.Config{RootCAs: nil, InsecureSkipVerify: true} //nolint:gosec // test needs self-signed cert
d := websocket.Dialer{TLSClientConfig: config}
connection, _, err := d.Dial(wsURL.String(), nil)
s.Require().NoError(err)
message, err := helpers.ReceiveNextWebSocketMessage(connection)
s.Require().NoError(err)
s.Equal(dto.WebSocketMetaStart, message.Type)
_, err = helpers.ReceiveAllWebSocketMessages(connection)
s.Require().Error(err)
s.True(websocket.IsCloseError(err, websocket.CloseNormalClosure))
s.NoError(r.Destroy(nil))
}
func (s *MainTestSuite) TestWebSocketProxyStopsReadingTheWebSocketAfterClosingIt() {
apiMock := &nomad.ExecutorAPIMock{}
executionID := tests.DefaultExecutionID
r, wsURL, cleanup := newRunnerWithNotMockedRunnerManager(s, apiMock, executionID)
defer cleanup()
logger, hook := test.NewNullLogger()
log = logger.WithField("pkg", "api")
r.StoreExecution(executionID, &executionRequestHead)
mockAPIExecute(apiMock, &executionRequestHead,
func(_ string, ctx context.Context, _ string, _ bool, _ io.Reader, _, _ io.Writer) (int, error) {
return 0, nil
})
connection, _, err := websocket.DefaultDialer.Dial(wsURL.String(), nil)
s.Require().NoError(err)
_, err = helpers.ReceiveAllWebSocketMessages(connection)
s.Require().Error(err)
s.True(websocket.IsCloseError(err, websocket.CloseNormalClosure))
for _, logMsg := range hook.Entries {
if logMsg.Level < logrus.InfoLevel {
s.Fail(logMsg.Message)
}
}
}
// --- Test suite specific test helpers ---
func newNomadAllocationWithMockedAPIClient(runnerID string) (runner.Runner, *nomad.ExecutorAPIMock) {
executorAPIMock := &nomad.ExecutorAPIMock{}
executorAPIMock.On("DeleteJob", mock.AnythingOfType("string")).Return(nil)
manager := &runner.ManagerMock{}
manager.On("Return", mock.Anything).Return(nil)
r := runner.NewNomadJob(runnerID, nil, executorAPIMock, nil)
return r, executorAPIMock
}
func newRunnerWithNotMockedRunnerManager(s *MainTestSuite, apiMock *nomad.ExecutorAPIMock, executionID string) (
r runner.Runner, wsURL *url.URL, cleanup func()) {
s.T().Helper()
apiMock.On("MarkRunnerAsUsed", mock.AnythingOfType("string"), mock.AnythingOfType("int")).Return(nil)
apiMock.On("LoadRunnerIDs", mock.AnythingOfType("string")).Return([]string{}, nil)
apiMock.On("DeleteJob", mock.AnythingOfType("string")).Return(nil)
apiMock.On("RegisterRunnerJob", mock.AnythingOfType("*api.Job")).Return(nil)
call := apiMock.On("WatchEventStream", mock.Anything, mock.Anything, mock.Anything)
call.Run(func(args mock.Arguments) {
<-s.TestCtx.Done()
call.ReturnArguments = mock.Arguments{nil}
})
runnerManager := runner.NewNomadRunnerManager(apiMock, s.TestCtx)
router := NewRouter(runnerManager, nil)
s.ExpectedGoroutineIncrease++ // The server is not closing properly. Therefore, we don't even try.
server := httptest.NewServer(router)
runnerID := tests.DefaultRunnerID
runnerJob := runner.NewNomadJob(runnerID, nil, apiMock, nil)
e, err := environment.NewNomadEnvironment(0, apiMock, "job \"template-0\" {}")
s.Require().NoError(err)
eID, err := nomad.EnvironmentIDFromRunnerID(runnerID)
s.Require().NoError(err)
e.SetID(eID)
e.SetPrewarmingPoolSize(0)
runnerManager.StoreEnvironment(e)
e.AddRunner(runnerJob)
r, err = runnerManager.Claim(e.ID(), int(tests.DefaultTestTimeout.Seconds()))
s.Require().NoError(err)
wsURL, err = webSocketURL("ws", server, router, r.ID(), executionID)
s.Require().NoError(err)
return r, wsURL, func() {
err = r.Destroy(tests.ErrCleanupDestroyReason)
s.NoError(err)
err = e.Delete(tests.ErrCleanupDestroyReason)
s.NoError(err)
}
}
func webSocketURL(scheme string, server *httptest.Server, router *mux.Router,
runnerID string, executionID string,
) (*url.URL, error) {
websocketURL, err := url.Parse(server.URL)
if err != nil {
return nil, err
}
path, err := router.Get(WebsocketPath).URL(RunnerIDKey, runnerID)
if err != nil {
return nil, err
}
websocketURL.Scheme = scheme
websocketURL.Path = path.Path
websocketURL.RawQuery = fmt.Sprintf("executionID=%s", executionID)
return websocketURL, nil
}
func (s *WebSocketTestSuite) webSocketURL(scheme, runnerID, executionID string) (*url.URL, error) {
return webSocketURL(scheme, s.server, s.router, runnerID, executionID)
}
var executionRequestLs = dto.ExecutionRequest{Command: "ls"}
// mockAPIExecuteLs mocks the ExecuteCommand of an ExecutorApi to act as if
// 'ls existing-file non-existing-file' was executed.
func mockAPIExecuteLs(api *nomad.ExecutorAPIMock) {
mockAPIExecute(api, &executionRequestLs,
func(_ string, _ context.Context, _ string, _ bool, _ io.Reader, stdout, stderr io.Writer) (int, error) {
_, _ = stdout.Write([]byte("existing-file\n"))
_, _ = stderr.Write([]byte("ls: cannot access 'non-existing-file': No such file or directory\n"))
return 0, nil
})
}
var executionRequestHead = dto.ExecutionRequest{Command: "head -n 1"}
// mockAPIExecuteHead mocks the ExecuteCommand of an ExecutorApi to act as if 'head -n 1' was executed.
func mockAPIExecuteHead(api *nomad.ExecutorAPIMock) {
mockAPIExecute(api, &executionRequestHead,
func(_ string, _ context.Context, _ string, _ bool,
stdin io.Reader, stdout io.Writer, stderr io.Writer,
) (int, error) {
scanner := bufio.NewScanner(stdin)
for !scanner.Scan() {
scanner = bufio.NewScanner(stdin)
}
_, _ = stdout.Write(scanner.Bytes())
return 0, nil
})
}
var executionRequestSleep = dto.ExecutionRequest{Command: "sleep infinity"}
// mockAPIExecuteSleep mocks the ExecuteCommand method of an ExecutorAPI to sleep
// until the execution receives a SIGQUIT.
func mockAPIExecuteSleep(api *nomad.ExecutorAPIMock) <-chan bool {
canceled := make(chan bool, 1)
mockAPIExecute(api, &executionRequestSleep,
func(_ string, ctx context.Context, _ string, _ bool,
stdin io.Reader, stdout io.Writer, stderr io.Writer,
) (int, error) {
var err error
buffer := make([]byte, 1) //nolint:makezero // if the length is zero, the Read call never reads anything
for n := 0; !(n == 1 && buffer[0] == runner.SIGQUIT); n, err = stdin.Read(buffer) {
if err != nil {
return 0, fmt.Errorf("error while reading stdin: %w", err)
}
}
close(canceled)
return 0, ctx.Err()
})
return canceled
}
var executionRequestError = dto.ExecutionRequest{Command: "error"}
// mockAPIExecuteError mocks the ExecuteCommand method of an ExecutorApi to return an error.
func mockAPIExecuteError(api *nomad.ExecutorAPIMock) {
mockAPIExecute(api, &executionRequestError,
func(_ string, _ context.Context, _ string, _ bool, _ io.Reader, _, _ io.Writer) (int, error) {
return 0, tests.ErrDefault
})
}
var executionRequestExitNonZero = dto.ExecutionRequest{Command: "exit 42"}
// mockAPIExecuteExitNonZero mocks the ExecuteCommand method of an ExecutorApi to exit with exit status 42.
func mockAPIExecuteExitNonZero(api *nomad.ExecutorAPIMock) {
mockAPIExecute(api, &executionRequestExitNonZero,
func(_ string, _ context.Context, _ string, _ bool, _ io.Reader, _, _ io.Writer) (int, error) {
return 42, nil
})
}
// mockAPIExecute mocks the ExecuteCommand method of an ExecutorApi to call the given method run when the command
// corresponding to the given ExecutionRequest is called.
func mockAPIExecute(api *nomad.ExecutorAPIMock, request *dto.ExecutionRequest,
run func(runnerId string, ctx context.Context, command string, tty bool,
stdin io.Reader, stdout, stderr io.Writer) (int, error)) {
tests.RemoveMethodFromMock(&api.Mock, "ExecuteCommand")
call := api.On("ExecuteCommand",
mock.AnythingOfType("string"),
mock.Anything,
request.FullCommand(),
mock.AnythingOfType("bool"),
mock.AnythingOfType("bool"),
mock.Anything,
mock.Anything,
mock.Anything)
call.Run(func(args mock.Arguments) {
exit, err := run(args.Get(0).(string),
args.Get(1).(context.Context),
args.Get(2).(string),
args.Get(3).(bool),
args.Get(5).(io.Reader),
args.Get(6).(io.Writer),
args.Get(7).(io.Writer))
call.ReturnArguments = mock.Arguments{exit, err}
})
}