github.com/newrelic/go-agent@v3.26.0+incompatible/_integrations/nrgin/v1/nrgin.go (about)

     1  // Copyright 2020 New Relic Corporation. All rights reserved.
     2  // SPDX-License-Identifier: Apache-2.0
     3  
     4  // Package nrgin instruments https://github.com/gin-gonic/gin applications.
     5  //
     6  // Use this package to instrument inbound requests handled by a gin.Engine.
     7  // Call nrgin.Middleware to get a gin.HandlerFunc which can be added to your
     8  // application as a middleware:
     9  //
    10  //	router := gin.Default()
    11  //	// Add the nrgin middleware before other middlewares or routes:
    12  //	router.Use(nrgin.Middleware(app))
    13  //
    14  // Example: https://github.com/newrelic/go-agent/tree/master/_integrations/nrgin/v1/example/main.go
    15  package nrgin
    16  
    17  import (
    18  	"net/http"
    19  
    20  	"github.com/gin-gonic/gin"
    21  	newrelic "github.com/newrelic/go-agent"
    22  	"github.com/newrelic/go-agent/internal"
    23  )
    24  
    25  func init() { internal.TrackUsage("integration", "framework", "gin", "v1") }
    26  
    27  // headerResponseWriter gives the transaction access to response headers and the
    28  // response code.
    29  type headerResponseWriter struct{ w gin.ResponseWriter }
    30  
    31  func (w *headerResponseWriter) Header() http.Header       { return w.w.Header() }
    32  func (w *headerResponseWriter) Write([]byte) (int, error) { return 0, nil }
    33  func (w *headerResponseWriter) WriteHeader(int)           {}
    34  
    35  var _ http.ResponseWriter = &headerResponseWriter{}
    36  
    37  // replacementResponseWriter mimics the behavior of gin.ResponseWriter which
    38  // buffers the response code rather than writing it when
    39  // gin.ResponseWriter.WriteHeader is called.
    40  type replacementResponseWriter struct {
    41  	gin.ResponseWriter
    42  	txn     newrelic.Transaction
    43  	code    int
    44  	written bool
    45  }
    46  
    47  var _ gin.ResponseWriter = &replacementResponseWriter{}
    48  
    49  func (w *replacementResponseWriter) flushHeader() {
    50  	if !w.written {
    51  		w.txn.WriteHeader(w.code)
    52  		w.written = true
    53  	}
    54  }
    55  
    56  func (w *replacementResponseWriter) WriteHeader(code int) {
    57  	w.code = code
    58  	w.ResponseWriter.WriteHeader(code)
    59  }
    60  
    61  func (w *replacementResponseWriter) Write(data []byte) (int, error) {
    62  	w.flushHeader()
    63  	return w.ResponseWriter.Write(data)
    64  }
    65  
    66  func (w *replacementResponseWriter) WriteString(s string) (int, error) {
    67  	w.flushHeader()
    68  	return w.ResponseWriter.WriteString(s)
    69  }
    70  
    71  func (w *replacementResponseWriter) WriteHeaderNow() {
    72  	w.flushHeader()
    73  	w.ResponseWriter.WriteHeaderNow()
    74  }
    75  
    76  // Context avoids making this package 1.7+ specific.
    77  type Context interface {
    78  	Value(key interface{}) interface{}
    79  }
    80  
    81  // Transaction returns the transaction stored inside the context, or nil if not
    82  // found.
    83  func Transaction(c Context) newrelic.Transaction {
    84  	if v := c.Value(internal.GinTransactionContextKey); nil != v {
    85  		if txn, ok := v.(newrelic.Transaction); ok {
    86  			return txn
    87  		}
    88  	}
    89  	if v := c.Value(internal.TransactionContextKey); nil != v {
    90  		if txn, ok := v.(newrelic.Transaction); ok {
    91  			return txn
    92  		}
    93  	}
    94  	return nil
    95  }
    96  
    97  // Middleware creates a Gin middleware that instruments requests.
    98  //
    99  //	router := gin.Default()
   100  //	// Add the nrgin middleware before other middlewares or routes:
   101  //	router.Use(nrgin.Middleware(app))
   102  //
   103  func Middleware(app newrelic.Application) gin.HandlerFunc {
   104  	return func(c *gin.Context) {
   105  		if app != nil {
   106  			name := c.HandlerName()
   107  			w := &headerResponseWriter{w: c.Writer}
   108  			txn := app.StartTransaction(name, w, c.Request)
   109  			defer txn.End()
   110  
   111  			repl := &replacementResponseWriter{
   112  				ResponseWriter: c.Writer,
   113  				txn:            txn,
   114  				code:           http.StatusOK,
   115  			}
   116  			c.Writer = repl
   117  			defer repl.flushHeader()
   118  
   119  			c.Set(internal.GinTransactionContextKey, txn)
   120  		}
   121  		c.Next()
   122  	}
   123  }