github.com/gramework/gramework@v1.8.1-0.20231027140105-82555c9057f5/middleware.go (about)

     1  // Copyright 2017-present Kirill Danshin and Gramework contributors
     2  // Copyright 2019-present Highload LTD (UK CN: 11893420)
     3  //
     4  // Licensed under the Apache License, Version 2.0 (the "License");
     5  // you may not use this file except in compliance with the License.
     6  // You may obtain a copy of the License at
     7  //
     8  //     http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  
    11  package gramework
    12  
    13  import (
    14  	"errors"
    15  
    16  	"github.com/valyala/fasthttp"
    17  )
    18  
    19  var (
    20  	// ErrEmptyMiddleware can be returned by App.Use*, if middleware is nil
    21  	ErrEmptyMiddleware = errors.New("can't use nil middleware")
    22  
    23  	// ErrUnsupportedMiddlewareType can be returned by App.Use*, if middleware type is unsupported
    24  	ErrUnsupportedMiddlewareType = errors.New("unsupported middleware type")
    25  )
    26  
    27  // CORSMiddleware provides gramework handler with ctx.CORS() call
    28  func (app *App) CORSMiddleware(domains ...string) func(*Context) {
    29  	return func(ctx *Context) {
    30  		ctx.CORS(domains...)
    31  	}
    32  }
    33  
    34  // Use the middleware before request processing
    35  func (app *App) Use(middleware interface{}) error {
    36  	if middleware == nil {
    37  		return ErrEmptyMiddleware
    38  	}
    39  	processor, err := app.middlewareProcessor(middleware)
    40  	app.middlewaresMu.Lock()
    41  	if err == nil {
    42  		app.middlewares = append(app.middlewares, processor)
    43  	}
    44  	app.middlewaresMu.Unlock()
    45  
    46  	return err
    47  }
    48  
    49  // UsePre registers middleware before any other middleware. Use only for metrics or access control!
    50  func (app *App) UsePre(middleware interface{}) error {
    51  	if middleware == nil {
    52  		return ErrEmptyMiddleware
    53  	}
    54  	processor, err := app.middlewareProcessor(middleware)
    55  	app.preMiddlewaresMu.Lock()
    56  	if err == nil {
    57  		app.preMiddlewares = append(app.preMiddlewares, processor)
    58  	}
    59  	app.preMiddlewaresMu.Unlock()
    60  
    61  	return err
    62  }
    63  
    64  // UseAfterRequest the middleware after request processing
    65  func (app *App) UseAfterRequest(middleware interface{}) error {
    66  	if middleware == nil {
    67  		return ErrEmptyMiddleware
    68  	}
    69  
    70  	processor, err := app.middlewareProcessor(middleware)
    71  	app.middlewaresAfterRequestMu.Lock()
    72  	if err == nil {
    73  		app.middlewaresAfterRequest = append(app.middlewaresAfterRequest, processor)
    74  	}
    75  	app.middlewaresAfterRequestMu.Unlock()
    76  
    77  	return nil
    78  }
    79  
    80  func (app *App) middlewareProcessor(middleware interface{}) (func(*Context), error) {
    81  	// we can register middlewares slowly to serve requests faster
    82  	switch m := middleware.(type) {
    83  	case func():
    84  		return func(*Context) {
    85  			m()
    86  		}, nil
    87  	case func(*Context):
    88  		// if middleware is that type, we can just return
    89  		// the middleware itself, to save some resources
    90  		// required to run the function via closures
    91  		return m, nil
    92  	case func(*Context) error:
    93  		return func(ctx *Context) {
    94  			if err := m(ctx); err != nil {
    95  				// if error occurred, we can stop processing even slowly
    96  				ctx.Logger.Errorf("Middleware error: %s", err)
    97  				ctx.Err500()
    98  				ctx.middlewaresShouldStopProcessing = true
    99  			}
   100  		}, nil
   101  	case func(*fasthttp.RequestCtx):
   102  		return func(ctx *Context) {
   103  			m(ctx.RequestCtx)
   104  		}, nil
   105  	}
   106  	return nil, ErrUnsupportedMiddlewareType
   107  }