forked from mirrors/homebox
103 lines
2.3 KiB
Go
103 lines
2.3 KiB
Go
|
package server
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"net/http"
|
||
|
|
|
||
|
|
"github.com/google/uuid"
|
||
|
|
)
|
||
|
|
|
||
|
|
type vkey int
|
||
|
|
|
||
|
|
const (
|
||
|
|
// Key is the key for the server in the request context.
|
||
|
|
key vkey = 1
|
||
|
|
)
|
||
|
|
|
||
|
|
type Values struct {
|
||
|
|
TraceID string
|
||
|
|
}
|
||
|
|
|
||
|
|
func GetTraceID(ctx context.Context) string {
|
||
|
|
v, ok := ctx.Value(key).(Values)
|
||
|
|
if !ok {
|
||
|
|
return ""
|
||
|
|
}
|
||
|
|
return v.TraceID
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Server) toHttpHandler(handler Handler, mw ...Middleware) http.HandlerFunc {
|
||
|
|
handler = wrapMiddleware(mw, handler)
|
||
|
|
|
||
|
|
handler = wrapMiddleware(s.mw, handler)
|
||
|
|
|
||
|
|
return func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
ctx := r.Context()
|
||
|
|
|
||
|
|
// Add the trace ID to the context
|
||
|
|
ctx = context.WithValue(ctx, key, Values{
|
||
|
|
TraceID: uuid.NewString(),
|
||
|
|
})
|
||
|
|
|
||
|
|
err := handler.ServeHTTP(w, r.WithContext(ctx))
|
||
|
|
if err != nil {
|
||
|
|
if IsShutdownError(err) {
|
||
|
|
_ = s.Shutdown("SIGTERM")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Server) handle(method, pattern string, handler Handler, mw ...Middleware) {
|
||
|
|
h := s.toHttpHandler(handler, mw...)
|
||
|
|
|
||
|
|
switch method {
|
||
|
|
case http.MethodGet:
|
||
|
|
s.mux.Get(pattern, h)
|
||
|
|
case http.MethodPost:
|
||
|
|
s.mux.Post(pattern, h)
|
||
|
|
case http.MethodPut:
|
||
|
|
s.mux.Put(pattern, h)
|
||
|
|
case http.MethodDelete:
|
||
|
|
s.mux.Delete(pattern, h)
|
||
|
|
case http.MethodPatch:
|
||
|
|
s.mux.Patch(pattern, h)
|
||
|
|
case http.MethodHead:
|
||
|
|
s.mux.Head(pattern, h)
|
||
|
|
case http.MethodOptions:
|
||
|
|
s.mux.Options(pattern, h)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Server) Get(pattern string, handler Handler, mw ...Middleware) {
|
||
|
|
s.handle(http.MethodGet, pattern, handler, mw...)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Server) Post(pattern string, handler Handler, mw ...Middleware) {
|
||
|
|
s.handle(http.MethodPost, pattern, handler, mw...)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Server) Put(pattern string, handler Handler, mw ...Middleware) {
|
||
|
|
s.handle(http.MethodPut, pattern, handler, mw...)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Server) Delete(pattern string, handler Handler, mw ...Middleware) {
|
||
|
|
s.handle(http.MethodDelete, pattern, handler, mw...)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Server) Patch(pattern string, handler Handler, mw ...Middleware) {
|
||
|
|
s.handle(http.MethodPatch, pattern, handler, mw...)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Server) Head(pattern string, handler Handler, mw ...Middleware) {
|
||
|
|
s.handle(http.MethodHead, pattern, handler, mw...)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Server) Options(pattern string, handler Handler, mw ...Middleware) {
|
||
|
|
s.handle(http.MethodOptions, pattern, handler, mw...)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Server) NotFound(handler Handler) {
|
||
|
|
s.mux.NotFound(s.toHttpHandler(handler))
|
||
|
|
}
|