chore: enforce max limit for webhook (#4975)

gateway/webhook/setup.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/rudderlabs/rudder-go-kit/logger"
 	"github.com/rudderlabs/rudder-go-kit/stats"
+
 	gwstats "github.com/rudderlabs/rudder-server/gateway/internal/stats"
 	gwtypes "github.com/rudderlabs/rudder-server/gateway/internal/types"
 	"github.com/rudderlabs/rudder-server/gateway/webhook/model"
@@ -56,6 +57,8 @@ func Setup(gwHandle Gateway, transformerFeaturesService transformer.FeaturesServ
 	maxTransformerProcess := config.GetIntVar(64, 1, "Gateway.webhook.maxTransformerProcess")
 	// Parse all query params from sources mentioned in this list
 	webhook.config.sourceListForParsingParams = config.GetStringSliceVar([]string{"Shopify", "adjust"}, "Gateway.webhook.sourceListForParsingParams")
+	// Maximum request size to gateway
+	webhook.config.maxReqSize = config.GetReloadableIntVar(4000, 1024, "Gateway.maxReqSizeInKB")
 
 	webhook.config.forwardGetRequestForSrcMap = lo.SliceToMap(
 		config.GetStringSliceVar([]string{"adjust"}, "Gateway.webhook.forwardGetRequestForSrcs"),

gateway/webhook/webhook.go

@@ -66,6 +66,7 @@ type HandleT struct {
 	backgroundCancel context.CancelFunc
 
 	config struct {
+		maxReqSize                 config.ValueLoader[int]
 		webhookBatchTimeout        config.ValueLoader[time.Duration]
 		maxWebhookBatchSize        config.ValueLoader[int]
 		sourceListForParsingParams []string
@@ -334,6 +335,13 @@ func (bt *batchWebhookTransformerT) batchTransformLoop() {
 				req.done <- transformerResponse{Err: response.GetStatus(response.InvalidJSON)}
 				continue
 			}
+			if len(body) > bt.webhook.config.maxReqSize.Load() {
+				req.done <- transformerResponse{
+					StatusCode: response.GetErrorStatusCode(response.RequestBodyTooLarge),
+					Err:        response.GetStatus(response.RequestBodyTooLarge),
+				}
+				continue
+			}
 
 			payload, err := sourceTransformAdapter.getTransformerEvent(req.authContext, body)
 			if err != nil {

gateway/webhook/webhook_test.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -15,13 +16,16 @@ import (
 
 	"go.uber.org/mock/gomock"
 
+	"github.com/rudderlabs/rudder-go-kit/bytesize"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/rudderlabs/rudder-go-kit/config"
 	"github.com/rudderlabs/rudder-go-kit/logger"
 	"github.com/rudderlabs/rudder-go-kit/stats"
 	"github.com/rudderlabs/rudder-go-kit/stats/memstats"
+
 	gwStats "github.com/rudderlabs/rudder-server/gateway/internal/stats"
 	gwtypes "github.com/rudderlabs/rudder-server/gateway/internal/types"
 	mockWebhook "github.com/rudderlabs/rudder-server/gateway/mocks"
@@ -61,11 +65,11 @@ type mockSourceTransformAdapter struct {
 	url string
 }
 
-func (v0 *mockSourceTransformAdapter) getTransformerEvent(authCtx *gwtypes.AuthRequestContext, body []byte) ([]byte, error) {
+func (v0 *mockSourceTransformAdapter) getTransformerEvent(_ *gwtypes.AuthRequestContext, body []byte) ([]byte, error) {
 	return body, nil
 }
 
-func (v0 *mockSourceTransformAdapter) getTransformerURL(sourceType string) (string, error) {
+func (v0 *mockSourceTransformAdapter) getTransformerURL(string) (string, error) {
 	return v0.url, nil
 }
 
@@ -77,13 +81,53 @@ func getMockSourceTransformAdapterFunc(url string) func(ctx context.Context) (so
 	}
 }
 
+func TestWebhookMaxRequestSize(t *testing.T) {
+	initWebhook()
+
+	ctrl := gomock.NewController(t)
+
+	mockGW := mockWebhook.NewMockGateway(ctrl)
+	mockGW.EXPECT().TrackRequestMetrics(gomock.Any()).Times(1)
+	mockGW.EXPECT().NewSourceStat(gomock.Any(), gomock.Any()).Return(&gwStats.SourceStat{}).Times(1)
+
+	mockTransformerFeaturesService := mock_features.NewMockFeaturesService(ctrl)
+
+	maxReqSizeInKB := 1
+
+	webhookHandler := Setup(mockGW, mockTransformerFeaturesService, stats.NOP, func(bt *batchWebhookTransformerT) {
+		bt.sourceTransformAdapter = func(ctx context.Context) (sourceTransformAdapter, error) {
+			return &mockSourceTransformAdapter{}, nil
+		}
+	})
+	webhookHandler.config.maxReqSize = config.SingleValueLoader(maxReqSizeInKB)
+	t.Cleanup(func() {
+		_ = webhookHandler.Shutdown()
+	})
+
+	webhookHandler.Register(sourceDefName)
+
+	payload := fmt.Sprintf(`{"hello":"world", "data": %q}`, strings.Repeat("a", 2*maxReqSizeInKB*int(bytesize.KB)))
+	require.Greater(t, len(payload), maxReqSizeInKB*int(bytesize.KB))
+
+	req := httptest.NewRequest(http.MethodPost, "/v1/webhook", bytes.NewBufferString(payload))
+	resp := httptest.NewRecorder()
+
+	reqCtx := context.WithValue(req.Context(), gwtypes.CtxParamCallType, "webhook")
+	reqCtx = context.WithValue(reqCtx, gwtypes.CtxParamAuthRequestContext, &gwtypes.AuthRequestContext{
+		SourceDefName: sourceDefName,
+	})
+
+	webhookHandler.RequestHandler(resp, req.WithContext(reqCtx))
+	require.Equal(t, http.StatusRequestEntityTooLarge, resp.Result().StatusCode)
+}
+
 func TestWebhookBlockTillFeaturesAreFetched(t *testing.T) {
 	initWebhook()
 	ctrl := gomock.NewController(t)
 	mockGW := mockWebhook.NewMockGateway(ctrl)
 	mockTransformerFeaturesService := mock_features.NewMockFeaturesService(ctrl)
 	mockTransformerFeaturesService.EXPECT().Wait().Return(make(chan struct{})).Times(1)
-	webhookHandler := Setup(mockGW, mockTransformerFeaturesService, stats.Default)
+	webhookHandler := Setup(mockGW, mockTransformerFeaturesService, stats.NOP)
 
 	mockGW.EXPECT().TrackRequestMetrics(gomock.Any()).Times(1)
 	mockGW.EXPECT().NewSourceStat(gomock.Any(), gomock.Any()).Return(&gwStats.SourceStat{}).Times(1)
@@ -112,7 +156,7 @@ func TestWebhookRequestHandlerWithTransformerBatchGeneralError(t *testing.T) {
 		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			http.Error(w, sampleError, http.StatusBadRequest)
 		}))
-	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.Default, func(bt *batchWebhookTransformerT) {
+	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.NOP, func(bt *batchWebhookTransformerT) {
 		bt.sourceTransformAdapter = getMockSourceTransformAdapterFunc(transformerServer.URL)
 	})
 
@@ -157,7 +201,7 @@ func TestWebhookRequestHandlerWithTransformerBatchPayloadLengthMismatchError(t *
 			respBody, _ := json.Marshal(responses)
 			_, _ = w.Write(respBody)
 		}))
-	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.Default, func(bt *batchWebhookTransformerT) {
+	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.NOP, func(bt *batchWebhookTransformerT) {
 		bt.sourceTransformAdapter = getMockSourceTransformAdapterFunc(transformerServer.URL)
 	})
 
@@ -200,7 +244,7 @@ func TestWebhookRequestHandlerWithTransformerRequestError(t *testing.T) {
 			respBody, _ := json.Marshal(responses)
 			_, _ = w.Write(respBody)
 		}))
-	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.Default, func(bt *batchWebhookTransformerT) {
+	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.NOP, func(bt *batchWebhookTransformerT) {
 		bt.sourceTransformAdapter = getMockSourceTransformAdapterFunc(transformerServer.URL)
 	})
 
@@ -243,7 +287,7 @@ func TestWebhookRequestHandlerWithOutputToSource(t *testing.T) {
 			respBody, _ := json.Marshal(responses)
 			_, _ = w.Write(respBody)
 		}))
-	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.Default, func(bt *batchWebhookTransformerT) {
+	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.NOP, func(bt *batchWebhookTransformerT) {
 		bt.sourceTransformAdapter = getMockSourceTransformAdapterFunc(transformerServer.URL)
 	})
 	mockGW.EXPECT().TrackRequestMetrics("").Times(1)
@@ -285,7 +329,7 @@ func TestWebhookRequestHandlerWithOutputToGateway(t *testing.T) {
 			respBody, _ := json.Marshal(responses)
 			_, _ = w.Write(respBody)
 		}))
-	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.Default, func(bt *batchWebhookTransformerT) {
+	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.NOP, func(bt *batchWebhookTransformerT) {
 		bt.sourceTransformAdapter = getMockSourceTransformAdapterFunc(transformerServer.URL)
 	})
 	mockGW.EXPECT().TrackRequestMetrics("").Times(1)
@@ -332,7 +376,7 @@ func TestWebhookRequestHandlerWithOutputToGatewayAndSource(t *testing.T) {
 			respBody, _ := json.Marshal(responses)
 			_, _ = w.Write(respBody)
 		}))
-	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.Default, func(bt *batchWebhookTransformerT) {
+	webhookHandler := Setup(mockGW, transformer.NewNoOpService(), stats.NOP, func(bt *batchWebhookTransformerT) {
 		bt.sourceTransformAdapter = getMockSourceTransformAdapterFunc(transformerServer.URL)
 	})
 	mockGW.EXPECT().TrackRequestMetrics("").Times(1)