chore: refactor middleware setup in server

Makefile

@@ -60,6 +60,7 @@ mockgen: install-mockgen
 	cd core; mockgen -source=pkg/sync/grpc/grpc_sync.go -destination=pkg/sync/grpc/mock/grpc.go -package=grpcmock
 	cd core; mockgen -source=pkg/sync/grpc/credentials/builder.go -destination=pkg/sync/grpc/credentials/mock/builder.go -package=credendialsmock
 	cd core; mockgen -source=pkg/eval/ievaluator.go -destination=pkg/eval/mock/ievaluator.go -package=evalmock
+	cd core; mockgen -source=pkg/service/middleware/interface.go -destination=pkg/service/middleware/mock/interface.go -package=middlewaremock
 generate-docs:
 	cd flagd; go run ./cmd/doc/main.go
 

core/pkg/service/flag-evaluation/connect_service.go

@@ -10,17 +10,18 @@ import (
 	"sync"
 	"time"
 
+	"github.com/open-feature/flagd/core/pkg/service/middleware"
+
 	schemaConnectV1 "buf.build/gen/go/open-feature/flagd/bufbuild/connect-go/schema/v1/schemav1connect"
 	"github.com/open-feature/flagd/core/pkg/eval"
 	"github.com/open-feature/flagd/core/pkg/logger"
 	"github.com/open-feature/flagd/core/pkg/otel"
 	"github.com/open-feature/flagd/core/pkg/service"
-	"github.com/open-feature/flagd/core/pkg/service/middleware"
+	corsmw "github.com/open-feature/flagd/core/pkg/service/middleware/cors"
+	h2cmw "github.com/open-feature/flagd/core/pkg/service/middleware/h2c"
+	metricsmw "github.com/open-feature/flagd/core/pkg/service/middleware/metrics"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
-	"github.com/rs/cors"
 	"go.uber.org/zap"
-	"golang.org/x/net/http2"
-	"golang.org/x/net/http2/h2c"
 )
 
 const ErrorPrefix = "FlagdError:"
@@ -65,6 +66,8 @@ func (s *ConnectService) Serve(ctx context.Context, eval eval.IEvaluator, svcCon
 		close(errChan)
 	}()
 
+	go s.startMetricsServer(svcConf)
+
 	select {
 	case err := <-errChan:
 		return err
@@ -94,30 +97,37 @@ func (s *ConnectService) setupServer(svcConf service.Configuration) (net.Listene
 	path, handler := schemaConnectV1.NewServiceHandler(fes)
 	mux.Handle(path, handler)
 
-	mdlw := middleware.NewHTTPMetric(middleware.Config{
-		Service:        "openfeature/flagd",
+	s.server = http.Server{
+		ReadHeaderTimeout: time.Second,
+		Handler:           handler,
+	}
+
+	// Add middlewares
+
+	metricsMiddleware := metricsmw.NewHTTPMetric(metricsmw.Config{
+		Service:        svcConf.ServiceName,
 		MetricRecorder: s.Metrics,
 		Logger:         s.Logger,
+		HandlerID:      "",
 	})
-	h := middleware.Handler("", mdlw, mux)
 
-	go bindMetrics(s, svcConf)
+	s.AddMiddleware(metricsMiddleware)
 
-	if svcConf.CertPath != "" && svcConf.KeyPath != "" {
-		handler = s.newCORS(svcConf).Handler(h)
-	} else {
-		handler = h2c.NewHandler(
-			s.newCORS(svcConf).Handler(h),
-			&http2.Server{},
-		)
-	}
-	s.server = http.Server{
-		ReadHeaderTimeout: time.Second,
-		Handler:           handler,
+	corsMiddleware := corsmw.New(svcConf.CORS)
+	s.AddMiddleware(corsMiddleware)
+
+	if svcConf.CertPath == "" || svcConf.KeyPath == "" {
+		h2cMiddleware := h2cmw.New()
+		s.AddMiddleware(h2cMiddleware)
 	}
+
 	return lis, nil
 }
 
+func (s *ConnectService) AddMiddleware(mw middleware.IMiddleware) {
+	s.server.Handler = mw.Handler(s.server.Handler)
+}
+
 func (s *ConnectService) Notify(n service.Notification) {
 	s.eventingConfiguration.mu.RLock()
 	defer s.eventingConfiguration.mu.RUnlock()
@@ -126,36 +136,7 @@ func (s *ConnectService) Notify(n service.Notification) {
 	}
 }
 
-func (s *ConnectService) newCORS(svcConf service.Configuration) *cors.Cors {
-	return cors.New(cors.Options{
-		AllowedMethods: []string{
-			http.MethodHead,
-			http.MethodGet,
-			http.MethodPost,
-			http.MethodPut,
-			http.MethodPatch,
-			http.MethodDelete,
-		},
-		AllowedOrigins: svcConf.CORS,
-		AllowedHeaders: []string{"*"},
-		ExposedHeaders: []string{
-			// Content-Type is in the default safelist.
-			"Accept",
-			"Accept-Encoding",
-			"Accept-Post",
-			"Connect-Accept-Encoding",
-			"Connect-Content-Encoding",
-			"Content-Encoding",
-			"Grpc-Accept-Encoding",
-			"Grpc-Encoding",
-			"Grpc-Message",
-			"Grpc-Status",
-			"Grpc-Status-Details-Bin",
-		},
-	})
-}
-
-func bindMetrics(s *ConnectService, svcConf service.Configuration) {
+func (s *ConnectService) startMetricsServer(svcConf service.Configuration) {
 	s.Logger.Info(fmt.Sprintf("metrics and probes listening at %d", svcConf.MetricsPort))
 	server := &http.Server{
 		Addr:              fmt.Sprintf(":%d", svcConf.MetricsPort),

core/pkg/service/flag-evaluation/connect_service_test.go

@@ -4,10 +4,13 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net/http"
 	"os"
 	"testing"
 	"time"
 
+	middlewaremock "github.com/open-feature/flagd/core/pkg/service/middleware/mock"
+
 	schemaGrpcV1 "buf.build/gen/go/open-feature/flagd/grpc/go/schema/v1/schemav1grpc"
 	schemaV1 "buf.build/gen/go/open-feature/flagd/protocolbuffers/go/schema/v1"
 	"github.com/golang/mock/gomock"
@@ -117,3 +120,53 @@ func TestConnectService_UnixConnection(t *testing.T) {
 		})
 	}
 }
+
+func TestAddMiddleware(t *testing.T) {
+	const port = 12345
+	ctrl := gomock.NewController(t)
+
+	mwMock := middlewaremock.NewMockIMiddleware(ctrl)
+
+	mwMock.EXPECT().Handler(gomock.Any()).Return(
+		http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
+			writer.WriteHeader(http.StatusOK)
+		}))
+
+	exp := metric.NewManualReader()
+	metricRecorder := otel.NewOTelRecorder(exp, "my-exporter")
+
+	svc := ConnectService{
+		Logger:  logger.NewLogger(nil, false),
+		Metrics: metricRecorder,
+	}
+
+	serveConf := iservice.Configuration{
+		ReadinessProbe: func() bool {
+			return true
+		},
+		Port: port,
+	}
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	go func() {
+		err := svc.Serve(ctx, nil, serveConf)
+		fmt.Println(err)
+	}()
+
+	require.Eventually(t, func() bool {
+		resp, err := http.Get(fmt.Sprintf("http://localhost:%d/schema.v1.Service/ResolveAll", port))
+		// with the default http handler we should get a method not allowed (405) when attempting a GET request
+		return err == nil && resp.StatusCode == http.StatusMethodNotAllowed
+	}, 3*time.Second, 100*time.Millisecond)
+
+	svc.AddMiddleware(mwMock)
+
+	// with the injected middleware, the GET method should work
+	resp, err := http.Get(fmt.Sprintf("http://localhost:%d/schema.v1.Service/ResolveAll", port))
+
+	require.Nil(t, err)
+	// verify that the status we return in the mocked middleware
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+}

core/pkg/service/middleware/cors/cors.go

@@ -0,0 +1,46 @@
+package cors
+
+import (
+	"net/http"
+
+	"github.com/rs/cors"
+)
+
+type Middleware struct {
+	cors *cors.Cors
+}
+
+func New(allowedOrigins []string) *Middleware {
+	return &Middleware{
+		cors: cors.New(cors.Options{
+			AllowedMethods: []string{
+				http.MethodHead,
+				http.MethodGet,
+				http.MethodPost,
+				http.MethodPut,
+				http.MethodPatch,
+				http.MethodDelete,
+			},
+			AllowedOrigins: allowedOrigins,
+			AllowedHeaders: []string{"*"},
+			ExposedHeaders: []string{
+				// Content-Type is in the default safelist.
+				"Accept",
+				"Accept-Encoding",
+				"Accept-Post",
+				"Connect-Accept-Encoding",
+				"Connect-Content-Encoding",
+				"Content-Encoding",
+				"Grpc-Accept-Encoding",
+				"Grpc-Encoding",
+				"Grpc-Message",
+				"Grpc-Status",
+				"Grpc-Status-Details-Bin",
+			},
+		}),
+	}
+}
+
+func (c Middleware) Handler(handler http.Handler) http.Handler {
+	return c.cors.Handler(handler)
+}

core/pkg/service/middleware/cors/cors_test.go

@@ -0,0 +1,44 @@
+package cors
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	middlewaremock "github.com/open-feature/flagd/core/pkg/service/middleware/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMiddleware(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mockMw := middlewaremock.NewMockIMiddleware(ctrl)
+
+	handlerFunc := http.HandlerFunc(
+		func(writer http.ResponseWriter, request *http.Request) {
+			writer.WriteHeader(http.StatusOK)
+		},
+	)
+
+	mockMw.EXPECT().Handler(gomock.Any()).Return(handlerFunc)
+
+	ts := httptest.NewServer(handlerFunc)
+
+	defer ts.Close()
+
+	mw := New([]string{"*"})
+	require.NotNil(t, mw)
+
+	// wrap the cors middleware around the mock to make sure the wrapped handler is called by the cors middleware
+	ts.Config.Handler = mw.Handler(mockMw.Handler(handlerFunc))
+
+	req, err := http.NewRequest(http.MethodGet, ts.URL, nil)
+
+	require.Nil(t, err)
+
+	client := http.DefaultClient
+	resp, err := client.Do(req)
+
+	require.Nil(t, err)
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+}

core/pkg/service/middleware/h2c/h2c.go

@@ -0,0 +1,18 @@
+package h2c
+
+import (
+	"net/http"
+
+	"golang.org/x/net/http2"
+	"golang.org/x/net/http2/h2c"
+)
+
+type Middleware struct{}
+
+func New() *Middleware {
+	return &Middleware{}
+}
+
+func (m Middleware) Handler(handler http.Handler) http.Handler {
+	return h2c.NewHandler(handler, &http2.Server{})
+}

core/pkg/service/middleware/h2c/h2c_test.go

@@ -0,0 +1,39 @@
+package h2c
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	middlewaremock "github.com/open-feature/flagd/core/pkg/service/middleware/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMiddleware(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	mockMw := middlewaremock.NewMockIMiddleware(ctrl)
+
+	handlerFunc := http.HandlerFunc(
+		func(writer http.ResponseWriter, request *http.Request) {
+			writer.WriteHeader(http.StatusOK)
+		},
+	)
+
+	mockMw.EXPECT().Handler(gomock.Any()).Return(handlerFunc)
+
+	ts := httptest.NewServer(handlerFunc)
+
+	defer ts.Close()
+
+	mw := New()
+	require.NotNil(t, mw)
+
+	// wrap the h2c middleware around the mock to make sure the wrapped handler is called by the h2c middleware
+	ts.Config.Handler = mw.Handler(mockMw.Handler(handlerFunc))
+
+	resp, err := http.Get(ts.URL)
+
+	require.Nil(t, err)
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+}

core/pkg/service/middleware/interface.go

@@ -0,0 +1,9 @@
+package middleware
+
+import (
+	"net/http"
+)
+
+type IMiddleware interface {
+	Handler(handler http.Handler) http.Handler
+}

core/pkg/service/middleware/http_metrics.go → core/pkg/service/middleware/metrics/http_metrics.go

@@ -1,4 +1,4 @@
-package middleware
+package metrics
 
 import (
 	"bufio"
@@ -21,6 +21,7 @@ type Config struct {
 	Service            string
 	GroupedStatus      bool
 	DisableMeasureSize bool
+	HandlerID          string
 }
 
 type Middleware struct {
@@ -90,7 +91,7 @@ func (m Middleware) Measure(ctx context.Context, handlerID string, reporter Repo
 }
 
 // Handler returns an measuring standard http.Handler.
-func Handler(handlerID string, m Middleware, h http.Handler) http.Handler {
+func (m Middleware) Handler(h http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		wi := &responseWriterInterceptor{
 			statusCode:     http.StatusOK,
@@ -100,7 +101,7 @@ func Handler(handlerID string, m Middleware, h http.Handler) http.Handler {
 			w: wi,
 			r: r,
 		}
-		m.Measure(r.Context(), handlerID, reporter, func() {
+		m.Measure(r.Context(), m.cfg.HandlerID, reporter, func() {
 			h.ServeHTTP(wi, r)
 		})
 	})