context.go 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. // Copyright 2012 The Gorilla Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package context
  5. import (
  6. "net/http"
  7. "sync"
  8. "time"
  9. )
  10. var (
  11. mutex sync.RWMutex
  12. data = make(map[*http.Request]map[interface{}]interface{})
  13. datat = make(map[*http.Request]int64)
  14. )
  15. // Set stores a value for a given key in a given request.
  16. func Set(r *http.Request, key, val interface{}) {
  17. mutex.Lock()
  18. if data[r] == nil {
  19. data[r] = make(map[interface{}]interface{})
  20. datat[r] = time.Now().Unix()
  21. }
  22. data[r][key] = val
  23. mutex.Unlock()
  24. }
  25. // Get returns a value stored for a given key in a given request.
  26. func Get(r *http.Request, key interface{}) interface{} {
  27. mutex.RLock()
  28. if ctx := data[r]; ctx != nil {
  29. value := ctx[key]
  30. mutex.RUnlock()
  31. return value
  32. }
  33. mutex.RUnlock()
  34. return nil
  35. }
  36. // GetOk returns stored value and presence state like multi-value return of map access.
  37. func GetOk(r *http.Request, key interface{}) (interface{}, bool) {
  38. mutex.RLock()
  39. if _, ok := data[r]; ok {
  40. value, ok := data[r][key]
  41. mutex.RUnlock()
  42. return value, ok
  43. }
  44. mutex.RUnlock()
  45. return nil, false
  46. }
  47. // GetAll returns all stored values for the request as a map. Nil is returned for invalid requests.
  48. func GetAll(r *http.Request) map[interface{}]interface{} {
  49. mutex.RLock()
  50. if context, ok := data[r]; ok {
  51. result := make(map[interface{}]interface{}, len(context))
  52. for k, v := range context {
  53. result[k] = v
  54. }
  55. mutex.RUnlock()
  56. return result
  57. }
  58. mutex.RUnlock()
  59. return nil
  60. }
  61. // GetAllOk returns all stored values for the request as a map and a boolean value that indicates if
  62. // the request was registered.
  63. func GetAllOk(r *http.Request) (map[interface{}]interface{}, bool) {
  64. mutex.RLock()
  65. context, ok := data[r]
  66. result := make(map[interface{}]interface{}, len(context))
  67. for k, v := range context {
  68. result[k] = v
  69. }
  70. mutex.RUnlock()
  71. return result, ok
  72. }
  73. // Delete removes a value stored for a given key in a given request.
  74. func Delete(r *http.Request, key interface{}) {
  75. mutex.Lock()
  76. if data[r] != nil {
  77. delete(data[r], key)
  78. }
  79. mutex.Unlock()
  80. }
  81. // Clear removes all values stored for a given request.
  82. //
  83. // This is usually called by a handler wrapper to clean up request
  84. // variables at the end of a request lifetime. See ClearHandler().
  85. func Clear(r *http.Request) {
  86. mutex.Lock()
  87. clear(r)
  88. mutex.Unlock()
  89. }
  90. // clear is Clear without the lock.
  91. func clear(r *http.Request) {
  92. delete(data, r)
  93. delete(datat, r)
  94. }
  95. // Purge removes request data stored for longer than maxAge, in seconds.
  96. // It returns the amount of requests removed.
  97. //
  98. // If maxAge <= 0, all request data is removed.
  99. //
  100. // This is only used for sanity check: in case context cleaning was not
  101. // properly set some request data can be kept forever, consuming an increasing
  102. // amount of memory. In case this is detected, Purge() must be called
  103. // periodically until the problem is fixed.
  104. func Purge(maxAge int) int {
  105. mutex.Lock()
  106. count := 0
  107. if maxAge <= 0 {
  108. count = len(data)
  109. data = make(map[*http.Request]map[interface{}]interface{})
  110. datat = make(map[*http.Request]int64)
  111. } else {
  112. min := time.Now().Unix() - int64(maxAge)
  113. for r := range data {
  114. if datat[r] < min {
  115. clear(r)
  116. count++
  117. }
  118. }
  119. }
  120. mutex.Unlock()
  121. return count
  122. }
  123. // ClearHandler wraps an http.Handler and clears request values at the end
  124. // of a request lifetime.
  125. func ClearHandler(h http.Handler) http.Handler {
  126. return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  127. defer Clear(r)
  128. h.ServeHTTP(w, r)
  129. })
  130. }