Add wrapper for using middleware (#178)

server/server.go

@@ -31,6 +31,7 @@ type Server struct {
 	initializers      []InitializerFunc
 	initializeTimeout time.Duration
 	registrars        []func(mux *http.ServeMux) error
+	middlewares       []Middleware
 
 	// GRPCServer will be started whenever this is served
 	GRPCServer *grpc.Server
@@ -39,6 +40,9 @@ type Server struct {
 	HTTPServer *http.Server
 }
 
+//Middleware wrapper
+type Middleware func(handler http.Handler) http.Handler
+
 // Option is a functional option for creating a Server
 type Option func(*Server) error
 
@@ -67,6 +71,15 @@ func NewServer(opts ...Option) (*Server, error) {
 	}
 	s.HTTPServer.Handler = mux
 
+	// Revert user input middlewares
+	for i, j := 0, len(s.middlewares)-1; i < j; i, j = i+1, j-1 {
+		s.middlewares[i], s.middlewares[j] = s.middlewares[j], s.middlewares[i]
+	}
+
+	for _, m := range s.middlewares {
+		s.HTTPServer.Handler = m(s.HTTPServer.Handler)
+	}
+
 	return s, nil
 }
 
@@ -129,6 +142,14 @@ func WithGateway(options ...gateway.Option) Option {
 	}
 }
 
+// WithMiddlewaries add opportunity to add different middleware
+func WithMiddlewares(middleware ...Middleware) Option {
+	return func(s *Server) error {
+		s.middlewares = append(s.middlewares, middleware...)
+		return nil
+	}
+}
+
 // Serve invokes all initializers then serves on the given listeners.
 //
 // If a listener is left blank, then that particular part will not be served.

server/server_test.go

@@ -127,6 +127,50 @@ func TestWithHandler(t *testing.T) {
 	}
 }
 
+func TestWithMiddlewares(t *testing.T) {
+	s, err := NewServer(WithMiddlewares(testParamSetter, (&testHttpServer{t}).testParamVerify))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	httpL, err := servertest.NewLocalListener()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	go s.Serve(nil, httpL)
+	defer s.Stop()
+
+	resp, err := http.Get(fmt.Sprint("http://", httpL.Addr().String(), "/v1/hello?name=test"))
+	if resp.Header.Get("status") != "ok" {
+		t.Error("Expected another status")
+	}
+	if resp.Header.Get("param") != "status" {
+		t.Error("Expected another status")
+	}
+}
+func testParamSetter(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("status", "ok")
+		h.ServeHTTP(w, r)
+	})
+}
+
+func (t *testHttpServer) testParamVerify(h http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		status := w.Header().Get("status")
+		if status != "ok" {
+			t.T.Error("Invalid status")
+		}
+		w.Header().Set("param", "status")
+		h.ServeHTTP(w, r)
+	})
+}
+
+type testHttpServer struct {
+	*testing.T
+}
+
 func TestWithGateway(t *testing.T) {
 	grpcL, err := servertest.NewLocalListener()
 	if err != nil {