server: use httpsnoop to wrap http.ResponseWriter

bmeg · Jan 17, 2020 · 4f2bf2a · 4f2bf2a
1 parent 6e7e7b5
commit 4f2bf2a
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 15 deletions.
diff --git a/cmd/server/main.go b/cmd/server/main.go
@@ -143,4 +143,5 @@ func init() {
 	flags.BoolVar(&conf.Server.ReadOnly, "read-only", conf.Server.ReadOnly, "Start server in read-only mode")
 	flags.StringVar(&conf.Logger.Level, "log-level", conf.Logger.Level, "Log level [info, debug, warn, error]")
 	flags.StringVar(&conf.Logger.Formatter, "log-format", conf.Logger.Formatter, "Log format [text, json]")
+	flags.BoolVar(&conf.Server.RequestLogging.Enable, "log-requests", conf.Server.RequestLogging.Enable, "Log all requests")
 }
diff --git a/go.mod b/go.mod
@@ -12,6 +12,7 @@ require (
 	github.com/dgraph-io/badger/v2 v2.0.1
 	github.com/dlclark/regexp2 v1.1.6 // indirect
 	github.com/dop251/goja v0.0.0-20190429205339-8d6ee3d16611
+	github.com/felixge/httpsnoop v1.0.1
 	github.com/ghodss/yaml v1.0.0
 	github.com/globalsign/mgo v0.0.0-20181015135952-eeefdecb41b8
 	github.com/go-sourcemap/sourcemap v2.1.2+incompatible // indirect

diff --git a/go.sum b/go.sum
@@ -65,6 +65,8 @@ github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7
 github.com/fatih/color v1.6.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
 github.com/fatih/color v1.7.0 h1:DkWD4oS2D8LGGgTQ6IvwJJXSL5Vp2ffcQg58nFV38Ys=
 github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
+github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
+github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/fortytw2/leaktest v1.2.0/go.mod h1:jDsjWgpAGjm2CA7WthBh/CdZYEPF31XHquHwclZch5g=
 github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=

diff --git a/server/logging_middleware.go b/server/logging_middleware.go
@@ -1,7 +1,6 @@
 package server
 
 import (
-	"net/http"
 	"strings"
 	"time"
 
@@ -98,13 +97,3 @@ func (l *loggingServerStream) RecvMsg(m interface{}) error {
 	l.request = m
 	return l.ServerStream.RecvMsg(m)
 }
-
-type loggingResponseWriter struct {
-	http.ResponseWriter
-	statusCode int
-}
-
-func (lrw *loggingResponseWriter) WriteHeader(code int) {
-	lrw.statusCode = code
-	lrw.ResponseWriter.WriteHeader(code)
-}
diff --git a/server/server.go b/server/server.go
@@ -16,6 +16,7 @@ import (
 	"github.com/bmeg/grip/gripql"
 	"github.com/bmeg/grip/log"
 	"github.com/bmeg/grip/util/rpc"
+	"github.com/felixge/httpsnoop"
 	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	"github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"golang.org/x/net/context"
@@ -155,8 +156,7 @@ func (server *GripServer) Serve(pctx context.Context) error {
 			}
 
 			// handle the request
-			lrw := &loggingResponseWriter{resp, http.StatusOK}
-			grpcMux.ServeHTTP(lrw, req)
+			m := httpsnoop.CaptureMetrics(grpcMux, resp, req)
 
 			if !server.conf.RequestLogging.Enable {
 				return
@@ -171,9 +171,9 @@ func (server *GripServer) Serve(pctx context.Context) error {
 				"request": string(body),
 				"header":  headers,
 				"latency": time.Since(start).String(),
-				"status":  lrw.statusCode,
+				"status":  m.Code,
 			})
-			if lrw.statusCode == http.StatusOK {
+			if m.Code == http.StatusOK {
 				entry.Info("HTTP server responded")
 			} else {
 				entry.Error("HTTP server responded")