From 19ce63e5ef91d3c8151bdc22db5c2bca06ca7353 Mon Sep 17 00:00:00 2001
From: Marcus Cobden <marcus@weave.works>
Date: Thu, 23 Aug 2018 14:20:14 +0100
Subject: [PATCH] Add HTTP tracing middleware

---
 middleware/http_tracing.go | 27 +++++++++++++++++++++++++++
 server/server.go           |  5 +----
 2 files changed, 28 insertions(+), 4 deletions(-)
 create mode 100644 middleware/http_tracing.go

diff --git a/middleware/http_tracing.go b/middleware/http_tracing.go
new file mode 100644
index 00000000..b270e3b8
--- /dev/null
+++ b/middleware/http_tracing.go
@@ -0,0 +1,27 @@
+package middleware
+
+import (
+	"net/http"
+
+	"github.com/opentracing-contrib/go-stdlib/nethttp"
+	"github.com/opentracing/opentracing-go"
+)
+
+// Tracer is a middleware which traces incoming requests.
+type Tracer struct{}
+
+// Wrap implements Interface
+func (t Tracer) Wrap(next http.Handler) http.Handler {
+	traceHandler := nethttp.Middleware(opentracing.GlobalTracer(), next)
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var maybeTracer http.Handler
+		// Don't try and trace websocket requests because nethttp.Middleware
+		// doesn't support http.Hijack yet
+		if IsWSHandshakeRequest(r) {
+			maybeTracer = next
+		} else {
+			maybeTracer = traceHandler
+		}
+		maybeTracer.ServeHTTP(w, r)
+	})
+}
diff --git a/server/server.go b/server/server.go
index d69dc4d5..98d17e63 100644
--- a/server/server.go
+++ b/server/server.go
@@ -11,7 +11,6 @@ import (
 	"github.com/gorilla/mux"
 	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
 	"github.com/mwitkow/go-grpc-middleware"
-	"github.com/opentracing-contrib/go-stdlib/nethttp"
 	"github.com/opentracing/opentracing-go"
 	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/net/context"
@@ -147,9 +146,7 @@ func New(cfg Config) (*Server, error) {
 			Duration:     requestDuration,
 			RouteMatcher: router,
 		},
-		middleware.Func(func(handler http.Handler) http.Handler {
-			return nethttp.Middleware(opentracing.GlobalTracer(), handler)
-		}),
+		middleware.Tracer{},
 	}
 	httpMiddleware = append(httpMiddleware, cfg.HTTPMiddleware...)
 	httpServer := &http.Server{