From c647984391f7d85a5d370767aed12dcba92653b1 Mon Sep 17 00:00:00 2001 From: Dj Gilcrease Date: Wed, 15 Feb 2017 18:35:10 -0800 Subject: [PATCH] Add the ability to specify ServerBefore & ServerAfter multiple times in the ServerOption section This makes it so you can set some global ServerBefore or After options without having to make sure you specify them individually on every handler --- transport/grpc/server.go | 4 +- transport/http/server.go | 4 +- transport/http/server_test.go | 93 +++++++++++++++++++++++++++++++++ transport/httprp/server.go | 2 +- transport/httprp/server_test.go | 42 +++++++++++++++ 5 files changed, 140 insertions(+), 5 deletions(-) diff --git a/transport/grpc/server.go b/transport/grpc/server.go index 84e27cfc9..742c1a086 100644 --- a/transport/grpc/server.go +++ b/transport/grpc/server.go @@ -59,13 +59,13 @@ type ServerOption func(*Server) // ServerBefore functions are executed on the HTTP request object before the // request is decoded. func ServerBefore(before ...RequestFunc) ServerOption { - return func(s *Server) { s.before = before } + return func(s *Server) { s.before = append(s.before, before...) } } // ServerAfter functions are executed on the HTTP response writer after the // endpoint is invoked, but before anything is written to the client. func ServerAfter(after ...ResponseFunc) ServerOption { - return func(s *Server) { s.after = after } + return func(s *Server) { s.after = append(s.after, after...) } } // ServerErrorLogger is used to log non-terminal errors. By default, no errors diff --git a/transport/http/server.go b/transport/http/server.go index ab2b22e56..de109beff 100644 --- a/transport/http/server.go +++ b/transport/http/server.go @@ -51,13 +51,13 @@ type ServerOption func(*Server) // ServerBefore functions are executed on the HTTP request object before the // request is decoded. func ServerBefore(before ...RequestFunc) ServerOption { - return func(s *Server) { s.before = before } + return func(s *Server) { s.before = append(s.before, before...) } } // ServerAfter functions are executed on the HTTP response writer after the // endpoint is invoked, but before anything is written to the client. func ServerAfter(after ...ServerResponseFunc) ServerOption { - return func(s *Server) { s.after = after } + return func(s *Server) { s.after = append(s.after, after...) } } // ServerErrorEncoder is used to encode errors to the http.ResponseWriter diff --git a/transport/http/server_test.go b/transport/http/server_test.go index 654fbecba..51e311ee8 100644 --- a/transport/http/server_test.go +++ b/transport/http/server_test.go @@ -93,6 +93,99 @@ func TestServerHappyPath(t *testing.T) { } } + +func TestMultipleServerBefore(t *testing.T) { + var ( + headerKey = "X-Henlo-Lizer" + headerVal = "Helllo you stinky lizard" + statusCode = http.StatusTeapot + responseBody = "go eat a fly ugly\n" + done = make(chan struct{}) + ) + handler := httptransport.NewServer( + context.Background(), + endpoint.Nop, + func(context.Context, *http.Request) (interface{}, error) { + return struct{}{}, nil + }, + func(_ context.Context, w http.ResponseWriter, _ interface{}) error { + w.Header().Set(headerKey, headerVal) + w.WriteHeader(statusCode) + w.Write([]byte(responseBody)) + return nil + }, + httptransport.ServerBefore(func(ctx context.Context, r *http.Request) context.Context { + ctx = context.WithValue(ctx, "one", 1) + + return ctx + }), + httptransport.ServerBefore(func(ctx context.Context, r *http.Request) context.Context { + if _, ok := ctx.Value("one").(int); !ok { + t.Error("Value was not set properly when multiple ServerBefores are used") + } + + close(done) + return ctx + }), + ) + + server := httptest.NewServer(handler) + defer server.Close() + go http.Get(server.URL) + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("timeout waiting for finalizer") + } +} + +func TestMultipleServerAfter(t *testing.T) { + var ( + headerKey = "X-Henlo-Lizer" + headerVal = "Helllo you stinky lizard" + statusCode = http.StatusTeapot + responseBody = "go eat a fly ugly\n" + done = make(chan struct{}) + ) + handler := httptransport.NewServer( + context.Background(), + endpoint.Nop, + func(context.Context, *http.Request) (interface{}, error) { + return struct{}{}, nil + }, + func(_ context.Context, w http.ResponseWriter, _ interface{}) error { + w.Header().Set(headerKey, headerVal) + w.WriteHeader(statusCode) + w.Write([]byte(responseBody)) + return nil + }, + httptransport.ServerAfter(func(ctx context.Context, w http.ResponseWriter) context.Context { + ctx = context.WithValue(ctx, "one", 1) + + return ctx + }), + httptransport.ServerAfter(func(ctx context.Context, w http.ResponseWriter) context.Context { + if _, ok := ctx.Value("one").(int); !ok { + t.Error("Value was not set properly when multiple ServerAfters are used") + } + + close(done) + return ctx + }), + ) + + server := httptest.NewServer(handler) + defer server.Close() + go http.Get(server.URL) + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("timeout waiting for finalizer") + } +} + func TestServerFinalizer(t *testing.T) { var ( headerKey = "X-Henlo-Lizer" diff --git a/transport/httprp/server.go b/transport/httprp/server.go index 25776ed8d..5c3ad1baf 100644 --- a/transport/httprp/server.go +++ b/transport/httprp/server.go @@ -45,7 +45,7 @@ type ServerOption func(*Server) // ServerBefore functions are executed on the HTTP request object before the // request is decoded. func ServerBefore(before ...RequestFunc) ServerOption { - return func(s *Server) { s.before = before } + return func(s *Server) { s.before = append(s.before, before...) } } // ServeHTTP implements http.Handler. diff --git a/transport/httprp/server_test.go b/transport/httprp/server_test.go index 45fd429d1..47d6f77f0 100644 --- a/transport/httprp/server_test.go +++ b/transport/httprp/server_test.go @@ -119,3 +119,45 @@ func TestServerOriginServerUnreachable(t *testing.T) { t.Errorf("want %d or %d, have %d", http.StatusBadGateway, http.StatusInternalServerError, resp.StatusCode) } } + +func TestMultipleServerBefore(t *testing.T) { + const ( + headerKey = "X-TEST-HEADER" + headerVal = "go-kit-proxy" + ) + + originServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if want, have := headerVal, r.Header.Get(headerKey); want != have { + t.Errorf("want %q, have %q", want, have) + } + + w.WriteHeader(http.StatusOK) + w.Write([]byte("hey")) + })) + defer originServer.Close() + originURL, _ := url.Parse(originServer.URL) + + handler := httptransport.NewServer( + context.Background(), + originURL, + httptransport.ServerBefore(func(ctx context.Context, r *http.Request) context.Context { + r.Header.Add(headerKey, headerVal) + return ctx + }), + httptransport.ServerBefore(func(ctx context.Context, r *http.Request) context.Context { + return ctx + }), + ) + proxyServer := httptest.NewServer(handler) + defer proxyServer.Close() + + resp, _ := http.Get(proxyServer.URL) + if want, have := http.StatusOK, resp.StatusCode; want != have { + t.Errorf("want %d, have %d", want, have) + } + + responseBody, _ := ioutil.ReadAll(resp.Body) + if want, have := "hey", string(responseBody); want != have { + t.Errorf("want %q, have %q", want, have) + } +}