Move addLoggerToContextFn to a middleware

http.go

@@ -71,14 +71,10 @@ func (t *HTTPTransport) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	t.httpmiddleware(t.mux).ServeHTTP(w, r)
 }
 
-// RegisterEndpoints registers all configured endpoints, wraps them with the m middleware and
-// configures the logger.
-func (t *HTTPTransport) RegisterEndpoints(m endpoint.Middleware, fn AddLoggerToContextFn) error {
+// RegisterEndpoints registers all configured endpoints, wraps them with the m middleware.
+func (t *HTTPTransport) RegisterEndpoints(m endpoint.Middleware) error {
 	opts := []kithttp.ServerOption{
 		kithttp.ServerBefore(kithttp.PopulateRequestContext),
-		kithttp.ServerBefore(func(ctx context.Context, _ *http.Request) context.Context {
-			return fn(ctx)
-		}),
 	}
 	opts = append(opts, t.opts...)
 

log.go

@@ -3,6 +3,7 @@ package kitty
 import (
 	"context"
 
+	"github.com/go-kit/kit/endpoint"
 	"github.com/go-kit/kit/log"
 	kithttp "github.com/go-kit/kit/transport/http"
 )
@@ -49,7 +50,14 @@ func (s *Server) addLoggerToContext(ctx context.Context) context.Context {
 	return context.WithValue(ctx, logKey, l)
 }
 
-type AddLoggerToContextFn func(ctx context.Context) context.Context
+func (s *Server) addLoggerToContextMiddleware(m endpoint.Middleware) endpoint.Middleware {
+	return func(e endpoint.Endpoint) endpoint.Endpoint {
+		e = m(e)
+		return func(ctx context.Context, request interface{}) (response interface{}, err error) {
+			return e(s.addLoggerToContext(ctx), request)
+		}
+	}
+}
 
 // Logger will return the logger that has been injected into the context by the kitty
 // server. This function can only be called from an endpoint.

server.go

@@ -45,8 +45,9 @@ func NewServer(t ...Transport) *Server {
 // Run starts the server.
 func (s *Server) Run(ctx context.Context) error {
 	ctx = s.addLoggerToContext(ctx)
+	m := s.addLoggerToContextMiddleware(s.middleware)
 	for _, t := range s.transports {
-		if err := t.RegisterEndpoints(s.middleware, s.addLoggerToContext); err != nil {
+		if err := t.RegisterEndpoints(m); err != nil {
 			return err
 		}
 	}

transport.go

@@ -8,8 +8,8 @@ import (
 
 // Transport is the interface all transports must implement.
 type Transport interface {
-	// RegisterEndpoints registers all endpoints, and injects the AddLoggerToContextFn call to configure the logger.
-	RegisterEndpoints(m endpoint.Middleware, fn AddLoggerToContextFn) error
+	// RegisterEndpoints registers all endpoints. Endpoints needs to be wrapped with the specified middleware.
+	RegisterEndpoints(m endpoint.Middleware) error
 	// Start starts the transport.
 	Start(ctx context.Context) error
 	// Shutdown shutdowns the transport.