github.com/epsagon/epsagon-go@v1.39.0/wrappers/gin/gin.go (about)

     1  package epsagongin
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net/http"
     7  
     8  	"github.com/epsagon/epsagon-go/epsagon"
     9  	"github.com/epsagon/epsagon-go/protocol"
    10  	"github.com/epsagon/epsagon-go/tracer"
    11  	epsagonhttp "github.com/epsagon/epsagon-go/wrappers/net/http"
    12  	"github.com/gin-gonic/gin"
    13  )
    14  
    15  // TracerKey is the key of the epsagon tracer in the gin.Context Keys map passed to the handlers
    16  const TracerKey = "EpsagonTracer"
    17  
    18  // EpsagonContext creates a context.Background() with epsagon's associated tracer for nexted instrumentations
    19  func EpsagonContext(c *gin.Context) context.Context {
    20  	return epsagon.ContextWithTracer(c.Keys[TracerKey].(tracer.Tracer))
    21  }
    22  
    23  // GinRouterWrapper is an epsagon instumentation wrapper for gin.RouterGroup
    24  type GinRouterWrapper struct {
    25  	gin.IRouter
    26  	Hostname string
    27  	Config   *epsagon.Config
    28  }
    29  
    30  type wrappedGinWriter struct {
    31  	gin.ResponseWriter
    32  	htrw http.ResponseWriter
    33  }
    34  
    35  func wrapGinWriter(c *gin.Context, triggerEvent *protocol.Event) {
    36  	wrappedResponseWriter := &wrappedGinWriter{
    37  		ResponseWriter: c.Writer,
    38  		htrw:           epsagonhttp.CreateWrappedResponseWriter(c.Writer, triggerEvent.Resource),
    39  	}
    40  	c.Writer = wrappedResponseWriter
    41  }
    42  
    43  func postExecutionUpdates(
    44  	wrapperTracer tracer.Tracer, triggerEvent *protocol.Event,
    45  	c *gin.Context, handlerWrapper *epsagon.GenericWrapper) {
    46  	runner := handlerWrapper.GetRunnerEvent()
    47  	if runner != nil {
    48  		runner.Resource.Type = "gin"
    49  	}
    50  	wrappedResponseWriter, ok := c.Writer.(*wrappedGinWriter)
    51  	if ok {
    52  		wrappedResponseWriter.htrw.(*epsagonhttp.WrappedResponseWriter).UpdateResource()
    53  	}
    54  	userError := recover()
    55  	if userError != nil {
    56  		triggerEvent.Resource.Metadata["status_code"] = "500"
    57  		panic(userError)
    58  	}
    59  }
    60  
    61  func wrapGinHandler(handler gin.HandlerFunc, hostname string, relativePath string, config *epsagon.Config) gin.HandlerFunc {
    62  	if config == nil {
    63  		config = &epsagon.Config{}
    64  	}
    65  	return func(c *gin.Context) {
    66  		wrapperTracer := tracer.CreateTracer(&config.Config)
    67  		wrapperTracer.Start()
    68  		defer wrapperTracer.SendStopSignal()
    69  
    70  		c.Set(TracerKey, wrapperTracer)
    71  		wrapper := epsagon.WrapGenericFunction(
    72  			handler, config, wrapperTracer, false, relativePath,
    73  		)
    74  		triggerEvent := epsagonhttp.CreateHTTPTriggerEvent(
    75  			wrapperTracer, c.Request, hostname)
    76  		wrapperTracer.AddEvent(triggerEvent)
    77  		if !config.MetadataOnly {
    78  			wrapGinWriter(c, triggerEvent)
    79  		}
    80  		defer postExecutionUpdates(wrapperTracer, triggerEvent, c, wrapper)
    81  		wrapper.Call(c)
    82  		triggerEvent.Resource.Metadata["status_code"] = fmt.Sprint(c.Writer.Status())
    83  	}
    84  }
    85  
    86  // Handle is a wrapper for gin.RouterGroup.Handle that adds epsagon instrumentaiton and event triggers
    87  // to all invocations of that handler
    88  func (router *GinRouterWrapper) Handle(httpMethod, relativePath string, handlers ...gin.HandlerFunc) gin.IRoutes {
    89  	if len(handlers) >= 1 {
    90  		handlers[0] = wrapGinHandler(handlers[0], router.Hostname, relativePath, router.Config)
    91  	}
    92  	return router.IRouter.Handle(httpMethod, relativePath, handlers...)
    93  }
    94  
    95  // POST is a shortcut for router.Handle("POST", path, handle).
    96  func (router *GinRouterWrapper) POST(relativePath string, handlers ...gin.HandlerFunc) gin.IRoutes {
    97  	return router.Handle(http.MethodPost, relativePath, handlers...)
    98  }
    99  
   100  // GET is a shortcut for router.Handle("GET", path, handle).
   101  func (router *GinRouterWrapper) GET(relativePath string, handlers ...gin.HandlerFunc) gin.IRoutes {
   102  	return router.Handle(http.MethodGet, relativePath, handlers...)
   103  }
   104  
   105  // DELETE is a shortcut for router.Handle("DELETE", path, handle).
   106  func (router *GinRouterWrapper) DELETE(relativePath string, handlers ...gin.HandlerFunc) gin.IRoutes {
   107  	return router.Handle(http.MethodDelete, relativePath, handlers...)
   108  }
   109  
   110  // PATCH is a shortcut for router.Handle("PATCH", path, handle).
   111  func (router *GinRouterWrapper) PATCH(relativePath string, handlers ...gin.HandlerFunc) gin.IRoutes {
   112  	return router.Handle(http.MethodPatch, relativePath, handlers...)
   113  }
   114  
   115  // PUT is a shortcut for router.Handle("PUT", path, handle).
   116  func (router *GinRouterWrapper) PUT(relativePath string, handlers ...gin.HandlerFunc) gin.IRoutes {
   117  	return router.Handle(http.MethodPut, relativePath, handlers...)
   118  }
   119  
   120  // OPTIONS is a shortcut for router.Handle("OPTIONS", path, handle).
   121  func (router *GinRouterWrapper) OPTIONS(relativePath string, handlers ...gin.HandlerFunc) gin.IRoutes {
   122  	return router.Handle(http.MethodOptions, relativePath, handlers...)
   123  }
   124  
   125  // HEAD is a shortcut for router.Handle("HEAD", path, handle).
   126  func (router *GinRouterWrapper) HEAD(relativePath string, handlers ...gin.HandlerFunc) gin.IRoutes {
   127  	return router.Handle(http.MethodHead, relativePath, handlers...)
   128  }
   129  
   130  // Run is a shortcut for router.IRouter.(*gin.Engine).Run()
   131  func (router *GinRouterWrapper) Run(addr ...string) error {
   132  	engine, ok := router.IRouter.(*gin.Engine)
   133  
   134  	if !ok {
   135  		return fmt.Errorf("Could not start Gin engine")
   136  	}
   137  
   138  	return engine.Run(addr...)
   139  }
   140  
   141  func (grw *wrappedGinWriter) Header() http.Header {
   142  	return grw.htrw.Header()
   143  }
   144  
   145  func (grw *wrappedGinWriter) Write(data []byte) (int, error) {
   146  	return grw.htrw.Write(data)
   147  }
   148  
   149  func (grw *wrappedGinWriter) WriteHeader(statusCode int) {
   150  	grw.htrw.WriteHeader(statusCode)
   151  }