serverimpl_test.go

package server

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"google.golang.org/protobuf/proto"

	sharedinternal "github.com/open-telemetry/opamp-go/internal"
	"github.com/open-telemetry/opamp-go/internal/testhelpers"
	"github.com/open-telemetry/opamp-go/protobufs"
	"github.com/open-telemetry/opamp-go/server/types"
)

func startServer(t *testing.T, settings *StartSettings) *server {
	srv := New(&sharedinternal.NopLogger{})
	require.NotNil(t, srv)
	if settings.ListenEndpoint == "" {
		// Find an avaiable port to listne on.
		settings.ListenEndpoint = testhelpers.GetAvailableLocalAddress()
	}
	if settings.ListenPath == "" {
		settings.ListenPath = "/"
	}
	err := srv.Start(*settings)
	require.NoError(t, err)

	return srv
}

func dialClient(serverSettings *StartSettings) (*websocket.Conn, *http.Response, error) {
	srvUrl := "ws://" + serverSettings.ListenEndpoint + serverSettings.ListenPath
	dailer := websocket.DefaultDialer
	dailer.EnableCompression = serverSettings.EnableCompression
	return websocket.DefaultDialer.Dial(srvUrl, nil)
}

func TestServerStartStop(t *testing.T) {
	srv := startServer(t, &StartSettings{})

	err := srv.Start(StartSettings{})
	assert.ErrorIs(t, err, errAlreadyStarted)

	err = srv.Stop(context.Background())
	assert.NoError(t, err)
}

func TestServerStartRejectConnection(t *testing.T) {
	callbacks := CallbacksStruct{
		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
			// Reject the incoming HTTP connection.
			return types.ConnectionResponse{
				Accept:             false,
				HTTPStatusCode:     503,
				HTTPResponseHeader: map[string]string{"Retry-After": "30"},
			}
		},
	}

	// Start a Server.
	settings := &StartSettings{Settings: Settings{Callbacks: callbacks}}
	srv := startServer(t, settings)
	defer srv.Stop(context.Background())

	// Try to connect to the Server.
	conn, resp, err := dialClient(settings)

	// Verify that the connection is rejected and rejection data is available to the client.
	assert.Nil(t, conn)
	assert.Error(t, err)
	assert.EqualValues(t, 503, resp.StatusCode)
	assert.EqualValues(t, "30", resp.Header.Get("Retry-After"))
}

func eventually(t *testing.T, f func() bool) {
	assert.Eventually(t, f, 5*time.Second, 10*time.Millisecond)
}

func TestServerStartAcceptConnection(t *testing.T) {
	connectedCalled := int32(0)
	connectionCloseCalled := int32(0)
	var srvConn types.Connection
	callbacks := CallbacksStruct{
		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
			return types.ConnectionResponse{Accept: true}
		},
		OnConnectedFunc: func(conn types.Connection) {
			srvConn = conn
			atomic.StoreInt32(&connectedCalled, 1)
		},
		OnConnectionCloseFunc: func(conn types.Connection) {
			atomic.StoreInt32(&connectionCloseCalled, 1)
			assert.EqualValues(t, srvConn, conn)
		},
	}

	// Start a Server.
	settings := &StartSettings{Settings: Settings{Callbacks: callbacks}}
	srv := startServer(t, settings)
	defer srv.Stop(context.Background())

	// Connect to the Server.
	conn, resp, err := dialClient(settings)

	// Verify that the connection is successful.
	assert.NoError(t, err)
	assert.NotNil(t, conn)
	require.NotNil(t, resp)
	assert.EqualValues(t, 101, resp.StatusCode)
	eventually(t, func() bool { return atomic.LoadInt32(&connectedCalled) == 1 })
	assert.True(t, atomic.LoadInt32(&connectionCloseCalled) == 0)

	// Verify that the RemoteAddr is correct.
	require.Equal(t, conn.LocalAddr().String(), srvConn.RemoteAddr().String())

	// Close the connection from client side.
	conn.Close()

	// Verify that the Server receives the closing notification.
	eventually(t, func() bool { return atomic.LoadInt32(&connectionCloseCalled) == 1 })
}

func TestDisconnectHttpConnection(t *testing.T) {
	// Verify Disconnect() results with Invalid HTTP Connection error
	err := httpConnection{}.Disconnect()
	assert.Error(t, err)
	assert.Equal(t, ErrInvalidHTTPConnection, err)
}

func TestDisconnectWSConnection(t *testing.T) {
	connectionCloseCalled := int32(0)
	callback := CallbacksStruct{
		OnConnectionCloseFunc: func(conn types.Connection) {
			atomic.StoreInt32(&connectionCloseCalled, 1)
		},
	}

	// Start a Server.
	settings := &StartSettings{Settings: Settings{Callbacks: callback}}
	srv := startServer(t, settings)
	defer srv.Stop(context.Background())

	// Connect to the Server.
	conn, _, err := dialClient(settings)

	// Verify that the connection is successful.
	assert.NoError(t, err)
	assert.True(t, atomic.LoadInt32(&connectionCloseCalled) == 0)

	// Close connection from server side
	srvConn := wsConnection{wsConn: conn}
	err = srvConn.Disconnect()
	assert.NoError(t, err)

	// Verify connection disconnected from server side
	eventually(t, func() bool { return atomic.LoadInt32(&connectionCloseCalled) == 1 })
	// Waiting for wsConnection to fail ReadMessage() over a Disconnected communication
	eventually(t, func() bool {
		_, _, err := conn.ReadMessage()
		return err != nil
	})
}

func TestServerReceiveSendMessage(t *testing.T) {
	var rcvMsg atomic.Value
	callbacks := CallbacksStruct{
		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
			return types.ConnectionResponse{Accept: true}
		},
		OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
			// Remember received message.
			rcvMsg.Store(message)

			// Send a response.
			response := protobufs.ServerToAgent{
				InstanceUid:  message.InstanceUid,
				Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
			}
			return &response
		},
	}

	// Start a Server.
	settings := &StartSettings{Settings: Settings{Callbacks: callbacks}}
	srv := startServer(t, settings)
	defer srv.Stop(context.Background())

	// Connect using a WebSocket client.
	conn, _, _ := dialClient(settings)
	require.NotNil(t, conn)
	defer conn.Close()

	// Send a message to the Server.
	sendMsg := protobufs.AgentToServer{
		InstanceUid: "12345678",
	}
	bytes, err := proto.Marshal(&sendMsg)
	require.NoError(t, err)
	err = conn.WriteMessage(websocket.BinaryMessage, bytes)
	require.NoError(t, err)

	// Wait until Server receives the message.
	eventually(t, func() bool { return rcvMsg.Load() != nil })
	assert.True(t, proto.Equal(rcvMsg.Load().(proto.Message), &sendMsg))

	// Read Server's response.
	mt, bytes, err := conn.ReadMessage()
	require.NoError(t, err)
	require.EqualValues(t, websocket.BinaryMessage, mt)

	// Decode the response.

	// Must start with a zero byte header until the end of grace period that ends Feb 1, 2023.
	require.EqualValues(t, 0, bytes[0])

	var response protobufs.ServerToAgent
	err = sharedinternal.DecodeWSMessage(bytes, &response)
	require.NoError(t, err)

	// Verify the response.
	assert.EqualValues(t, sendMsg.InstanceUid, response.InstanceUid)
	assert.EqualValues(t, protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus, response.Capabilities)
}

func TestServerReceiveSendMessageWithCompression(t *testing.T) {
	// Use highly compressible config body.
	uncompressedCfg := []byte(strings.Repeat("test", 10000))
	tests := []bool{false, true}
	for _, withCompression := range tests {
		t.Run(fmt.Sprintf("%v", withCompression), func(t *testing.T) {
			var rcvMsg atomic.Value
			callbacks := CallbacksStruct{
				OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
					return types.ConnectionResponse{Accept: true}
				},
				OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
					// Remember received message.
					rcvMsg.Store(message)

					// Send a response.
					response := protobufs.ServerToAgent{
						InstanceUid:  message.InstanceUid,
						Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
						RemoteConfig: &protobufs.AgentRemoteConfig{
							Config: &protobufs.AgentConfigMap{
								ConfigMap: map[string]*protobufs.AgentConfigFile{
									"": {Body: uncompressedCfg},
								},
							},
						},
					}
					return &response
				},
			}

			// Start a Server.
			settings := &StartSettings{Settings: Settings{Callbacks: callbacks, EnableCompression: withCompression}}
			srv := startServer(t, settings)
			defer srv.Stop(context.Background())

			// We use a transparent TCP proxy to be able to count the actual bytes transferred so that
			// we can test the number of actual bytes vs number of expected bytes with and without compression.
			proxy := testhelpers.NewProxy(settings.ListenEndpoint)
			assert.NoError(t, proxy.Start())

			serverSettings := *settings
			serverSettings.ListenEndpoint = proxy.IncomingEndpoint()
			// Connect using a WebSocket client.
			conn, _, _ := dialClient(&serverSettings)
			require.NotNil(t, conn)
			defer conn.Close()

			// Send a message to the Server.
			sendMsg := protobufs.AgentToServer{
				InstanceUid: "10000000",
				EffectiveConfig: &protobufs.EffectiveConfig{
					ConfigMap: &protobufs.AgentConfigMap{
						ConfigMap: map[string]*protobufs.AgentConfigFile{
							"": {Body: uncompressedCfg},
						},
					},
				},
			}
			bytes, err := proto.Marshal(&sendMsg)
			require.NoError(t, err)
			err = conn.WriteMessage(websocket.BinaryMessage, bytes)
			require.NoError(t, err)

			// Wait until Server receives the message.
			eventually(t, func() bool { return rcvMsg.Load() != nil })
			assert.True(t, proto.Equal(rcvMsg.Load().(proto.Message), &sendMsg))

			// Read Server's response.
			mt, bytes, err := conn.ReadMessage()
			require.NoError(t, err)
			require.EqualValues(t, websocket.BinaryMessage, mt)

			// Decode the response.

			// Must start with a zero byte header until the end of grace period that ends Feb 1, 2023.
			require.EqualValues(t, 0, bytes[0])

			var response protobufs.ServerToAgent
			err = sharedinternal.DecodeWSMessage(bytes, &response)
			require.NoError(t, err)

			fmt.Printf("sent %d, received %d\n", proxy.ClientToServerBytes(), proxy.ServerToClientBytes())

			// Verify the response.
			assert.EqualValues(t, sendMsg.InstanceUid, response.InstanceUid)
			assert.EqualValues(t, protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus, response.Capabilities)
			if withCompression {
				// With compression the entire bytes exchanged should be less than the config body.
				// This is only possible if there is any compression happening.
				assert.Less(t, proxy.ClientToServerBytes(), len(uncompressedCfg))
				assert.Less(t, proxy.ServerToClientBytes(), len(uncompressedCfg))
			} else {
				// Without compression the entire bytes exchanged should be more than the config body.
				assert.Greater(t, proxy.ClientToServerBytes(), len(uncompressedCfg))
				assert.Greater(t, proxy.ServerToClientBytes(), len(uncompressedCfg))
			}
		})
	}
}

func TestServerReceiveSendMessagePlainHTTP(t *testing.T) {
	var rcvMsg atomic.Value
	var onConnectedCalled, onCloseCalled int32
	callbacks := CallbacksStruct{
		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
			return types.ConnectionResponse{Accept: true}
		},
		OnConnectedFunc: func(conn types.Connection) {
			atomic.StoreInt32(&onConnectedCalled, 1)
		},
		OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
			// Remember received message.
			rcvMsg.Store(message)

			// Send a response.
			response := protobufs.ServerToAgent{
				InstanceUid:  message.InstanceUid,
				Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
			}
			return &response
		},
		OnConnectionCloseFunc: func(conn types.Connection) {
			atomic.StoreInt32(&onCloseCalled, 1)
		},
	}

	// Start a Server.
	settings := &StartSettings{Settings: Settings{Callbacks: callbacks}}
	srv := startServer(t, settings)
	defer srv.Stop(context.Background())

	// Send a message to the Server.
	sendMsg := protobufs.AgentToServer{
		InstanceUid: "12345678",
	}
	b, err := proto.Marshal(&sendMsg)
	require.NoError(t, err)
	resp, err := http.Post("http://"+settings.ListenEndpoint+settings.ListenPath, contentTypeProtobuf, bytes.NewReader(b))
	require.NoError(t, err)

	// Wait until Server receives the message.
	eventually(t, func() bool { return rcvMsg.Load() != nil })
	assert.True(t, atomic.LoadInt32(&onConnectedCalled) == 1)

	// Verify the received message is what was sent.
	assert.True(t, proto.Equal(rcvMsg.Load().(proto.Message), &sendMsg))

	// Read Server's response.
	b, err = io.ReadAll(resp.Body)
	require.NoError(t, err)

	assert.EqualValues(t, http.StatusOK, resp.StatusCode)
	assert.EqualValues(t, contentTypeProtobuf, resp.Header.Get(headerContentType))

	// Decode the response.
	var response protobufs.ServerToAgent
	err = proto.Unmarshal(b, &response)
	require.NoError(t, err)

	// Verify the response.
	assert.EqualValues(t, sendMsg.InstanceUid, response.InstanceUid)
	assert.EqualValues(t, protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus, response.Capabilities)

	eventually(t, func() bool { return atomic.LoadInt32(&onCloseCalled) == 1 })
}

func TestServerAttachAcceptConnection(t *testing.T) {
	connectedCalled := int32(0)
	connectionCloseCalled := int32(0)
	var srvConn types.Connection
	callbacks := CallbacksStruct{
		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
			return types.ConnectionResponse{Accept: true}
		},
		OnConnectedFunc: func(conn types.Connection) {
			atomic.StoreInt32(&connectedCalled, 1)
			srvConn = conn
		},
		OnConnectionCloseFunc: func(conn types.Connection) {
			atomic.StoreInt32(&connectionCloseCalled, 1)
			assert.EqualValues(t, srvConn, conn)
		},
	}

	// Prepare to attach OpAMP Server to an HTTP Server created separately.
	settings := Settings{Callbacks: callbacks}
	srv := New(&sharedinternal.NopLogger{})
	require.NotNil(t, srv)
	handlerFunc, err := srv.Attach(settings)
	require.NoError(t, err)

	// Create an HTTP Server and make it handle OpAMP connections.
	mux := http.NewServeMux()
	path := "/opamppath"
	mux.HandleFunc(path, handlerFunc)
	hs := httptest.NewServer(mux)
	defer hs.Close()

	// Connect using WebSocket client.
	srvUrl := "ws://" + hs.Listener.Addr().String() + path
	conn, resp, err := websocket.DefaultDialer.Dial(srvUrl, nil)

	// Verify that connection is successful.
	assert.NoError(t, err)
	assert.NotNil(t, conn)
	assert.EqualValues(t, 101, resp.StatusCode)
	eventually(t, func() bool { return atomic.LoadInt32(&connectedCalled) == 1 })
	assert.True(t, atomic.LoadInt32(&connectionCloseCalled) == 0)

	conn.Close()
	eventually(t, func() bool { return atomic.LoadInt32(&connectionCloseCalled) == 1 })
}

func TestServerAttachSendMessagePlainHTTP(t *testing.T) {
	connectedCalled := int32(0)
	connectionCloseCalled := int32(0)
	var srvConn types.Connection
	callbacks := CallbacksStruct{
		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
			return types.ConnectionResponse{Accept: true}
		},
		OnConnectedFunc: func(conn types.Connection) {
			atomic.StoreInt32(&connectedCalled, 1)
			srvConn = conn
		},
		OnConnectionCloseFunc: func(conn types.Connection) {
			atomic.StoreInt32(&connectionCloseCalled, 1)
			assert.EqualValues(t, srvConn, conn)
		},
	}

	// Prepare to attach OpAMP Server to an HTTP Server created separately.
	settings := Settings{Callbacks: callbacks}
	srv := New(&sharedinternal.NopLogger{})
	require.NotNil(t, srv)
	handlerFunc, err := srv.Attach(settings)
	require.NoError(t, err)

	// Create an HTTP Server and make it handle OpAMP connections.
	mux := http.NewServeMux()
	path := "/opamppath"
	mux.HandleFunc(path, handlerFunc)
	hs := httptest.NewServer(mux)
	defer hs.Close()

	//verify message can be sent
	//Send a message to the Server.
	sendMsg := protobufs.AgentToServer{
		InstanceUid: "12345678",
	}
	b, err := proto.Marshal(&sendMsg)
	require.NoError(t, err)
	resp, err := http.Post("http://"+hs.Listener.Addr().String()+path, contentTypeProtobuf, bytes.NewReader(b))
	require.NoError(t, err)
	require.NotNil(t, resp)
}

func TestServerHonoursClientRequestContentEncoding(t *testing.T) {

	hc := http.Client{}
	var rcvMsg atomic.Value
	var onConnectedCalled, onCloseCalled int32
	callbacks := CallbacksStruct{
		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
			return types.ConnectionResponse{Accept: true}
		},
		OnConnectedFunc: func(conn types.Connection) {
			atomic.StoreInt32(&onConnectedCalled, 1)
		},
		OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
			// Remember received message.
			rcvMsg.Store(message)

			// Send a response.
			response := protobufs.ServerToAgent{
				InstanceUid:  message.InstanceUid,
				Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
			}
			return &response
		},
		OnConnectionCloseFunc: func(conn types.Connection) {
			atomic.StoreInt32(&onCloseCalled, 1)
		},
	}

	// Start a Server.
	settings := &StartSettings{Settings: Settings{Callbacks: callbacks}}
	srv := startServer(t, settings)
	defer srv.Stop(context.Background())

	// Send a message to the Server.
	sendMsg := protobufs.AgentToServer{
		InstanceUid: "12345678",
	}
	b, err := proto.Marshal(&sendMsg)
	require.NoError(t, err)

	// gzip compress the request payload
	b, err = compressGzip(b)
	require.NoError(t, err)

	req, err := http.NewRequest("POST", "http://"+settings.ListenEndpoint+settings.ListenPath, bytes.NewReader(b))
	req.Header.Set(headerContentType, contentTypeProtobuf)
	req.Header.Set(headerContentEncoding, contentEncodingGzip)
	resp, err := hc.Do(req)
	require.NoError(t, err)

	// Wait until Server receives the message.
	eventually(t, func() bool { return rcvMsg.Load() != nil })
	assert.True(t, atomic.LoadInt32(&onConnectedCalled) == 1)

	// Verify the received message is what was sent.
	assert.True(t, proto.Equal(rcvMsg.Load().(proto.Message), &sendMsg))

	// Read Server's response.
	b, err = io.ReadAll(resp.Body)
	require.NoError(t, err)

	assert.EqualValues(t, http.StatusOK, resp.StatusCode)
	assert.EqualValues(t, contentTypeProtobuf, resp.Header.Get(headerContentType))

	// Decode the response.
	var response protobufs.ServerToAgent
	err = proto.Unmarshal(b, &response)
	require.NoError(t, err)

	// Verify the response.
	assert.EqualValues(t, sendMsg.InstanceUid, response.InstanceUid)
	assert.EqualValues(t, protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus, response.Capabilities)

	eventually(t, func() bool { return atomic.LoadInt32(&onCloseCalled) == 1 })
}

func TestServerHonoursAcceptEncoding(t *testing.T) {

	hc := http.Client{}
	var rcvMsg atomic.Value
	var onConnectedCalled, onCloseCalled int32
	callbacks := CallbacksStruct{
		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
			return types.ConnectionResponse{Accept: true}
		},
		OnConnectedFunc: func(conn types.Connection) {
			atomic.StoreInt32(&onConnectedCalled, 1)
		},
		OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
			// Remember received message.
			rcvMsg.Store(message)

			// Send a response.
			response := protobufs.ServerToAgent{
				InstanceUid:  message.InstanceUid,
				Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
			}
			return &response
		},
		OnConnectionCloseFunc: func(conn types.Connection) {
			atomic.StoreInt32(&onCloseCalled, 1)
		},
	}

	// Start a Server.
	settings := &StartSettings{Settings: Settings{Callbacks: callbacks}}
	srv := startServer(t, settings)
	defer srv.Stop(context.Background())

	// Send a message to the Server.
	sendMsg := protobufs.AgentToServer{
		InstanceUid: "12345678",
	}
	b, err := proto.Marshal(&sendMsg)
	require.NoError(t, err)
	req, err := http.NewRequest("POST", "http://"+settings.ListenEndpoint+settings.ListenPath, bytes.NewReader(b))
	req.Header.Set(headerContentType, contentTypeProtobuf)
	req.Header.Set(headerAcceptEncoding, contentEncodingGzip)
	resp, err := hc.Do(req)
	require.NoError(t, err)

	// Wait until Server receives the message.
	eventually(t, func() bool { return rcvMsg.Load() != nil })
	assert.True(t, atomic.LoadInt32(&onConnectedCalled) == 1)

	// Verify the received message is what was sent.
	assert.True(t, proto.Equal(rcvMsg.Load().(proto.Message), &sendMsg))

	// Read Server's response.
	b, err = io.ReadAll(resp.Body)
	require.NoError(t, err)
	// Decompress the gzip response
	b, err = decompressGzip(b)
	require.NoError(t, err)

	assert.EqualValues(t, http.StatusOK, resp.StatusCode)
	assert.EqualValues(t, contentTypeProtobuf, resp.Header.Get(headerContentType))

	// Decode the response.
	var response protobufs.ServerToAgent
	err = proto.Unmarshal(b, &response)
	require.NoError(t, err)

	// Verify the response.
	assert.EqualValues(t, sendMsg.InstanceUid, response.InstanceUid)
	assert.EqualValues(t, protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus, response.Capabilities)

	eventually(t, func() bool { return atomic.LoadInt32(&onCloseCalled) == 1 })
}

func TestDecodeMessage(t *testing.T) {
	msgsToTest := []*protobufs.AgentToServer{
		{}, // Empty message
		{
			InstanceUid: "abcd",
			SequenceNum: 123,
		},
	}

	// Try with and without header byte. This is only necessary until the
	// end of grace period that ends Feb 1, 2023. After that the header is
	// no longer optional.
	withHeaderTests := []bool{false, true}

	for _, msg := range msgsToTest {
		for _, withHeader := range withHeaderTests {
			bytes, err := proto.Marshal(msg)
			require.NoError(t, err)

			if withHeader {
				// Prepend zero header byte.
				bytes = append([]byte{0}, bytes...)
			}

			var decoded protobufs.AgentToServer
			err = sharedinternal.DecodeWSMessage(bytes, &decoded)
			require.NoError(t, err)

			assert.True(t, proto.Equal(msg, &decoded))
		}
	}
}