...

Source file src/github.com/rs/zerolog/hlog/hlog.go

Documentation: github.com/rs/zerolog/hlog

     1  // Package hlog provides a set of http.Handler helpers for zerolog.
     2  package hlog
     3  
     4  import (
     5  	"context"
     6  	"net/http"
     7  	"time"
     8  
     9  	"github.com/rs/xid"
    10  	"github.com/rs/zerolog"
    11  	"github.com/rs/zerolog/hlog/internal/mutil"
    12  	"github.com/rs/zerolog/log"
    13  )
    14  
    15  // FromRequest gets the logger in the request's context.
    16  // This is a shortcut for log.Ctx(r.Context())
    17  func FromRequest(r *http.Request) *zerolog.Logger {
    18  	return log.Ctx(r.Context())
    19  }
    20  
    21  // NewHandler injects log into requests context.
    22  func NewHandler(log zerolog.Logger) func(http.Handler) http.Handler {
    23  	return func(next http.Handler) http.Handler {
    24  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    25  			// Create a copy of the logger (including internal context slice)
    26  			// to prevent data race when using UpdateContext.
    27  			l := log.With().Logger()
    28  			r = r.WithContext(l.WithContext(r.Context()))
    29  			next.ServeHTTP(w, r)
    30  		})
    31  	}
    32  }
    33  
    34  // URLHandler adds the requested URL as a field to the context's logger
    35  // using fieldKey as field key.
    36  func URLHandler(fieldKey string) func(next http.Handler) http.Handler {
    37  	return func(next http.Handler) http.Handler {
    38  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    39  			log := zerolog.Ctx(r.Context())
    40  			log.UpdateContext(func(c zerolog.Context) zerolog.Context {
    41  				return c.Str(fieldKey, r.URL.String())
    42  			})
    43  			next.ServeHTTP(w, r)
    44  		})
    45  	}
    46  }
    47  
    48  // MethodHandler adds the request method as a field to the context's logger
    49  // using fieldKey as field key.
    50  func MethodHandler(fieldKey string) func(next http.Handler) http.Handler {
    51  	return func(next http.Handler) http.Handler {
    52  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    53  			log := zerolog.Ctx(r.Context())
    54  			log.UpdateContext(func(c zerolog.Context) zerolog.Context {
    55  				return c.Str(fieldKey, r.Method)
    56  			})
    57  			next.ServeHTTP(w, r)
    58  		})
    59  	}
    60  }
    61  
    62  // RequestHandler adds the request method and URL as a field to the context's logger
    63  // using fieldKey as field key.
    64  func RequestHandler(fieldKey string) func(next http.Handler) http.Handler {
    65  	return func(next http.Handler) http.Handler {
    66  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    67  			log := zerolog.Ctx(r.Context())
    68  			log.UpdateContext(func(c zerolog.Context) zerolog.Context {
    69  				return c.Str(fieldKey, r.Method+" "+r.URL.String())
    70  			})
    71  			next.ServeHTTP(w, r)
    72  		})
    73  	}
    74  }
    75  
    76  // RemoteAddrHandler adds the request's remote address as a field to the context's logger
    77  // using fieldKey as field key.
    78  func RemoteAddrHandler(fieldKey string) func(next http.Handler) http.Handler {
    79  	return func(next http.Handler) http.Handler {
    80  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    81  			if r.RemoteAddr != "" {
    82  				log := zerolog.Ctx(r.Context())
    83  				log.UpdateContext(func(c zerolog.Context) zerolog.Context {
    84  					return c.Str(fieldKey, r.RemoteAddr)
    85  				})
    86  			}
    87  			next.ServeHTTP(w, r)
    88  		})
    89  	}
    90  }
    91  
    92  // UserAgentHandler adds the request's user-agent as a field to the context's logger
    93  // using fieldKey as field key.
    94  func UserAgentHandler(fieldKey string) func(next http.Handler) http.Handler {
    95  	return func(next http.Handler) http.Handler {
    96  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    97  			if ua := r.Header.Get("User-Agent"); ua != "" {
    98  				log := zerolog.Ctx(r.Context())
    99  				log.UpdateContext(func(c zerolog.Context) zerolog.Context {
   100  					return c.Str(fieldKey, ua)
   101  				})
   102  			}
   103  			next.ServeHTTP(w, r)
   104  		})
   105  	}
   106  }
   107  
   108  // RefererHandler adds the request's referer as a field to the context's logger
   109  // using fieldKey as field key.
   110  func RefererHandler(fieldKey string) func(next http.Handler) http.Handler {
   111  	return func(next http.Handler) http.Handler {
   112  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   113  			if ref := r.Header.Get("Referer"); ref != "" {
   114  				log := zerolog.Ctx(r.Context())
   115  				log.UpdateContext(func(c zerolog.Context) zerolog.Context {
   116  					return c.Str(fieldKey, ref)
   117  				})
   118  			}
   119  			next.ServeHTTP(w, r)
   120  		})
   121  	}
   122  }
   123  
   124  // ProtoHandler adds the requests protocol version as a field to the context logger
   125  // using fieldKey as field Key.
   126  func ProtoHandler(fieldKey string) func(next http.Handler) http.Handler {
   127  	return func(next http.Handler) http.Handler {
   128  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   129  			log := zerolog.Ctx(r.Context())
   130  			log.UpdateContext(func(c zerolog.Context) zerolog.Context {
   131  				return c.Str(fieldKey, r.Proto)
   132  			})
   133  			next.ServeHTTP(w, r)
   134  		})
   135  	}
   136  }
   137  
   138  type idKey struct{}
   139  
   140  // IDFromRequest returns the unique id associated to the request if any.
   141  func IDFromRequest(r *http.Request) (id xid.ID, ok bool) {
   142  	if r == nil {
   143  		return
   144  	}
   145  	return IDFromCtx(r.Context())
   146  }
   147  
   148  // IDFromCtx returns the unique id associated to the context if any.
   149  func IDFromCtx(ctx context.Context) (id xid.ID, ok bool) {
   150  	id, ok = ctx.Value(idKey{}).(xid.ID)
   151  	return
   152  }
   153  
   154  // CtxWithID adds the given xid.ID to the context
   155  func CtxWithID(ctx context.Context, id xid.ID) context.Context {
   156  	return context.WithValue(ctx, idKey{}, id)
   157  }
   158  
   159  // RequestIDHandler returns a handler setting a unique id to the request which can
   160  // be gathered using IDFromRequest(req). This generated id is added as a field to the
   161  // logger using the passed fieldKey as field name. The id is also added as a response
   162  // header if the headerName is not empty.
   163  //
   164  // The generated id is a URL safe base64 encoded mongo object-id-like unique id.
   165  // Mongo unique id generation algorithm has been selected as a trade-off between
   166  // size and ease of use: UUID is less space efficient and snowflake requires machine
   167  // configuration.
   168  func RequestIDHandler(fieldKey, headerName string) func(next http.Handler) http.Handler {
   169  	return func(next http.Handler) http.Handler {
   170  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   171  			ctx := r.Context()
   172  			id, ok := IDFromRequest(r)
   173  			if !ok {
   174  				id = xid.New()
   175  				ctx = CtxWithID(ctx, id)
   176  				r = r.WithContext(ctx)
   177  			}
   178  			if fieldKey != "" {
   179  				log := zerolog.Ctx(ctx)
   180  				log.UpdateContext(func(c zerolog.Context) zerolog.Context {
   181  					return c.Str(fieldKey, id.String())
   182  				})
   183  			}
   184  			if headerName != "" {
   185  				w.Header().Set(headerName, id.String())
   186  			}
   187  			next.ServeHTTP(w, r)
   188  		})
   189  	}
   190  }
   191  
   192  // CustomHeaderHandler adds given header from request's header as a field to
   193  // the context's logger using fieldKey as field key.
   194  func CustomHeaderHandler(fieldKey, header string) func(next http.Handler) http.Handler {
   195  	return func(next http.Handler) http.Handler {
   196  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   197  			if val := r.Header.Get(header); val != "" {
   198  				log := zerolog.Ctx(r.Context())
   199  				log.UpdateContext(func(c zerolog.Context) zerolog.Context {
   200  					return c.Str(fieldKey, val)
   201  				})
   202  			}
   203  			next.ServeHTTP(w, r)
   204  		})
   205  	}
   206  }
   207  
   208  // AccessHandler returns a handler that call f after each request.
   209  func AccessHandler(f func(r *http.Request, status, size int, duration time.Duration)) func(next http.Handler) http.Handler {
   210  	return func(next http.Handler) http.Handler {
   211  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   212  			start := time.Now()
   213  			lw := mutil.WrapWriter(w)
   214  			next.ServeHTTP(lw, r)
   215  			f(r, lw.Status(), lw.BytesWritten(), time.Since(start))
   216  		})
   217  	}
   218  }
   219  

View as plain text