fix: Add Server header to all responses

keboola · Sep 3, 2024 · 58ff0a5 · 58ff0a5
1 parent 4da1cae
commit 58ff0a5
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 16 deletions.
diff --git a/internal/pkg/service/stream/source/httpsource/error.go b/internal/pkg/service/stream/source/httpsource/error.go
@@ -28,6 +28,8 @@ func newErrorHandler(cfg Config, logger log.Logger) func(c *fasthttp.RequestCtx,
 	// Error handler
 	errorWriter := newErrorWriter(logger)
 	return func(c *fasthttp.RequestCtx, err error) {
+		c.Response.Header.Set("Server", ServerHeader)
+
 		var smallBufferErr *fasthttp.ErrSmallBuffer
 
 		switch {

diff --git a/internal/pkg/service/stream/source/httpsource/httpsource.go b/internal/pkg/service/stream/source/httpsource/httpsource.go
@@ -26,6 +26,7 @@ import (
 )
 
 const (
+	ServerHeader            = "Keboola/Stream/HTTPSource"
 	gracefulShutdownTimeout = 30 * time.Second
 )
 
@@ -53,6 +54,10 @@ func Start(ctx context.Context, d dependencies, cfg Config) error {
 
 	// Static routes
 	router := routing.New()
+	router.Use(func(c *routing.Context) error {
+		c.Response.Header.Set("Server", ServerHeader)
+		return nil
+	})
 	router.NotFound(routing.MethodNotAllowedHandler, func(c *routing.Context) error {
 		errorHandler(c.RequestCtx, svcErrors.NewRouteNotFound(errors.New("not found, please send data using POST /stream/<projectID>/<sourceID>/<secret>")))
 		return nil
@@ -70,7 +75,6 @@ func Start(ctx context.Context, d dependencies, cfg Config) error {
 
 	// Route import requests to the dispatcher
 	router.Post("/stream/<projectID>/<sourceID>/<secret>", func(c *routing.Context) error {
-		c.Response.Header.Set("Server", "Keboola stream HTTP source")
 		// Get parameters
 		projectIDStr := c.Param("projectID")
 		projectIDInt, err := strconv.Atoi(projectIDStr)

diff --git a/internal/pkg/service/stream/source/httpsource/httpsource_test.go b/internal/pkg/service/stream/source/httpsource/httpsource_test.go
@@ -23,6 +23,7 @@ import (
 	"github.com/keboola/keboola-as-code/internal/pkg/service/stream/definition/key"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/stream/dependencies"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/stream/sink/pipeline"
+	"github.com/keboola/keboola-as-code/internal/pkg/service/stream/source/httpsource"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/stream/storage/test"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/stream/storage/test/dummy"
 	"github.com/keboola/keboola-as-code/internal/pkg/utils/errors"
@@ -166,13 +167,15 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			Method:             http.MethodGet,
 			Path:               "/health-check",
 			ExpectedStatusCode: http.StatusOK,
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
 			ExpectedBody:       "OK\n",
 		},
 		{
 			Name:               "not found",
 			Method:             http.MethodGet,
 			Path:               "/foo",
 			ExpectedStatusCode: http.StatusNotFound,
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
 			ExpectedLogs:       `{"level":"info","message":"not found, please send data using POST /stream/<projectID>/<sourceID>/<secret>"}`,
 			ExpectedBody: `
 {
@@ -189,6 +192,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			ExpectedHeaders: map[string]string{
 				"Allow":          "OPTIONS, POST",
 				"Content-Length": "0",
+				"Server":         httpsource.ServerHeader,
 			},
 		},
 		{
@@ -197,6 +201,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			Path:               "/stream/foo/my-source/my-secret",
 			Body:               strings.NewReader(strings.Repeat(".", ts.maxBodySize)),
 			ExpectedStatusCode: http.StatusBadRequest,
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
 			ExpectedBody: `
 {
   "statusCode": 400,
@@ -210,6 +215,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			Path:               "/stream/1111/my-source/my-secret",
 			Body:               strings.NewReader(strings.Repeat(".", ts.maxBodySize)),
 			ExpectedStatusCode: http.StatusNotFound,
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
 			ExpectedBody: `
 {
   "statusCode": 404,
@@ -223,6 +229,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			Path:               "/stream/123/my-source-1/" + ts.invalidSecret,
 			Body:               strings.NewReader(strings.Repeat(".", ts.maxBodySize)),
 			ExpectedStatusCode: http.StatusNotFound,
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
 			ExpectedBody: `
 {
   "statusCode": 404,
@@ -236,6 +243,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			Path:               "/stream/123/my-source-2/" + ts.validSecret,
 			Body:               strings.NewReader(strings.Repeat(".", ts.maxBodySize)),
 			ExpectedStatusCode: http.StatusNotFound,
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
 			ExpectedBody: `
 {
   "statusCode": 404,
@@ -257,6 +265,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			ExpectedStatusCode: http.StatusInternalServerError,
 			ExpectedHeaders: map[string]string{
 				"Content-Type": "application/json",
+				"Server":       httpsource.ServerHeader,
 			},
 			ExpectedLogs: `
 {"level":"error","message":"write record error: cannot open sink pipeline: some open error, next attempt after %s","component":"sink.router"}
@@ -318,6 +327,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			ExpectedStatusCode: http.StatusInternalServerError,
 			ExpectedHeaders: map[string]string{
 				"Content-Type": "application/json",
+				"Server":       httpsource.ServerHeader,
 			},
 			ExpectedLogs: `
 {"level":"error","message":"write record error: some write error","component":"sink.router"}
@@ -378,7 +388,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			ExpectedStatusCode: http.StatusAccepted,
 			ExpectedHeaders: map[string]string{
 				"Content-Type": "text/plain",
-				"Server":       "Keboola stream HTTP source",
+				"Server":       httpsource.ServerHeader,
 			},
 			ExpectedBody: "OK",
 		},
@@ -396,7 +406,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			ExpectedStatusCode: http.StatusOK,
 			ExpectedHeaders: map[string]string{
 				"Content-Type": "text/plain",
-				"Server":       "Keboola stream HTTP source",
+				"Server":       httpsource.ServerHeader,
 			},
 			ExpectedBody: "OK",
 		},
@@ -413,9 +423,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			Query:              "verbose=true",
 			Body:               strings.NewReader(strings.Repeat(".", ts.maxBodySize)),
 			ExpectedStatusCode: http.StatusAccepted,
-			ExpectedHeaders: map[string]string{
-				"Server": "Keboola stream HTTP source",
-			},
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
 			ExpectedBody: `
 {
   "statusCode": 202,
@@ -465,9 +473,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			Query:              "verbose=true",
 			Body:               strings.NewReader(strings.Repeat(".", ts.maxBodySize)),
 			ExpectedStatusCode: http.StatusOK,
-			ExpectedHeaders: map[string]string{
-				"Server": "Keboola stream HTTP source",
-			},
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
 			ExpectedBody: `
 {
   "statusCode": 200,
@@ -510,8 +516,8 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			Path:               "/stream/123/my-source-1/" + ts.validSecret,
 			Headers:            map[string]string{"foo": strings.Repeat(".", ts.maxHeaderSize+1)},
 			ExpectedStatusCode: http.StatusRequestEntityTooLarge,
-			// No expected Server headers, fasthttp internal return
-			ExpectedLogs: `{"level":"info","message":"request header size is over the maximum \"2000B\"","error.type":"%s/errors.HeaderTooLargeError"}`,
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
+			ExpectedLogs:       `{"level":"info","message":"request header size is over the maximum \"2000B\"","error.type":"%s/errors.HeaderTooLargeError"}`,
 			ExpectedBody: `
 {
   "statusCode": 413,
@@ -525,8 +531,8 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			Path:               "/stream/123/my-source/" + ts.validSecret,
 			Body:               strings.NewReader(strings.Repeat(".", ts.maxBodySize+1)),
 			ExpectedStatusCode: http.StatusRequestEntityTooLarge,
-			// No expected Server headers, fasthttp internal return
-			ExpectedLogs: `{"level":"info","message":"request body size is over the maximum \"8000B\"","error.type":"%s/errors.BodyTooLargeError"}`,
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
+			ExpectedLogs:       `{"level":"info","message":"request body size is over the maximum \"8000B\"","error.type":"%s/errors.BodyTooLargeError"}`,
 			ExpectedBody: `
 {
   "statusCode": 413,
@@ -554,9 +560,7 @@ func testCases(t *testing.T, ts *testState) []TestCase {
 			Query:              "verbose=true",
 			Body:               strings.NewReader("foo"),
 			ExpectedStatusCode: http.StatusOK,
-			ExpectedHeaders: map[string]string{
-				"Server": "Keboola stream HTTP source",
-			},
+			ExpectedHeaders:    map[string]string{"Server": httpsource.ServerHeader},
 			ExpectedBody: `
 {
   "statusCode": 200,