new: support for custom response writer

context.go

@@ -21,23 +21,24 @@ import (
 )
 
 type bcontext struct {
-	claims        []string
-	claimsMap     map[string]string
-	count         int
-	ctx           context.Context
-	events        elemental.Events
-	eventsLock    *sync.Mutex
-	id            string
-	inputData     interface{}
-	messages      []string
-	messagesLock  *sync.Mutex
-	metadata      map[interface{}]interface{}
-	outputData    interface{}
-	redirect      string
-	request       *elemental.Request
-	statusCode    int
-	next          string
-	outputCookies []*http.Cookie
+	claims         []string
+	claimsMap      map[string]string
+	count          int
+	ctx            context.Context
+	events         elemental.Events
+	eventsLock     *sync.Mutex
+	id             string
+	inputData      interface{}
+	messages       []string
+	messagesLock   *sync.Mutex
+	metadata       map[interface{}]interface{}
+	next           string
+	outputCookies  []*http.Cookie
+	outputData     interface{}
+	redirect       string
+	request        *elemental.Request
+	responseWriter ResponseWriter
+	statusCode     int
 }
 
 // NewContext creates a new *Context.
@@ -95,9 +96,23 @@ func (c *bcontext) OutputData() interface{} {
 }
 
 func (c *bcontext) SetOutputData(data interface{}) {
+
+	if c.responseWriter != nil {
+		panic("you cannot use SetOutputData after using SetResponseWriter")
+	}
+
 	c.outputData = data
 }
 
+func (c *bcontext) SetResponseWriter(writer ResponseWriter) {
+
+	if c.outputData != nil {
+		panic("you cannot use SetResponseWriter after using SetOutputData")
+	}
+
+	c.responseWriter = writer
+}
+
 func (c *bcontext) StatusCode() int {
 	return c.statusCode
 }
@@ -193,6 +208,7 @@ func (c *bcontext) Duplicate() Context {
 	c2.messages = append(c2.messages, c.messages...)
 	c2.next = c.next
 	c2.outputCookies = append(c2.outputCookies, c.outputCookies...)
+	c2.responseWriter = c.responseWriter
 
 	for k, v := range c.claimsMap {
 		c2.claimsMap[k] = v

context_test.go

@@ -126,6 +126,7 @@ func TestContext_Duplicate(t *testing.T) {
 		}
 
 		cookies := []*http.Cookie{&http.Cookie{}, &http.Cookie{}}
+		rwriter := func(http.ResponseWriter) int { return 0 }
 
 		ctx := newContext(context.Background(), req)
 		ctx.SetCount(10)
@@ -139,6 +140,7 @@ func TestContext_Duplicate(t *testing.T) {
 		ctx.SetClaims([]string{"ouais=yes"})
 		ctx.SetNext("next")
 		ctx.AddOutputCookies(cookies[0], cookies[1])
+		ctx.SetResponseWriter(rwriter)
 
 		Convey("When I call the Duplicate method", func() {
 
@@ -159,6 +161,7 @@ func TestContext_Duplicate(t *testing.T) {
 				So(ctx.messages, ShouldResemble, ctx2.(*bcontext).messages)
 				So(ctx.outputCookies, ShouldResemble, ctx2.(*bcontext).outputCookies)
 				So(ctx.outputCookies, ShouldResemble, cookies)
+				So(ctx.responseWriter, ShouldEqual, rwriter)
 			})
 		})
 	})
@@ -218,3 +221,30 @@ func TestContext_GetClaims(t *testing.T) {
 		})
 	})
 }
+
+func TestOutputDataVSResponseWriter(t *testing.T) {
+
+	Convey("Given I have a bcontext", t, func() {
+
+		rwriter := func(http.ResponseWriter) int { return 0 }
+		ctx := newContext(context.Background(), elemental.NewRequest())
+
+		Convey("When I call SetResponseWriter after SetOutputData", func() {
+
+			ctx.SetOutputData("hello")
+
+			Convey("Then it should panic", func() {
+				So(func() { ctx.SetResponseWriter(rwriter) }, ShouldPanicWith, "you cannot use SetResponseWriter after using SetOutputData")
+			})
+		})
+
+		Convey("When I call SetOutputData after SetResponseWriter", func() {
+
+			ctx.SetResponseWriter(rwriter)
+
+			Convey("Then it should panic", func() {
+				So(func() { ctx.SetOutputData("hello") }, ShouldPanicWith, "you cannot use SetOutputData after using SetResponseWriter")
+			})
+		})
+	})
+}

gateway/gateway.go

@@ -368,6 +368,15 @@ func (s *gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		goto HANDLE_INTERCEPTION
 	}
 
+	// If we reach here, we check for suffix match
+	if interceptAction, upstream, err = s.checkInterceptor(
+		s.gatewayConfig.suffixInterceptors,
+		func(path string, key string) bool { return strings.HasSuffix(path, key) },
+		w, r, path,
+	); interceptAction != 0 {
+		goto HANDLE_INTERCEPTION
+	}
+
 HANDLE_INTERCEPTION:
 	if err != nil {
 		writeError(w, r, makeError(http.StatusInternalServerError, "Internal Server Error", fmt.Sprintf("unable to run interceptor: %s", err)))
@@ -414,7 +423,22 @@ HANDLE_INTERCEPTION:
 			mm.UnregisterWSConnection()
 		}
 
+	case InterceptorActionForwardDirect:
+
+		var finish bahamut.FinishMeasurementFunc
+
+		if mm := s.gatewayConfig.metricsManager; mm != nil {
+			finish = mm.MeasureRequest(r.Method, path)
+		}
+
+		s.forwarder.ServeHTTP(w, r)
+
+		if finish != nil {
+			finish(0, nil)
+		}
+
 	default:
+
 		var finish bahamut.FinishMeasurementFunc
 
 		if mm := s.gatewayConfig.metricsManager; mm != nil {

gateway/gateway_test.go

@@ -284,9 +284,10 @@ func TestGateway(t *testing.T) {
 			gw, err := New(
 				"127.0.0.1:7765",
 				u,
+				OptionMetricsManager(&fakeMetricManager{}),
 				OptionUpstreamTLSConfig(&tls.Config{InsecureSkipVerify: true}),
 				OptionRegisterExactInterceptor("/ups1", func(w http.ResponseWriter, req *http.Request, ew ErrorWriter) (InterceptorAction, string, error) {
-					return InterceptorActionForward, strings.Replace(u.ups2.URL, "https://", "", 1), nil
+					return InterceptorActionForwardWS, strings.Replace(u.ups2.URL, "https://", "", 1), nil
 				}),
 			)
 			defer gw.Stop()
@@ -312,7 +313,75 @@ func TestGateway(t *testing.T) {
 			})
 		})
 
-		Convey("When I start the gateway with a custom exact handler that returns an error", func() {
+		Convey("When I start the gateway with a custom suffix handler that handles the request", func() {
+
+			gw, err := New(
+				"127.0.0.1:7765",
+				u,
+				OptionUpstreamTLSConfig(&tls.Config{InsecureSkipVerify: true}),
+				OptionRegisterSuffixInterceptor("/hello", func(w http.ResponseWriter, req *http.Request, ew ErrorWriter) (InterceptorAction, string, error) {
+					w.WriteHeader(604)
+					return InterceptorActionStop, "", nil
+				}),
+			)
+			defer gw.Stop()
+
+			So(err, ShouldBeNil)
+			So(gw, ShouldNotBeNil)
+
+			testclient := &http.Client{
+				Transport: &http.Transport{
+					TLSClientConfig: &tls.Config{
+						InsecureSkipVerify: true,
+					},
+				},
+			}
+
+			gw.Start(context.Background())
+
+			Convey("Then we I call existing ep 1", func() {
+				req, _ := http.NewRequest(http.MethodGet, "http://127.0.0.1:7765/chien/hello", nil)
+				resp, err := testclient.Do(req)
+				So(err, ShouldBeNil)
+				So(resp.StatusCode, ShouldEqual, 604)
+			})
+		})
+
+		Convey("When I start the gateway with a custom suffix handler that modifies the request", func() {
+
+			gw, err := New(
+				"127.0.0.1:7765",
+				u,
+				OptionMetricsManager(&fakeMetricManager{}),
+				OptionUpstreamTLSConfig(&tls.Config{InsecureSkipVerify: true}),
+				OptionRegisterSuffixInterceptor("/ups1", func(w http.ResponseWriter, req *http.Request, ew ErrorWriter) (InterceptorAction, string, error) {
+					return InterceptorActionForwardDirect, strings.Replace(u.ups2.URL, "https://", "", 1), nil
+				}),
+			)
+			defer gw.Stop()
+
+			So(err, ShouldBeNil)
+			So(gw, ShouldNotBeNil)
+
+			testclient := &http.Client{
+				Transport: &http.Transport{
+					TLSClientConfig: &tls.Config{
+						InsecureSkipVerify: true,
+					},
+				},
+			}
+
+			gw.Start(context.Background())
+
+			Convey("Then we I call existing ep 1", func() {
+				req, _ := http.NewRequest(http.MethodGet, "http://127.0.0.1:7765/chien/ups1", nil)
+				resp, err := testclient.Do(req)
+				So(err, ShouldBeNil)
+				So(resp.StatusCode, ShouldEqual, 602)
+			})
+		})
+
+		Convey("When I start the gateway with a custom prefix handler that returns an error", func() {
 
 			gw, err := New(
 				"127.0.0.1:7765",
@@ -341,7 +410,7 @@ func TestGateway(t *testing.T) {
 			gw.Start(context.Background())
 
 			Convey("Then we I call existing ep 1", func() {
-				req, _ := http.NewRequest(http.MethodGet, "http://127.0.0.1:7765/ups1", nil)
+				req, _ := http.NewRequest(http.MethodGet, "http://127.0.0.1:7765/ups1/chien", nil)
 				resp, err := testclient.Do(req)
 				So(err, ShouldBeNil)
 				So(resp.StatusCode, ShouldEqual, 500)

gateway/options.go

@@ -45,6 +45,11 @@ const (
 	// the HTTP response writer.
 	InterceptorActionForwardWS
 
+	// InterceptorActionForwardDirect means the Gateway will continue forwarding the request directly.
+	// In that case the Interceptor must only modify the request, and MUST NOT use
+	// the HTTP response writer.
+	InterceptorActionForwardDirect
+
 	// InterceptorActionStop means the interceptor handled the request
 	// and the gateway will not do anything more.
 	InterceptorActionStop
@@ -63,6 +68,7 @@ type gwconfig struct {
 	maintenance                 bool
 	metricsManager              bahamut.MetricsManager
 	prefixInterceptors          map[string]InterceptorFunc
+	suffixInterceptors          map[string]InterceptorFunc
 	proxyProtocolEnabled        bool
 	proxyProtocolSubnet         string
 	rateLimitingBurst           int64
@@ -89,6 +95,7 @@ func newGatewayConfig() *gwconfig {
 	return &gwconfig{
 		corsOrigin:                  "*",
 		prefixInterceptors:          map[string]InterceptorFunc{},
+		suffixInterceptors:          map[string]InterceptorFunc{},
 		exactInterceptors:           map[string]InterceptorFunc{},
 		tcpRateLimitingBurst:        100,
 		tcpRateLimitingCPS:          200.0,
@@ -212,6 +219,13 @@ func OptionRegisterPrefixInterceptor(prefix string, f InterceptorFunc) Option {
 	}
 }
 
+// OptionRegisterSuffixInterceptor registers a given InterceptorFunc for the given path suffix.
+func OptionRegisterSuffixInterceptor(prefix string, f InterceptorFunc) Option {
+	return func(cfg *gwconfig) {
+		cfg.suffixInterceptors[prefix] = f
+	}
+}
+
 // OptionRegisterExactInterceptor registers a given InterceptorFunc for the given path.
 func OptionRegisterExactInterceptor(path string, f InterceptorFunc) Option {
 	return func(cfg *gwconfig) {

gateway/options_test.go

@@ -82,6 +82,14 @@ func Test_Options(t *testing.T) {
 		So(c.prefixInterceptors["/prefix"], ShouldEqual, f)
 	})
 
+	Convey("Calling OptionRegisterSuffixInterceptor should work", t, func() {
+		f := func(http.ResponseWriter, *http.Request, ErrorWriter) (InterceptorAction, string, error) {
+			return InterceptorActionForward, "", nil
+		}
+		OptionRegisterSuffixInterceptor("/suffix", f)(c)
+		So(c.suffixInterceptors["/suffix"], ShouldEqual, f)
+	})
+
 	Convey("Calling OptionRegisterExactInterceptor should work", t, func() {
 		f := func(http.ResponseWriter, *http.Request, ErrorWriter) (InterceptorAction, string, error) {
 			return InterceptorActionForward, "", nil

interfaces.go

@@ -80,6 +80,12 @@ type Server interface {
 	Run(context.Context)
 }
 
+// A ResponseWriter is a function you can use in
+// the Context to handle the writing of the response by
+// yourself. You are responsible for the full handling of the response,
+// including encoding, setting the CORS headers etc.
+type ResponseWriter func(w http.ResponseWriter) int
+
 // A Context contains all information about a current operation.
 type Context interface {
 
@@ -102,8 +108,22 @@ type Context interface {
 	OutputData() interface{}
 
 	// SetOutputData sets the data that will be returned to the client.
+	//
+	// If you use SetOutputData after having already used SetResponseWriter,
+	// the call will panic.
 	SetOutputData(interface{})
 
+	// SetResponseWriter sets the ResponseWriter function to use to write the response back to the client.
+	//
+	// No additional operation or check will be performed by Bahamut. You are responsible
+	// for correctly encoding the response, setting the header etc. This is useful when
+	// you want to handle a route that is note really fitting in the handling of an elemental Model
+	// like for instance handling file download, response streaming etc.
+	//
+	// If you use SetResponseWriter after having already used SetOutputData,
+	// the call will panic.
+	SetResponseWriter(ResponseWriter)
+
 	// Set count sets the count.
 	SetCount(int)
 

rest_server.go

@@ -285,7 +285,17 @@ func (a *restServer) makeHandler(handler handlerFunc) http.HandlerFunc {
 			}
 		}
 
-		code := writeHTTPResponse(a.cfg.security.CORSOrigin, w, handler(newContext(ctx, request), a.cfg, a.processorFinder, a.pusher))
+		bctx := newContext(ctx, request)
+		resp := handler(bctx, a.cfg, a.processorFinder, a.pusher)
+		var code int
+
+		switch {
+		case bctx.responseWriter != nil:
+			code = bctx.responseWriter(w)
+		default:
+			code = writeHTTPResponse(a.cfg.security.CORSOrigin, w, resp)
+		}
+
 		if measure != nil {
 			measure(code, opentracing.SpanFromContext(ctx))
 		}