1// Copyright 2018, OpenCensus Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package ochttp
16
17import (
18	"context"
19	"io"
20	"net/http"
21	"strconv"
22	"sync"
23	"time"
24
25	"go.opencensus.io/stats"
26	"go.opencensus.io/tag"
27	"go.opencensus.io/trace"
28	"go.opencensus.io/trace/propagation"
29)
30
31// Handler is an http.Handler wrapper to instrument your HTTP server with
32// OpenCensus. It supports both stats and tracing.
33//
34// Tracing
35//
36// This handler is aware of the incoming request's span, reading it from request
37// headers as configured using the Propagation field.
38// The extracted span can be accessed from the incoming request's
39// context.
40//
41//    span := trace.FromContext(r.Context())
42//
43// The server span will be automatically ended at the end of ServeHTTP.
44type Handler struct {
45	// Propagation defines how traces are propagated. If unspecified,
46	// B3 propagation will be used.
47	Propagation propagation.HTTPFormat
48
49	// Handler is the handler used to handle the incoming request.
50	Handler http.Handler
51
52	// StartOptions are applied to the span started by this Handler around each
53	// request.
54	//
55	// StartOptions.SpanKind will always be set to trace.SpanKindServer
56	// for spans started by this transport.
57	StartOptions trace.StartOptions
58
59	// GetStartOptions allows to set start options per request. If set,
60	// StartOptions is going to be ignored.
61	GetStartOptions func(*http.Request) trace.StartOptions
62
63	// IsPublicEndpoint should be set to true for publicly accessible HTTP(S)
64	// servers. If true, any trace metadata set on the incoming request will
65	// be added as a linked trace instead of being added as a parent of the
66	// current trace.
67	IsPublicEndpoint bool
68
69	// FormatSpanName holds the function to use for generating the span name
70	// from the information found in the incoming HTTP Request. By default the
71	// name equals the URL Path.
72	FormatSpanName func(*http.Request) string
73
74	// IsHealthEndpoint holds the function to use for determining if the
75	// incoming HTTP request should be considered a health check. This is in
76	// addition to the private isHealthEndpoint func which may also indicate
77	// tracing should be skipped.
78	IsHealthEndpoint func(*http.Request) bool
79}
80
81func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
82	var tags addedTags
83	r, traceEnd := h.startTrace(w, r)
84	defer traceEnd()
85	w, statsEnd := h.startStats(w, r)
86	defer statsEnd(&tags)
87	handler := h.Handler
88	if handler == nil {
89		handler = http.DefaultServeMux
90	}
91	r = r.WithContext(context.WithValue(r.Context(), addedTagsKey{}, &tags))
92	handler.ServeHTTP(w, r)
93}
94
95func (h *Handler) startTrace(w http.ResponseWriter, r *http.Request) (*http.Request, func()) {
96	if h.IsHealthEndpoint != nil && h.IsHealthEndpoint(r) || isHealthEndpoint(r.URL.Path) {
97		return r, func() {}
98	}
99	var name string
100	if h.FormatSpanName == nil {
101		name = spanNameFromURL(r)
102	} else {
103		name = h.FormatSpanName(r)
104	}
105	ctx := r.Context()
106
107	startOpts := h.StartOptions
108	if h.GetStartOptions != nil {
109		startOpts = h.GetStartOptions(r)
110	}
111
112	var span *trace.Span
113	sc, ok := h.extractSpanContext(r)
114	if ok && !h.IsPublicEndpoint {
115		ctx, span = trace.StartSpanWithRemoteParent(ctx, name, sc,
116			trace.WithSampler(startOpts.Sampler),
117			trace.WithSpanKind(trace.SpanKindServer))
118	} else {
119		ctx, span = trace.StartSpan(ctx, name,
120			trace.WithSampler(startOpts.Sampler),
121			trace.WithSpanKind(trace.SpanKindServer),
122		)
123		if ok {
124			span.AddLink(trace.Link{
125				TraceID:    sc.TraceID,
126				SpanID:     sc.SpanID,
127				Type:       trace.LinkTypeParent,
128				Attributes: nil,
129			})
130		}
131	}
132	span.AddAttributes(requestAttrs(r)...)
133	if r.Body == nil {
134		// TODO: Handle cases where ContentLength is not set.
135	} else if r.ContentLength > 0 {
136		span.AddMessageReceiveEvent(0, /* TODO: messageID */
137			r.ContentLength, -1)
138	}
139	return r.WithContext(ctx), span.End
140}
141
142func (h *Handler) extractSpanContext(r *http.Request) (trace.SpanContext, bool) {
143	if h.Propagation == nil {
144		return defaultFormat.SpanContextFromRequest(r)
145	}
146	return h.Propagation.SpanContextFromRequest(r)
147}
148
149func (h *Handler) startStats(w http.ResponseWriter, r *http.Request) (http.ResponseWriter, func(tags *addedTags)) {
150	ctx, _ := tag.New(r.Context(),
151		tag.Upsert(Host, r.Host),
152		tag.Upsert(Path, r.URL.Path),
153		tag.Upsert(Method, r.Method))
154	track := &trackingResponseWriter{
155		start:  time.Now(),
156		ctx:    ctx,
157		writer: w,
158	}
159	if r.Body == nil {
160		// TODO: Handle cases where ContentLength is not set.
161		track.reqSize = -1
162	} else if r.ContentLength > 0 {
163		track.reqSize = r.ContentLength
164	}
165	stats.Record(ctx, ServerRequestCount.M(1))
166	return track.wrappedResponseWriter(), track.end
167}
168
169type trackingResponseWriter struct {
170	ctx        context.Context
171	reqSize    int64
172	respSize   int64
173	start      time.Time
174	statusCode int
175	statusLine string
176	endOnce    sync.Once
177	writer     http.ResponseWriter
178}
179
180// Compile time assertion for ResponseWriter interface
181var _ http.ResponseWriter = (*trackingResponseWriter)(nil)
182
183func (t *trackingResponseWriter) end(tags *addedTags) {
184	t.endOnce.Do(func() {
185		if t.statusCode == 0 {
186			t.statusCode = 200
187		}
188
189		span := trace.FromContext(t.ctx)
190		span.SetStatus(TraceStatus(t.statusCode, t.statusLine))
191		span.AddAttributes(trace.Int64Attribute(StatusCodeAttribute, int64(t.statusCode)))
192
193		m := []stats.Measurement{
194			ServerLatency.M(float64(time.Since(t.start)) / float64(time.Millisecond)),
195			ServerResponseBytes.M(t.respSize),
196		}
197		if t.reqSize >= 0 {
198			m = append(m, ServerRequestBytes.M(t.reqSize))
199		}
200		allTags := make([]tag.Mutator, len(tags.t)+1)
201		allTags[0] = tag.Upsert(StatusCode, strconv.Itoa(t.statusCode))
202		copy(allTags[1:], tags.t)
203		stats.RecordWithTags(t.ctx, allTags, m...)
204	})
205}
206
207func (t *trackingResponseWriter) Header() http.Header {
208	return t.writer.Header()
209}
210
211func (t *trackingResponseWriter) Write(data []byte) (int, error) {
212	n, err := t.writer.Write(data)
213	t.respSize += int64(n)
214	// Add message event for request bytes sent.
215	span := trace.FromContext(t.ctx)
216	span.AddMessageSendEvent(0 /* TODO: messageID */, int64(n), -1)
217	return n, err
218}
219
220func (t *trackingResponseWriter) WriteHeader(statusCode int) {
221	t.writer.WriteHeader(statusCode)
222	t.statusCode = statusCode
223	t.statusLine = http.StatusText(t.statusCode)
224}
225
226// wrappedResponseWriter returns a wrapped version of the original
227//  ResponseWriter and only implements the same combination of additional
228// interfaces as the original.
229// This implementation is based on https://github.com/felixge/httpsnoop.
230func (t *trackingResponseWriter) wrappedResponseWriter() http.ResponseWriter {
231	var (
232		hj, i0 = t.writer.(http.Hijacker)
233		cn, i1 = t.writer.(http.CloseNotifier)
234		pu, i2 = t.writer.(http.Pusher)
235		fl, i3 = t.writer.(http.Flusher)
236		rf, i4 = t.writer.(io.ReaderFrom)
237	)
238
239	switch {
240	case !i0 && !i1 && !i2 && !i3 && !i4:
241		return struct {
242			http.ResponseWriter
243		}{t}
244	case !i0 && !i1 && !i2 && !i3 && i4:
245		return struct {
246			http.ResponseWriter
247			io.ReaderFrom
248		}{t, rf}
249	case !i0 && !i1 && !i2 && i3 && !i4:
250		return struct {
251			http.ResponseWriter
252			http.Flusher
253		}{t, fl}
254	case !i0 && !i1 && !i2 && i3 && i4:
255		return struct {
256			http.ResponseWriter
257			http.Flusher
258			io.ReaderFrom
259		}{t, fl, rf}
260	case !i0 && !i1 && i2 && !i3 && !i4:
261		return struct {
262			http.ResponseWriter
263			http.Pusher
264		}{t, pu}
265	case !i0 && !i1 && i2 && !i3 && i4:
266		return struct {
267			http.ResponseWriter
268			http.Pusher
269			io.ReaderFrom
270		}{t, pu, rf}
271	case !i0 && !i1 && i2 && i3 && !i4:
272		return struct {
273			http.ResponseWriter
274			http.Pusher
275			http.Flusher
276		}{t, pu, fl}
277	case !i0 && !i1 && i2 && i3 && i4:
278		return struct {
279			http.ResponseWriter
280			http.Pusher
281			http.Flusher
282			io.ReaderFrom
283		}{t, pu, fl, rf}
284	case !i0 && i1 && !i2 && !i3 && !i4:
285		return struct {
286			http.ResponseWriter
287			http.CloseNotifier
288		}{t, cn}
289	case !i0 && i1 && !i2 && !i3 && i4:
290		return struct {
291			http.ResponseWriter
292			http.CloseNotifier
293			io.ReaderFrom
294		}{t, cn, rf}
295	case !i0 && i1 && !i2 && i3 && !i4:
296		return struct {
297			http.ResponseWriter
298			http.CloseNotifier
299			http.Flusher
300		}{t, cn, fl}
301	case !i0 && i1 && !i2 && i3 && i4:
302		return struct {
303			http.ResponseWriter
304			http.CloseNotifier
305			http.Flusher
306			io.ReaderFrom
307		}{t, cn, fl, rf}
308	case !i0 && i1 && i2 && !i3 && !i4:
309		return struct {
310			http.ResponseWriter
311			http.CloseNotifier
312			http.Pusher
313		}{t, cn, pu}
314	case !i0 && i1 && i2 && !i3 && i4:
315		return struct {
316			http.ResponseWriter
317			http.CloseNotifier
318			http.Pusher
319			io.ReaderFrom
320		}{t, cn, pu, rf}
321	case !i0 && i1 && i2 && i3 && !i4:
322		return struct {
323			http.ResponseWriter
324			http.CloseNotifier
325			http.Pusher
326			http.Flusher
327		}{t, cn, pu, fl}
328	case !i0 && i1 && i2 && i3 && i4:
329		return struct {
330			http.ResponseWriter
331			http.CloseNotifier
332			http.Pusher
333			http.Flusher
334			io.ReaderFrom
335		}{t, cn, pu, fl, rf}
336	case i0 && !i1 && !i2 && !i3 && !i4:
337		return struct {
338			http.ResponseWriter
339			http.Hijacker
340		}{t, hj}
341	case i0 && !i1 && !i2 && !i3 && i4:
342		return struct {
343			http.ResponseWriter
344			http.Hijacker
345			io.ReaderFrom
346		}{t, hj, rf}
347	case i0 && !i1 && !i2 && i3 && !i4:
348		return struct {
349			http.ResponseWriter
350			http.Hijacker
351			http.Flusher
352		}{t, hj, fl}
353	case i0 && !i1 && !i2 && i3 && i4:
354		return struct {
355			http.ResponseWriter
356			http.Hijacker
357			http.Flusher
358			io.ReaderFrom
359		}{t, hj, fl, rf}
360	case i0 && !i1 && i2 && !i3 && !i4:
361		return struct {
362			http.ResponseWriter
363			http.Hijacker
364			http.Pusher
365		}{t, hj, pu}
366	case i0 && !i1 && i2 && !i3 && i4:
367		return struct {
368			http.ResponseWriter
369			http.Hijacker
370			http.Pusher
371			io.ReaderFrom
372		}{t, hj, pu, rf}
373	case i0 && !i1 && i2 && i3 && !i4:
374		return struct {
375			http.ResponseWriter
376			http.Hijacker
377			http.Pusher
378			http.Flusher
379		}{t, hj, pu, fl}
380	case i0 && !i1 && i2 && i3 && i4:
381		return struct {
382			http.ResponseWriter
383			http.Hijacker
384			http.Pusher
385			http.Flusher
386			io.ReaderFrom
387		}{t, hj, pu, fl, rf}
388	case i0 && i1 && !i2 && !i3 && !i4:
389		return struct {
390			http.ResponseWriter
391			http.Hijacker
392			http.CloseNotifier
393		}{t, hj, cn}
394	case i0 && i1 && !i2 && !i3 && i4:
395		return struct {
396			http.ResponseWriter
397			http.Hijacker
398			http.CloseNotifier
399			io.ReaderFrom
400		}{t, hj, cn, rf}
401	case i0 && i1 && !i2 && i3 && !i4:
402		return struct {
403			http.ResponseWriter
404			http.Hijacker
405			http.CloseNotifier
406			http.Flusher
407		}{t, hj, cn, fl}
408	case i0 && i1 && !i2 && i3 && i4:
409		return struct {
410			http.ResponseWriter
411			http.Hijacker
412			http.CloseNotifier
413			http.Flusher
414			io.ReaderFrom
415		}{t, hj, cn, fl, rf}
416	case i0 && i1 && i2 && !i3 && !i4:
417		return struct {
418			http.ResponseWriter
419			http.Hijacker
420			http.CloseNotifier
421			http.Pusher
422		}{t, hj, cn, pu}
423	case i0 && i1 && i2 && !i3 && i4:
424		return struct {
425			http.ResponseWriter
426			http.Hijacker
427			http.CloseNotifier
428			http.Pusher
429			io.ReaderFrom
430		}{t, hj, cn, pu, rf}
431	case i0 && i1 && i2 && i3 && !i4:
432		return struct {
433			http.ResponseWriter
434			http.Hijacker
435			http.CloseNotifier
436			http.Pusher
437			http.Flusher
438		}{t, hj, cn, pu, fl}
439	case i0 && i1 && i2 && i3 && i4:
440		return struct {
441			http.ResponseWriter
442			http.Hijacker
443			http.CloseNotifier
444			http.Pusher
445			http.Flusher
446			io.ReaderFrom
447		}{t, hj, cn, pu, fl, rf}
448	default:
449		return struct {
450			http.ResponseWriter
451		}{t}
452	}
453}
454