session_test.go 5.26 KB
Newer Older
Francisco Javier López's avatar
Francisco Javier López committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
package session

import (
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"

	"gitlab.com/gitlab-org/gitlab-runner/session/terminal"
)

func TestExec(t *testing.T) {
	cases := []struct {
		name               string
		authorization      string
		attachTerminal     bool
		isWebsocketUpgrade bool
		connectionErr      error
		expectedStatusCode int
	}{
		{
			name:               "Interactive terminal not available",
			attachTerminal:     false,
			isWebsocketUpgrade: true,
			expectedStatusCode: http.StatusServiceUnavailable,
		},
		{
			name:               "Request is not websocket upgraded",
			attachTerminal:     true,
			isWebsocketUpgrade: false,
			expectedStatusCode: http.StatusMethodNotAllowed,
		},
		{
			name:               "Request no authorized",
			attachTerminal:     true,
			isWebsocketUpgrade: true,
			authorization:      "invalidToken",
			expectedStatusCode: http.StatusUnauthorized,
		},
		{
			name:               "Terminal connected successfully",
			attachTerminal:     true,
			isWebsocketUpgrade: true,
			authorization:      "validToken",
			expectedStatusCode: http.StatusOK,
		},
		{
			name:               "Failed to start terminal",
			attachTerminal:     true,
			isWebsocketUpgrade: true,
			authorization:      "validToken",
			connectionErr:      errors.New("failed to connect to terminal"),
			expectedStatusCode: http.StatusInternalServerError,
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			session, err := NewSession(nil)
			require.NoError(t, err)
			session.Token = "validToken"

Tomasz Maczukin's avatar
Tomasz Maczukin committed
67 68
			mockTerminalConn := new(terminal.MockConn)
			defer mockTerminalConn.AssertExpectations(t)
Francisco Javier López's avatar
Francisco Javier López committed
69

Tomasz Maczukin's avatar
Tomasz Maczukin committed
70 71 72 73 74 75 76 77 78 79 80
			if c.connectionErr == nil && c.authorization == "validToken" && c.isWebsocketUpgrade {
				mockTerminalConn.On("Start", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Once()
				mockTerminalConn.On("Close").Return(nil).Once()
			}

			mockTerminal := new(terminal.MockInteractiveTerminal)
			defer mockTerminal.AssertExpectations(t)

			if c.authorization == "validToken" && c.isWebsocketUpgrade {
				mockTerminal.On("Connect").Return(mockTerminalConn, c.connectionErr).Once()
			}
Francisco Javier López's avatar
Francisco Javier López committed
81 82

			if c.attachTerminal {
Tomasz Maczukin's avatar
Tomasz Maczukin committed
83
				session.SetInteractiveTerminal(mockTerminal)
Francisco Javier López's avatar
Francisco Javier López committed
84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109
			}

			req := httptest.NewRequest(http.MethodPost, session.Endpoint+"/exec", nil)

			if c.isWebsocketUpgrade {
				req.Header.Add("Connection", "upgrade")
				req.Header.Add("Upgrade", "websocket")
			}
			req.Header.Add("Authorization", c.authorization)

			w := httptest.NewRecorder()

			session.Mux().ServeHTTP(w, req)

			resp := w.Result()

			assert.Equal(t, c.expectedStatusCode, resp.StatusCode)
		})
	}
}

func TestDoNotAllowMultipleConnections(t *testing.T) {
	session, err := NewSession(nil)
	require.NoError(t, err)
	session.Token = "validToken"

Tomasz Maczukin's avatar
Tomasz Maczukin committed
110 111
	mockTerminalConn := new(terminal.MockConn)
	defer mockTerminalConn.AssertExpectations(t)
Francisco Javier López's avatar
Francisco Javier López committed
112

Tomasz Maczukin's avatar
Tomasz Maczukin committed
113 114 115
	mockTerminal := new(terminal.MockInteractiveTerminal)
	defer mockTerminal.AssertExpectations(t)
	mockTerminal.On("Connect").Return(mockTerminalConn, nil).Once()
Francisco Javier López's avatar
Francisco Javier López committed
116

Tomasz Maczukin's avatar
Tomasz Maczukin committed
117
	session.SetInteractiveTerminal(mockTerminal)
Francisco Javier López's avatar
Francisco Javier López committed
118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151

	// Simulating another connection has already started.
	conn, err := session.newTerminalConn()
	require.NotNil(t, conn)
	require.NoError(t, err)

	req := httptest.NewRequest(http.MethodPost, session.Endpoint+"/exec", nil)
	req.Header.Add("Connection", "upgrade")
	req.Header.Add("Upgrade", "websocket")
	req.Header.Add("Authorization", "validToken")

	w := httptest.NewRecorder()
	session.Mux().ServeHTTP(w, req)
	resp := w.Result()
	assert.Equal(t, http.StatusLocked, resp.StatusCode)
}

func TestConnected(t *testing.T) {
	sess, err := NewSession(nil)
	require.NoError(t, err)

	assert.False(t, sess.Connected())
	sess.terminalConn = &terminal.MockConn{}
	assert.True(t, sess.Connected())
}

func TestKill(t *testing.T) {
	sess, err := NewSession(nil)
	require.NoError(t, err)

	// No connection attached
	err = sess.Kill()
	assert.NoError(t, err)

Tomasz Maczukin's avatar
Tomasz Maczukin committed
152 153
	mockConn := new(terminal.MockConn)
	defer mockConn.AssertExpectations(t)
Francisco Javier López's avatar
Francisco Javier López committed
154 155
	mockConn.On("Close").Return(nil).Once()

Tomasz Maczukin's avatar
Tomasz Maczukin committed
156
	sess.terminalConn = mockConn
Francisco Javier López's avatar
Francisco Javier López committed
157 158 159 160 161 162 163 164 165 166

	err = sess.Kill()
	assert.NoError(t, err)
	assert.Nil(t, sess.terminalConn)
}

func TestKillFailedToClose(t *testing.T) {
	sess, err := NewSession(nil)
	require.NoError(t, err)

Tomasz Maczukin's avatar
Tomasz Maczukin committed
167 168
	mockConn := new(terminal.MockConn)
	defer mockConn.AssertExpectations(t)
Francisco Javier López's avatar
Francisco Javier López committed
169 170
	mockConn.On("Close").Return(errors.New("some error")).Once()

Tomasz Maczukin's avatar
Tomasz Maczukin committed
171
	sess.terminalConn = mockConn
Francisco Javier López's avatar
Francisco Javier López committed
172 173 174 175 176 177 178

	err = sess.Kill()
	assert.Error(t, err)

	// Even though an error occurred closing it still is removed.
	assert.Nil(t, sess.terminalConn)
}
179 180 181 182 183 184 185 186 187 188 189 190 191

type fakeTerminalConn struct {
	commands []string
}

func (fakeTerminalConn) Close() error {
	return nil
}

func (fakeTerminalConn) Start(w http.ResponseWriter, r *http.Request, timeoutCh, disconnectCh chan error) {
}

func TestCloseTerminalConn(t *testing.T) {
192
	conn := &fakeTerminalConn{
193 194 195
		commands: []string{"command", "-c", "random"},
	}

196
	mockConn := new(terminal.MockConn)
Tomasz Maczukin's avatar
Tomasz Maczukin committed
197
	defer mockConn.AssertExpectations(t)
198 199
	mockConn.On("Close").Return(nil).Once()

200 201 202 203
	sess, err := NewSession(nil)
	sess.terminalConn = conn
	require.NoError(t, err)

204
	sess.closeTerminalConn(mockConn)
205 206
	assert.NotNil(t, sess.terminalConn)

207
	sess.closeTerminalConn(conn)
208 209
	assert.Nil(t, sess.terminalConn)
}