1// Copyright 2015 The Prometheus Authors
2// Licensed under the Apache License, Version 2.0 (the "License");
3// you may not use this file except in compliance with the License.
4// You may obtain a copy of the License at
5//
6// http://www.apache.org/licenses/LICENSE-2.0
7//
8// Unless required by applicable law or agreed to in writing, software
9// distributed under the License is distributed on an "AS IS" BASIS,
10// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11// See the License for the specific language governing permissions and
12// limitations under the License.
13
14package route
15
16import (
17	"context"
18	"net/http"
19
20	"github.com/julienschmidt/httprouter"
21)
22
23type param string
24
25// Param returns param p for the context, or the empty string when
26// param does not exist in context.
27func Param(ctx context.Context, p string) string {
28	if v := ctx.Value(param(p)); v != nil {
29		return v.(string)
30	}
31	return ""
32}
33
34// WithParam returns a new context with param p set to v.
35func WithParam(ctx context.Context, p, v string) context.Context {
36	return context.WithValue(ctx, param(p), v)
37}
38
39// Router wraps httprouter.Router and adds support for prefixed sub-routers,
40// per-request context injections and instrumentation.
41type Router struct {
42	rtr    *httprouter.Router
43	prefix string
44	instrh func(handlerName string, handler http.HandlerFunc) http.HandlerFunc
45}
46
47// New returns a new Router.
48func New() *Router {
49	return &Router{
50		rtr: httprouter.New(),
51	}
52}
53
54// WithInstrumentation returns a router with instrumentation support.
55func (r *Router) WithInstrumentation(instrh func(handlerName string, handler http.HandlerFunc) http.HandlerFunc) *Router {
56	if r.instrh != nil {
57		newInstrh := instrh
58		instrh = func(handlerName string, handler http.HandlerFunc) http.HandlerFunc {
59			return newInstrh(handlerName, r.instrh(handlerName, handler))
60		}
61	}
62	return &Router{rtr: r.rtr, prefix: r.prefix, instrh: instrh}
63}
64
65// WithPrefix returns a router that prefixes all registered routes with prefix.
66func (r *Router) WithPrefix(prefix string) *Router {
67	return &Router{rtr: r.rtr, prefix: r.prefix + prefix, instrh: r.instrh}
68}
69
70// handle turns a HandlerFunc into an httprouter.Handle.
71func (r *Router) handle(handlerName string, h http.HandlerFunc) httprouter.Handle {
72	if r.instrh != nil {
73		// This needs to be outside the closure to avoid data race when reading and writing to 'h'.
74		h = r.instrh(handlerName, h)
75	}
76	return func(w http.ResponseWriter, req *http.Request, params httprouter.Params) {
77		ctx, cancel := context.WithCancel(req.Context())
78		defer cancel()
79
80		for _, p := range params {
81			ctx = context.WithValue(ctx, param(p.Key), p.Value)
82		}
83		h(w, req.WithContext(ctx))
84	}
85}
86
87// Get registers a new GET route.
88func (r *Router) Get(path string, h http.HandlerFunc) {
89	r.rtr.GET(r.prefix+path, r.handle(path, h))
90}
91
92// Options registers a new OPTIONS route.
93func (r *Router) Options(path string, h http.HandlerFunc) {
94	r.rtr.OPTIONS(r.prefix+path, r.handle(path, h))
95}
96
97// Del registers a new DELETE route.
98func (r *Router) Del(path string, h http.HandlerFunc) {
99	r.rtr.DELETE(r.prefix+path, r.handle(path, h))
100}
101
102// Put registers a new PUT route.
103func (r *Router) Put(path string, h http.HandlerFunc) {
104	r.rtr.PUT(r.prefix+path, r.handle(path, h))
105}
106
107// Post registers a new POST route.
108func (r *Router) Post(path string, h http.HandlerFunc) {
109	r.rtr.POST(r.prefix+path, r.handle(path, h))
110}
111
112// Redirect takes an absolute path and sends an internal HTTP redirect for it,
113// prefixed by the router's path prefix. Note that this method does not include
114// functionality for handling relative paths or full URL redirects.
115func (r *Router) Redirect(w http.ResponseWriter, req *http.Request, path string, code int) {
116	http.Redirect(w, req, r.prefix+path, code)
117}
118
119// ServeHTTP implements http.Handler.
120func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
121	r.rtr.ServeHTTP(w, req)
122}
123
124// FileServe returns a new http.HandlerFunc that serves files from dir.
125// Using routes must provide the *filepath parameter.
126func FileServe(dir string) http.HandlerFunc {
127	fs := http.FileServer(http.Dir(dir))
128
129	return func(w http.ResponseWriter, r *http.Request) {
130		r.URL.Path = Param(r.Context(), "filepath")
131		fs.ServeHTTP(w, r)
132	}
133}
134