internal_response_writer.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121
  1. package newrelic
  2. import (
  3. "bufio"
  4. "io"
  5. "net"
  6. "net/http"
  7. )
  8. const (
  9. hasC = 1 << iota // CloseNotifier
  10. hasF // Flusher
  11. hasH // Hijacker
  12. hasR // ReaderFrom
  13. )
  14. type wrap struct{ *txn }
  15. type wrapR struct{ *txn }
  16. type wrapH struct{ *txn }
  17. type wrapHR struct{ *txn }
  18. type wrapF struct{ *txn }
  19. type wrapFR struct{ *txn }
  20. type wrapFH struct{ *txn }
  21. type wrapFHR struct{ *txn }
  22. type wrapC struct{ *txn }
  23. type wrapCR struct{ *txn }
  24. type wrapCH struct{ *txn }
  25. type wrapCHR struct{ *txn }
  26. type wrapCF struct{ *txn }
  27. type wrapCFR struct{ *txn }
  28. type wrapCFH struct{ *txn }
  29. type wrapCFHR struct{ *txn }
  30. func (x wrapC) CloseNotify() <-chan bool { return x.W.(http.CloseNotifier).CloseNotify() }
  31. func (x wrapCR) CloseNotify() <-chan bool { return x.W.(http.CloseNotifier).CloseNotify() }
  32. func (x wrapCH) CloseNotify() <-chan bool { return x.W.(http.CloseNotifier).CloseNotify() }
  33. func (x wrapCHR) CloseNotify() <-chan bool { return x.W.(http.CloseNotifier).CloseNotify() }
  34. func (x wrapCF) CloseNotify() <-chan bool { return x.W.(http.CloseNotifier).CloseNotify() }
  35. func (x wrapCFR) CloseNotify() <-chan bool { return x.W.(http.CloseNotifier).CloseNotify() }
  36. func (x wrapCFH) CloseNotify() <-chan bool { return x.W.(http.CloseNotifier).CloseNotify() }
  37. func (x wrapCFHR) CloseNotify() <-chan bool { return x.W.(http.CloseNotifier).CloseNotify() }
  38. func (x wrapF) Flush() { x.W.(http.Flusher).Flush() }
  39. func (x wrapFR) Flush() { x.W.(http.Flusher).Flush() }
  40. func (x wrapFH) Flush() { x.W.(http.Flusher).Flush() }
  41. func (x wrapFHR) Flush() { x.W.(http.Flusher).Flush() }
  42. func (x wrapCF) Flush() { x.W.(http.Flusher).Flush() }
  43. func (x wrapCFR) Flush() { x.W.(http.Flusher).Flush() }
  44. func (x wrapCFH) Flush() { x.W.(http.Flusher).Flush() }
  45. func (x wrapCFHR) Flush() { x.W.(http.Flusher).Flush() }
  46. func (x wrapH) Hijack() (net.Conn, *bufio.ReadWriter, error) { return x.W.(http.Hijacker).Hijack() }
  47. func (x wrapHR) Hijack() (net.Conn, *bufio.ReadWriter, error) { return x.W.(http.Hijacker).Hijack() }
  48. func (x wrapFH) Hijack() (net.Conn, *bufio.ReadWriter, error) { return x.W.(http.Hijacker).Hijack() }
  49. func (x wrapFHR) Hijack() (net.Conn, *bufio.ReadWriter, error) { return x.W.(http.Hijacker).Hijack() }
  50. func (x wrapCH) Hijack() (net.Conn, *bufio.ReadWriter, error) { return x.W.(http.Hijacker).Hijack() }
  51. func (x wrapCHR) Hijack() (net.Conn, *bufio.ReadWriter, error) { return x.W.(http.Hijacker).Hijack() }
  52. func (x wrapCFH) Hijack() (net.Conn, *bufio.ReadWriter, error) { return x.W.(http.Hijacker).Hijack() }
  53. func (x wrapCFHR) Hijack() (net.Conn, *bufio.ReadWriter, error) { return x.W.(http.Hijacker).Hijack() }
  54. func (x wrapR) ReadFrom(r io.Reader) (int64, error) { return x.W.(io.ReaderFrom).ReadFrom(r) }
  55. func (x wrapHR) ReadFrom(r io.Reader) (int64, error) { return x.W.(io.ReaderFrom).ReadFrom(r) }
  56. func (x wrapFR) ReadFrom(r io.Reader) (int64, error) { return x.W.(io.ReaderFrom).ReadFrom(r) }
  57. func (x wrapFHR) ReadFrom(r io.Reader) (int64, error) { return x.W.(io.ReaderFrom).ReadFrom(r) }
  58. func (x wrapCR) ReadFrom(r io.Reader) (int64, error) { return x.W.(io.ReaderFrom).ReadFrom(r) }
  59. func (x wrapCHR) ReadFrom(r io.Reader) (int64, error) { return x.W.(io.ReaderFrom).ReadFrom(r) }
  60. func (x wrapCFR) ReadFrom(r io.Reader) (int64, error) { return x.W.(io.ReaderFrom).ReadFrom(r) }
  61. func (x wrapCFHR) ReadFrom(r io.Reader) (int64, error) { return x.W.(io.ReaderFrom).ReadFrom(r) }
  62. func upgradeTxn(txn *txn) Transaction {
  63. x := 0
  64. if _, ok := txn.W.(http.CloseNotifier); ok {
  65. x |= hasC
  66. }
  67. if _, ok := txn.W.(http.Flusher); ok {
  68. x |= hasF
  69. }
  70. if _, ok := txn.W.(http.Hijacker); ok {
  71. x |= hasH
  72. }
  73. if _, ok := txn.W.(io.ReaderFrom); ok {
  74. x |= hasR
  75. }
  76. switch x {
  77. default:
  78. // Wrap the transaction even when there are no methods needed to
  79. // ensure consistent error stack trace depth.
  80. return wrap{txn}
  81. case hasR:
  82. return wrapR{txn}
  83. case hasH:
  84. return wrapH{txn}
  85. case hasH | hasR:
  86. return wrapHR{txn}
  87. case hasF:
  88. return wrapF{txn}
  89. case hasF | hasR:
  90. return wrapFR{txn}
  91. case hasF | hasH:
  92. return wrapFH{txn}
  93. case hasF | hasH | hasR:
  94. return wrapFHR{txn}
  95. case hasC:
  96. return wrapC{txn}
  97. case hasC | hasR:
  98. return wrapCR{txn}
  99. case hasC | hasH:
  100. return wrapCH{txn}
  101. case hasC | hasH | hasR:
  102. return wrapCHR{txn}
  103. case hasC | hasF:
  104. return wrapCF{txn}
  105. case hasC | hasF | hasR:
  106. return wrapCFR{txn}
  107. case hasC | hasF | hasH:
  108. return wrapCFH{txn}
  109. case hasC | hasF | hasH | hasR:
  110. return wrapCFHR{txn}
  111. }
  112. }