github.com/iron-io/functions@v0.0.0-20180820112432-d59d7d1c40b2/api/server/middleware.go (about)

     1  // TODO: it would be nice to move these into the top level folder so people can use these with the "functions" package, eg: functions.AddMiddleware(...)
     2  package server
     3  
     4  import (
     5  	"context"
     6  	"net/http"
     7  
     8  	"github.com/Sirupsen/logrus"
     9  	"github.com/gin-gonic/gin"
    10  	"github.com/iron-io/functions/api/models"
    11  )
    12  
    13  // Middleware is the interface required for implementing functions middlewar
    14  type Middleware interface {
    15  	// Serve is what the Middleware must implement. Can modify the request, write output, etc.
    16  	// todo: should we abstract the HTTP out of this?  In case we want to support other protocols.
    17  	Serve(ctx MiddlewareContext, w http.ResponseWriter, r *http.Request, app *models.App) error
    18  }
    19  
    20  // MiddlewareFunc func form of Middleware
    21  type MiddlewareFunc func(ctx MiddlewareContext, w http.ResponseWriter, r *http.Request, app *models.App) error
    22  
    23  // Serve wrapper
    24  func (f MiddlewareFunc) Serve(ctx MiddlewareContext, w http.ResponseWriter, r *http.Request, app *models.App) error {
    25  	return f(ctx, w, r, app)
    26  }
    27  
    28  // MiddlewareContext extends context.Context for Middleware
    29  type MiddlewareContext interface {
    30  	context.Context
    31  	// Middleware can call Next() explicitly to call the next middleware in the chain. If Next() is not called and an error is not returned, Next() will automatically be called.
    32  	Next()
    33  	// Index returns the index of where we're at in the chain
    34  	Index() int
    35  }
    36  
    37  type middlewareContextImpl struct {
    38  	context.Context
    39  
    40  	ginContext  *gin.Context
    41  	nextCalled  bool
    42  	index       int
    43  	middlewares []Middleware
    44  }
    45  
    46  func (c *middlewareContextImpl) Next() {
    47  	c.nextCalled = true
    48  	c.index++
    49  	c.serveNext()
    50  }
    51  
    52  func (c *middlewareContextImpl) serveNext() {
    53  	if c.Index() >= len(c.middlewares) {
    54  		return
    55  	}
    56  	// make shallow copy:
    57  	fctx2 := *c
    58  	fctx2.nextCalled = false
    59  	r := c.ginContext.Request.WithContext(fctx2)
    60  	err := c.middlewares[c.Index()].Serve(&fctx2, c.ginContext.Writer, r, nil)
    61  	if err != nil {
    62  		logrus.WithError(err).Warnln("Middleware error")
    63  		// todo: might be a good idea to check if anything is written yet, and if not, output the error: simpleError(err)
    64  		// see: http://stackoverflow.com/questions/39415827/golang-http-check-if-responsewriter-has-been-written
    65  		c.ginContext.Abort()
    66  		return
    67  	}
    68  	if !fctx2.nextCalled {
    69  		// then we automatically call next
    70  		fctx2.Next()
    71  	}
    72  
    73  }
    74  
    75  func (c *middlewareContextImpl) Index() int {
    76  	return c.index
    77  }
    78  
    79  func (s *Server) middlewareWrapperFunc(ctx context.Context) gin.HandlerFunc {
    80  	return func(c *gin.Context) {
    81  		if len(s.middlewares) == 0 {
    82  			return
    83  		}
    84  		ctx = c.MustGet("ctx").(context.Context)
    85  		fctx := &middlewareContextImpl{Context: ctx}
    86  		fctx.index = -1
    87  		fctx.ginContext = c
    88  		fctx.middlewares = s.middlewares
    89  		// start the chain:
    90  		fctx.Next()
    91  	}
    92  }
    93  
    94  // AddAppEndpoint adds an endpoints to /v1/apps/:app/x
    95  func (s *Server) AddMiddleware(m Middleware) {
    96  	s.middlewares = append(s.middlewares, m)
    97  }
    98  
    99  // AddAppEndpoint adds an endpoints to /v1/apps/:app/x
   100  func (s *Server) AddMiddlewareFunc(m func(ctx MiddlewareContext, w http.ResponseWriter, r *http.Request, app *models.App) error) {
   101  	s.AddMiddleware(MiddlewareFunc(m))
   102  }