Function testProxyWebsocket

proxy/proxy_test.go:194–256

(proxy connection.OriginProxy)

}

func testProxyWebsocket(proxy connection.OriginProxy) func(t *testing.T) {
	return func(t *testing.T) {
		// WSRoute is a websocket echo handler
		const testTimeout = 5 * time.Second * 1000
		ctx, cancel := context.WithTimeout(t.Context(), testTimeout)
		defer cancel()
		readPipe, writePipe := io.Pipe()
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("http://localhost:8080%s", hello.WSRoute), readPipe)
		req.Header.Set("Sec-Websocket-Key", "dGhlIHNhbXBsZSBub25jZQ==")
		req.Header.Set("Connection", "Upgrade")
		req.Header.Set("Upgrade", "websocket")
		responseWriter := newMockWSRespWriter(nil)

		finished := make(chan struct{})

		errGroup, ctx := errgroup.WithContext(ctx)
		errGroup.Go(func() error {
			log := zerolog.Nop()
			err = proxy.ProxyHTTP(responseWriter, tracing.NewTracedHTTPRequest(req, 0, &log), true)
			require.NoError(t, err)

			require.Equal(t, http.StatusSwitchingProtocols, responseWriter.Code)
			return nil
		})

		errGroup.Go(func() error {
			select {
			case <-finished:
			case <-ctx.Done():
			}
			if ctx.Err() == context.DeadlineExceeded {
				t.Errorf("Test timed out")
				readPipe.Close()
				writePipe.Close()
				responseWriter.Close()
			}
			return nil
		})

		msg := []byte("test websocket")
		err = wsutil.WriteClientText(writePipe, msg)
		require.NoError(t, err)

		// ReadServerText reads next data message from rw, considering that caller represents proxy side.
		returnedMsg, err := wsutil.ReadServerText(responseWriter.respBody())
		require.NoError(t, err)
		require.Equal(t, msg, returnedMsg)

		err = wsutil.WriteClientBinary(writePipe, msg)
		require.NoError(t, err)

		returnedMsg, err = wsutil.ReadServerBinary(responseWriter.respBody())
		require.NoError(t, err)
		require.Equal(t, msg, returnedMsg)

		_ = readPipe.Close()
		_ = writePipe.Close()
		_ = responseWriter.Close()