upstream-centrifugo/internal/proxy/publish_handler_test.go

package proxy

import (
	"context"
	"encoding/base64"
	"fmt"
	"log"
	"net/http"
	"testing"
	"time"

	"github.com/centrifugal/centrifugo/v6/internal/configtypes"
	"github.com/centrifugal/centrifugo/v6/internal/tools"

	"github.com/centrifugal/centrifuge"
	"github.com/stretchr/testify/require"
)

type grpcPublishHandleTestCase struct {
	*tools.CommonGRPCProxyTestCase
	publishProxyHandler *PublishHandler
	channelOpts         configtypes.ChannelOptions
}

func newPublishHandleGRPCTestCase(ctx context.Context, proxyGRPCServer proxyGRPCTestServer, opts configtypes.ChannelOptions) grpcPublishHandleTestCase {
	commonProxyTestCase := tools.NewCommonGRPCProxyTestCase(ctx, proxyGRPCServer)

	publishProxy, err := NewGRPCPublishProxy("default", getTestGrpcProxy(commonProxyTestCase))
	if err != nil {
		log.Fatalln("could not create grpc publish proxy: ", err)
	}

	publishProxyHandler := NewPublishHandler(PublishHandlerConfig{
		Proxies: map[string]PublishProxy{
			"test": publishProxy,
		},
	})

	return grpcPublishHandleTestCase{commonProxyTestCase, publishProxyHandler, opts}
}

type httpPublishHandleTestCase struct {
	*tools.CommonHTTPProxyTestCase
	publishProxyHandler *PublishHandler
	channelOpts         configtypes.ChannelOptions
}

func newPublishHandleHTTPTestCase(ctx context.Context, endpoint string, opts configtypes.ChannelOptions) httpPublishHandleTestCase {
	commonProxyTestCase := tools.NewCommonHTTPProxyTestCase(ctx)

	publishProxy, err := NewHTTPPublishProxy(getTestHttpProxy(commonProxyTestCase, endpoint))
	if err != nil {
		log.Fatalln("could not create http publish proxy: ", err)
	}

	publishProxyHandler := NewPublishHandler(PublishHandlerConfig{
		Proxies: map[string]PublishProxy{
			"test": publishProxy,
		},
	})

	return httpPublishHandleTestCase{commonProxyTestCase, publishProxyHandler, opts}
}

type publishHandleTestCase struct {
	publishProxyHandler *PublishHandler
	protocol            string
	node                *centrifuge.Node
	client              *centrifuge.Client
	channelOpts         configtypes.ChannelOptions
}

func (c publishHandleTestCase) invokeHandle() (reply centrifuge.PublishReply, err error) {
	publishHandler := c.publishProxyHandler.Handle(c.node)
	reply, err = publishHandler(c.client, centrifuge.PublishEvent{}, c.channelOpts, PerCallData{})

	return reply, err
}

func newPublishHandleTestCases(httpTestCase httpPublishHandleTestCase, grpcTestCase grpcPublishHandleTestCase) []publishHandleTestCase {
	return []publishHandleTestCase{
		{
			publishProxyHandler: grpcTestCase.publishProxyHandler,
			node:                grpcTestCase.Node,
			client:              grpcTestCase.Client,
			channelOpts:         grpcTestCase.channelOpts,
			protocol:            "grpc",
		},
		{
			publishProxyHandler: httpTestCase.publishProxyHandler,
			node:                httpTestCase.Node,
			client:              httpTestCase.Client,
			channelOpts:         httpTestCase.channelOpts,
			protocol:            "http",
		},
	}
}

func TestHandlePublishWithResult(t *testing.T) {
	customData := "test"
	customDataB64 := base64.StdEncoding.EncodeToString([]byte(customData))
	chOpts := configtypes.ChannelOptions{
		HistoryTTL:          configtypes.Duration(1 * time.Second),
		HistorySize:         1,
		PublishProxyEnabled: true,
		PublishProxyName:    "test",
	}

	opts := proxyGRPCTestServerOptions{
		B64Data: customDataB64,
	}
	grpcTestCase := newPublishHandleGRPCTestCase(context.Background(), newProxyGRPCTestServer("result", opts), chOpts)
	defer grpcTestCase.Teardown()

	httpTestCase := newPublishHandleHTTPTestCase(context.Background(), "/publish", chOpts)
	httpTestCase.Mux.HandleFunc("/publish", func(w http.ResponseWriter, req *http.Request) {
		_, _ = w.Write([]byte(fmt.Sprintf(`{"result": {"b64data": "%s"}}`, customDataB64)))
	})
	defer httpTestCase.Teardown()

	cases := newPublishHandleTestCases(httpTestCase, grpcTestCase)
	for _, c := range cases {
		reply, err := c.invokeHandle()
		require.NoError(t, err, c.protocol)
		require.Equal(t, uint64(1), reply.Result.Offset, c.protocol)
		require.NotEmpty(t, reply.Result.Epoch, c.protocol)
	}
}

func TestHandlePublishWithSkipHistory(t *testing.T) {
	customData := "test"
	customDataB64 := base64.StdEncoding.EncodeToString([]byte(customData))
	chOpts := configtypes.ChannelOptions{
		HistoryTTL:          configtypes.Duration(1 * time.Second),
		HistorySize:         1,
		PublishProxyEnabled: true,
		PublishProxyName:    "test",
	}

	opts := proxyGRPCTestServerOptions{
		B64Data: customDataB64,
	}
	grpcTestCase := newPublishHandleGRPCTestCase(context.Background(), newProxyGRPCTestServer("skip history", opts), chOpts)
	defer grpcTestCase.Teardown()

	httpTestCase := newPublishHandleHTTPTestCase(context.Background(), "/publish", chOpts)
	httpTestCase.Mux.HandleFunc("/publish", func(w http.ResponseWriter, req *http.Request) {
		_, _ = w.Write([]byte(fmt.Sprintf(`{"result": {"b64data": "%s", "skip_history": true}}`, customDataB64)))
	})
	defer httpTestCase.Teardown()

	cases := newPublishHandleTestCases(httpTestCase, grpcTestCase)
	for _, c := range cases {
		reply, err := c.invokeHandle()
		require.NoError(t, err, c.protocol)
		require.Equal(t, uint64(0), reply.Result.Offset, c.protocol)
		require.Empty(t, reply.Result.Epoch, c.protocol)
	}
}

func TestHandlePublishWithContextCancel(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	grpcTestCase := newPublishHandleGRPCTestCase(ctx, proxyGRPCTestServer{}, configtypes.ChannelOptions{
		PublishProxyEnabled: true,
		PublishProxyName:    "test",
	})
	defer grpcTestCase.Teardown()

	httpTestCase := newPublishHandleHTTPTestCase(ctx, "/publish", configtypes.ChannelOptions{
		PublishProxyEnabled: true,
		PublishProxyName:    "test",
	})
	httpTestCase.Mux.HandleFunc("/publish", func(w http.ResponseWriter, req *http.Request) {
		_, _ = w.Write([]byte(`{}`))
	})
	defer httpTestCase.Teardown()

	cases := newPublishHandleTestCases(httpTestCase, grpcTestCase)
	for _, c := range cases {
		reply, err := c.invokeHandle()
		require.ErrorIs(t, err, centrifuge.DisconnectConnectionClosed, c.protocol)
		require.Equal(t, centrifuge.PublishReply{}, reply, c.protocol)
	}
}

func TestHandlePublishWithoutProxyServerStart(t *testing.T) {
	grpcTestCase := newPublishHandleGRPCTestCase(context.Background(), proxyGRPCTestServer{}, configtypes.ChannelOptions{
		PublishProxyEnabled: true,
		PublishProxyName:    "test",
	})
	grpcTestCase.Teardown()

	httpTestCase := newPublishHandleHTTPTestCase(context.Background(), "/publish", configtypes.ChannelOptions{
		PublishProxyEnabled: true,
		PublishProxyName:    "test",
	})
	httpTestCase.Teardown()

	cases := newPublishHandleTestCases(httpTestCase, grpcTestCase)
	for _, c := range cases {
		reply, err := c.invokeHandle()
		require.Error(t, err, c.protocol)
		require.Equal(t, centrifuge.PublishReply{}, reply, c.protocol)
	}
}

func TestHandlePublishWithProxyServerCustomDisconnect(t *testing.T) {
	chOpts := configtypes.ChannelOptions{
		PublishProxyEnabled: true,
		PublishProxyName:    "test",
	}
	grpcTestCase := newPublishHandleGRPCTestCase(context.Background(), newProxyGRPCTestServer("custom disconnect", proxyGRPCTestServerOptions{}), chOpts)
	defer grpcTestCase.Teardown()

	httpTestCase := newPublishHandleHTTPTestCase(context.Background(), "/publish", chOpts)
	httpTestCase.Mux.HandleFunc("/publish", func(w http.ResponseWriter, req *http.Request) {
		_, _ = w.Write([]byte(`{"disconnect": {"code": 4000, "reconnect": false, "reason": "custom disconnect"}}`))
	})
	defer httpTestCase.Teardown()

	expectedErr := centrifuge.Disconnect{
		Code:   4000,
		Reason: "custom disconnect",
	}

	cases := newPublishHandleTestCases(httpTestCase, grpcTestCase)
	for _, c := range cases {
		reply, err := c.invokeHandle()
		require.NotNil(t, err, c.protocol)
		require.Equal(t, expectedErr.Error(), err.Error(), c.protocol)
		require.Equal(t, centrifuge.PublishReply{}, reply, c.protocol)
	}
}

func TestHandlePublishWithProxyServerCustomError(t *testing.T) {
	chOpts := configtypes.ChannelOptions{
		PublishProxyEnabled: true,
		PublishProxyName:    "test",
	}
	grpcTestCase := newPublishHandleGRPCTestCase(context.Background(), newProxyGRPCTestServer("custom error", proxyGRPCTestServerOptions{}), chOpts)
	defer grpcTestCase.Teardown()

	httpTestCase := newPublishHandleHTTPTestCase(context.Background(), "/publish", chOpts)
	httpTestCase.Mux.HandleFunc("/publish", func(w http.ResponseWriter, req *http.Request) {
		_, _ = w.Write([]byte(`{"error": {"code": 1000, "message": "custom error"}}`))
	})
	defer httpTestCase.Teardown()

	expectedErr := centrifuge.Error{
		Code:    1000,
		Message: "custom error",
	}

	cases := newPublishHandleTestCases(httpTestCase, grpcTestCase)
	for _, c := range cases {
		reply, err := c.invokeHandle()
		require.NotNil(t, err, c.protocol)
		require.Equal(t, expectedErr.Error(), err.Error(), c.protocol)
		require.Equal(t, centrifuge.PublishReply{}, reply, c.protocol)
	}
}

func TestHandlePublishWithInvalidCustomData(t *testing.T) {
	chOpts := configtypes.ChannelOptions{
		PublishProxyEnabled: true,
		PublishProxyName:    "test",
	}
	opts := proxyGRPCTestServerOptions{
		B64Data: "invalid data",
	}
	grpcTestCase := newPublishHandleGRPCTestCase(context.Background(), newProxyGRPCTestServer("result", opts), chOpts)
	defer grpcTestCase.Teardown()

	httpTestCase := newPublishHandleHTTPTestCase(context.Background(), "/publish", chOpts)
	httpTestCase.Mux.HandleFunc("/publish", func(w http.ResponseWriter, req *http.Request) {
		_, _ = w.Write([]byte(`{"result": {"b64data": "invalid data"}}`))
	})
	defer httpTestCase.Teardown()

	cases := newPublishHandleTestCases(httpTestCase, grpcTestCase)
	for _, c := range cases {
		reply, err := c.invokeHandle()
		require.ErrorIs(t, centrifuge.ErrorInternal, err, c.protocol)
		require.Equal(t, centrifuge.PublishReply{}, reply, c.protocol)
	}
}

// TestHandlePublishProxyResultDataValidated ensures data returned by a publish
// proxy is validated against the channel publication_data_format — the proxy can
// replace what the client sent, so the client-side check is not enough.
func TestHandlePublishProxyResultDataValidated(t *testing.T) {
	// Not a JSON object, while the channel requires one.
	badDataB64 := base64.StdEncoding.EncodeToString([]byte(`"just a string"`))
	chOpts := configtypes.ChannelOptions{
		PublishProxyEnabled:   true,
		PublishProxyName:      "test",
		PublicationDataFormat: configtypes.PublicationDataFormatJSONObject,
	}

	opts := proxyGRPCTestServerOptions{
		B64Data: badDataB64,
	}
	grpcTestCase := newPublishHandleGRPCTestCase(context.Background(), newProxyGRPCTestServer("result", opts), chOpts)
	defer grpcTestCase.Teardown()

	httpTestCase := newPublishHandleHTTPTestCase(context.Background(), "/publish", chOpts)
	httpTestCase.Mux.HandleFunc("/publish", func(w http.ResponseWriter, req *http.Request) {
		_, _ = w.Write([]byte(fmt.Sprintf(`{"result": {"b64data": "%s"}}`, badDataB64)))
	})
	defer httpTestCase.Teardown()

	cases := newPublishHandleTestCases(httpTestCase, grpcTestCase)
	for _, c := range cases {
		_, err := c.invokeHandle()
		require.Equal(t, centrifuge.ErrorBadRequest, err, c.protocol)
	}
}