|
|
|
@ -1,10 +1,14 @@ |
|
|
|
|
package http |
|
|
|
|
|
|
|
|
|
import ( |
|
|
|
|
"bytes" |
|
|
|
|
"net" |
|
|
|
|
"net/http" |
|
|
|
|
"path/filepath" |
|
|
|
|
"sync" |
|
|
|
|
"time" |
|
|
|
|
|
|
|
|
|
lru "github.com/hashicorp/golang-lru" |
|
|
|
|
"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" |
|
|
|
@ -52,33 +56,33 @@ func setLoggerMiddleware(logger *mlog.Logger) middleware { |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
type logResponseWriter struct { |
|
|
|
|
type wrappedResponseWriter struct { |
|
|
|
|
http.ResponseWriter |
|
|
|
|
http.Hijacker |
|
|
|
|
statusCode int |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func newLogResponseWriter(rw http.ResponseWriter) *logResponseWriter { |
|
|
|
|
func newWrappedResponseWriter(rw http.ResponseWriter) *wrappedResponseWriter { |
|
|
|
|
h, _ := rw.(http.Hijacker) |
|
|
|
|
return &logResponseWriter{ |
|
|
|
|
return &wrappedResponseWriter{ |
|
|
|
|
ResponseWriter: rw, |
|
|
|
|
Hijacker: h, |
|
|
|
|
statusCode: 200, |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func (lrw *logResponseWriter) WriteHeader(statusCode int) { |
|
|
|
|
lrw.statusCode = statusCode |
|
|
|
|
lrw.ResponseWriter.WriteHeader(statusCode) |
|
|
|
|
func (rw *wrappedResponseWriter) WriteHeader(statusCode int) { |
|
|
|
|
rw.statusCode = statusCode |
|
|
|
|
rw.ResponseWriter.WriteHeader(statusCode) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func logReqMiddleware(h http.Handler) http.Handler { |
|
|
|
|
return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { |
|
|
|
|
|
|
|
|
|
lrw := newLogResponseWriter(rw) |
|
|
|
|
wrw := newWrappedResponseWriter(rw) |
|
|
|
|
|
|
|
|
|
started := time.Now() |
|
|
|
|
h.ServeHTTP(lrw, r) |
|
|
|
|
h.ServeHTTP(wrw, r) |
|
|
|
|
took := time.Since(started) |
|
|
|
|
|
|
|
|
|
type logCtxKey string |
|
|
|
@ -86,7 +90,7 @@ func logReqMiddleware(h http.Handler) http.Handler { |
|
|
|
|
ctx := r.Context() |
|
|
|
|
ctx = mctx.Annotate(ctx, |
|
|
|
|
logCtxKey("took"), took.String(), |
|
|
|
|
logCtxKey("response_code"), lrw.statusCode, |
|
|
|
|
logCtxKey("response_code"), wrw.statusCode, |
|
|
|
|
) |
|
|
|
|
|
|
|
|
|
apiutil.GetRequestLogger(r).Info(ctx, "handled HTTP request") |
|
|
|
@ -106,3 +110,80 @@ func disallowGetMiddleware(h http.Handler) http.Handler { |
|
|
|
|
rw.WriteHeader(405) |
|
|
|
|
}) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
type cacheResponseWriter struct { |
|
|
|
|
*wrappedResponseWriter |
|
|
|
|
buf *bytes.Buffer |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func newCacheResponseWriter(rw http.ResponseWriter) *cacheResponseWriter { |
|
|
|
|
return &cacheResponseWriter{ |
|
|
|
|
wrappedResponseWriter: newWrappedResponseWriter(rw), |
|
|
|
|
buf: new(bytes.Buffer), |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func (rw *cacheResponseWriter) Write(b []byte) (int, error) { |
|
|
|
|
if _, err := rw.buf.Write(b); err != nil { |
|
|
|
|
panic(err) |
|
|
|
|
} |
|
|
|
|
return rw.wrappedResponseWriter.Write(b) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func cacheMiddleware(cache *lru.Cache) middleware { |
|
|
|
|
|
|
|
|
|
type entry struct { |
|
|
|
|
body []byte |
|
|
|
|
createdAt time.Time |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
pool := sync.Pool{ |
|
|
|
|
New: func() interface{} { return new(bytes.Reader) }, |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
return func(h http.Handler) http.Handler { |
|
|
|
|
return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { |
|
|
|
|
|
|
|
|
|
id := r.URL.RequestURI() |
|
|
|
|
|
|
|
|
|
if val, ok := cache.Get(id); ok { |
|
|
|
|
|
|
|
|
|
entry := val.(entry) |
|
|
|
|
|
|
|
|
|
reader := pool.Get().(*bytes.Reader) |
|
|
|
|
defer pool.Put(reader) |
|
|
|
|
|
|
|
|
|
reader.Reset(entry.body) |
|
|
|
|
|
|
|
|
|
http.ServeContent( |
|
|
|
|
rw, r, filepath.Base(r.URL.Path), entry.createdAt, reader, |
|
|
|
|
) |
|
|
|
|
return |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
cacheRW := newCacheResponseWriter(rw) |
|
|
|
|
h.ServeHTTP(cacheRW, r) |
|
|
|
|
|
|
|
|
|
if cacheRW.statusCode == 200 { |
|
|
|
|
cache.Add(id, entry{ |
|
|
|
|
body: cacheRW.buf.Bytes(), |
|
|
|
|
createdAt: time.Now(), |
|
|
|
|
}) |
|
|
|
|
} |
|
|
|
|
}) |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func purgeCacheOnOKMiddleware(cache *lru.Cache) middleware { |
|
|
|
|
return func(h http.Handler) http.Handler { |
|
|
|
|
return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { |
|
|
|
|
|
|
|
|
|
wrw := newWrappedResponseWriter(rw) |
|
|
|
|
h.ServeHTTP(wrw, r) |
|
|
|
|
|
|
|
|
|
if wrw.statusCode == 200 { |
|
|
|
|
cache.Purge() |
|
|
|
|
} |
|
|
|
|
}) |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|