summaryrefslogtreecommitdiff
path: root/srv/src/http/middleware.go
blob: 02d156b14604a08e979e3217d1df95395f056083 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
package http

import (
	"net"
	"net/http"
	"time"

	"github.com/mediocregopher/blog.mediocregopher.com/srv/http/apiutil"
	"github.com/mediocregopher/mediocre-go-lib/v2/mctx"
	"github.com/mediocregopher/mediocre-go-lib/v2/mlog"
)

type middleware func(http.Handler) http.Handler

func applyMiddlewares(h http.Handler, middlewares ...middleware) http.Handler {
	for _, m := range middlewares {
		h = m(h)
	}
	return h
}

func addResponseHeadersMiddleware(headers map[string]string) middleware {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
			for k, v := range headers {
				rw.Header().Set(k, v)
			}
			h.ServeHTTP(rw, r)
		})
	}
}

func setLoggerMiddleware(logger *mlog.Logger) middleware {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {

			type logCtxKey string

			ip, _, _ := net.SplitHostPort(r.RemoteAddr)

			ctx := r.Context()
			ctx = mctx.Annotate(ctx,
				logCtxKey("remote_ip"), ip,
				logCtxKey("url"), r.URL,
				logCtxKey("method"), r.Method,
			)

			r = r.WithContext(ctx)
			r = apiutil.SetRequestLogger(r, logger)
			h.ServeHTTP(rw, r)
		})
	}
}

type logResponseWriter struct {
	http.ResponseWriter
	http.Hijacker
	statusCode int
}

func newLogResponseWriter(rw http.ResponseWriter) *logResponseWriter {
	h, _ := rw.(http.Hijacker)
	return &logResponseWriter{
		ResponseWriter: rw,
		Hijacker:       h,
		statusCode:     200,
	}
}

func (lrw *logResponseWriter) WriteHeader(statusCode int) {
	lrw.statusCode = statusCode
	lrw.ResponseWriter.WriteHeader(statusCode)
}

func logReqMiddleware(h http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {

		lrw := newLogResponseWriter(rw)

		started := time.Now()
		h.ServeHTTP(lrw, r)
		took := time.Since(started)

		type logCtxKey string

		ctx := r.Context()
		ctx = mctx.Annotate(ctx,
			logCtxKey("took"), took.String(),
			logCtxKey("response_code"), lrw.statusCode,
		)

		apiutil.GetRequestLogger(r).Info(ctx, "handled HTTP request")
	})
}

func disallowGetMiddleware(h http.Handler) http.Handler {
	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {

		// we allow websockets to be GETs because, well, they must be
		if r.Method != "GET" || r.Header.Get("Upgrade") == "websocket" {
			h.ServeHTTP(rw, r)
			return
		}

		apiutil.GetRequestLogger(r).WarnString(r.Context(), "method not allowed")
		rw.WriteHeader(405)
	})
}