middleware.go 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960
  1. package xhandler
  2. import (
  3. "net/http"
  4. "time"
  5. "golang.org/x/net/context"
  6. )
  7. // CloseHandler returns a Handler, cancelling the context when the client
  8. // connection closes unexpectedly.
  9. func CloseHandler(next HandlerC) HandlerC {
  10. return HandlerFuncC(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
  11. // Cancel the context if the client closes the connection
  12. if wcn, ok := w.(http.CloseNotifier); ok {
  13. var cancel context.CancelFunc
  14. ctx, cancel = context.WithCancel(ctx)
  15. defer cancel()
  16. notify := wcn.CloseNotify()
  17. go func() {
  18. select {
  19. case <-notify:
  20. cancel()
  21. case <-ctx.Done():
  22. }
  23. }()
  24. }
  25. next.ServeHTTPC(ctx, w, r)
  26. })
  27. }
  28. // TimeoutHandler returns a Handler which adds a timeout to the context.
  29. //
  30. // Child handlers have the responsability of obeying the context deadline and to return
  31. // an appropriate error (or not) response in case of timeout.
  32. func TimeoutHandler(timeout time.Duration) func(next HandlerC) HandlerC {
  33. return func(next HandlerC) HandlerC {
  34. return HandlerFuncC(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
  35. ctx, _ = context.WithTimeout(ctx, timeout)
  36. next.ServeHTTPC(ctx, w, r)
  37. })
  38. }
  39. }
  40. // If is a special handler that will skip insert the condNext handler only if a condition
  41. // applies at runtime.
  42. func If(cond func(ctx context.Context, w http.ResponseWriter, r *http.Request) bool, condNext func(next HandlerC) HandlerC) func(next HandlerC) HandlerC {
  43. return func(next HandlerC) HandlerC {
  44. return HandlerFuncC(func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
  45. if cond(ctx, w, r) {
  46. condNext(next).ServeHTTPC(ctx, w, r)
  47. } else {
  48. next.ServeHTTPC(ctx, w, r)
  49. }
  50. })
  51. }
  52. }