context.go 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230
  1. // Copyright 2014 beego Author. All Rights Reserved.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. // Package context provide the context utils
  15. // Usage:
  16. //
  17. // import "github.com/astaxie/beego/context"
  18. //
  19. // ctx := context.Context{Request:req,ResponseWriter:rw}
  20. //
  21. // more docs http://beego.me/docs/module/context.md
  22. package context
  23. import (
  24. "bufio"
  25. "crypto/hmac"
  26. "crypto/sha1"
  27. "encoding/base64"
  28. "errors"
  29. "fmt"
  30. "net"
  31. "net/http"
  32. "strconv"
  33. "strings"
  34. "time"
  35. "github.com/astaxie/beego/utils"
  36. )
  37. // NewContext return the Context with Input and Output
  38. func NewContext() *Context {
  39. return &Context{
  40. Input: NewInput(),
  41. Output: NewOutput(),
  42. }
  43. }
  44. // Context Http request context struct including BeegoInput, BeegoOutput, http.Request and http.ResponseWriter.
  45. // BeegoInput and BeegoOutput provides some api to operate request and response more easily.
  46. type Context struct {
  47. Input *BeegoInput
  48. Output *BeegoOutput
  49. Request *http.Request
  50. ResponseWriter *Response
  51. _xsrfToken string
  52. }
  53. // Reset init Context, BeegoInput and BeegoOutput
  54. func (ctx *Context) Reset(rw http.ResponseWriter, r *http.Request) {
  55. ctx.Request = r
  56. if ctx.ResponseWriter == nil {
  57. ctx.ResponseWriter = &Response{}
  58. }
  59. ctx.ResponseWriter.reset(rw)
  60. ctx.Input.Reset(ctx)
  61. ctx.Output.Reset(ctx)
  62. ctx._xsrfToken = ""
  63. }
  64. // Redirect does redirection to localurl with http header status code.
  65. func (ctx *Context) Redirect(status int, localurl string) {
  66. http.Redirect(ctx.ResponseWriter, ctx.Request, localurl, status)
  67. }
  68. // Abort stops this request.
  69. // if beego.ErrorMaps exists, panic body.
  70. func (ctx *Context) Abort(status int, body string) {
  71. ctx.Output.SetStatus(status)
  72. panic(body)
  73. }
  74. // WriteString Write string to response body.
  75. // it sends response body.
  76. func (ctx *Context) WriteString(content string) {
  77. ctx.ResponseWriter.Write([]byte(content))
  78. }
  79. // GetCookie Get cookie from request by a given key.
  80. // It's alias of BeegoInput.Cookie.
  81. func (ctx *Context) GetCookie(key string) string {
  82. return ctx.Input.Cookie(key)
  83. }
  84. // SetCookie Set cookie for response.
  85. // It's alias of BeegoOutput.Cookie.
  86. func (ctx *Context) SetCookie(name string, value string, others ...interface{}) {
  87. ctx.Output.Cookie(name, value, others...)
  88. }
  89. // GetSecureCookie Get secure cookie from request by a given key.
  90. func (ctx *Context) GetSecureCookie(Secret, key string) (string, bool) {
  91. val := ctx.Input.Cookie(key)
  92. if val == "" {
  93. return "", false
  94. }
  95. parts := strings.SplitN(val, "|", 3)
  96. if len(parts) != 3 {
  97. return "", false
  98. }
  99. vs := parts[0]
  100. timestamp := parts[1]
  101. sig := parts[2]
  102. h := hmac.New(sha1.New, []byte(Secret))
  103. fmt.Fprintf(h, "%s%s", vs, timestamp)
  104. if fmt.Sprintf("%02x", h.Sum(nil)) != sig {
  105. return "", false
  106. }
  107. res, _ := base64.URLEncoding.DecodeString(vs)
  108. return string(res), true
  109. }
  110. // SetSecureCookie Set Secure cookie for response.
  111. func (ctx *Context) SetSecureCookie(Secret, name, value string, others ...interface{}) {
  112. vs := base64.URLEncoding.EncodeToString([]byte(value))
  113. timestamp := strconv.FormatInt(time.Now().UnixNano(), 10)
  114. h := hmac.New(sha1.New, []byte(Secret))
  115. fmt.Fprintf(h, "%s%s", vs, timestamp)
  116. sig := fmt.Sprintf("%02x", h.Sum(nil))
  117. cookie := strings.Join([]string{vs, timestamp, sig}, "|")
  118. ctx.Output.Cookie(name, cookie, others...)
  119. }
  120. // XSRFToken creates a xsrf token string and returns.
  121. func (ctx *Context) XSRFToken(key string, expire int64) string {
  122. if ctx._xsrfToken == "" {
  123. token, ok := ctx.GetSecureCookie(key, "_xsrf")
  124. if !ok {
  125. token = string(utils.RandomCreateBytes(32))
  126. ctx.SetSecureCookie(key, "_xsrf", token, expire)
  127. }
  128. ctx._xsrfToken = token
  129. }
  130. return ctx._xsrfToken
  131. }
  132. // CheckXSRFCookie checks xsrf token in this request is valid or not.
  133. // the token can provided in request header "X-Xsrftoken" and "X-CsrfToken"
  134. // or in form field value named as "_xsrf".
  135. func (ctx *Context) CheckXSRFCookie() bool {
  136. token := ctx.Input.Query("_xsrf")
  137. if token == "" {
  138. token = ctx.Request.Header.Get("X-Xsrftoken")
  139. }
  140. if token == "" {
  141. token = ctx.Request.Header.Get("X-Csrftoken")
  142. }
  143. if token == "" {
  144. ctx.Abort(403, "'_xsrf' argument missing from POST")
  145. return false
  146. }
  147. if ctx._xsrfToken != token {
  148. ctx.Abort(403, "XSRF cookie does not match POST argument")
  149. return false
  150. }
  151. return true
  152. }
  153. //Response is a wrapper for the http.ResponseWriter
  154. //started set to true if response was written to then don't execute other handler
  155. type Response struct {
  156. http.ResponseWriter
  157. Started bool
  158. Status int
  159. }
  160. func (r *Response) reset(rw http.ResponseWriter) {
  161. r.ResponseWriter = rw
  162. r.Status = 0
  163. r.Started = false
  164. }
  165. // Write writes the data to the connection as part of an HTTP reply,
  166. // and sets `started` to true.
  167. // started means the response has sent out.
  168. func (r *Response) Write(p []byte) (int, error) {
  169. r.Started = true
  170. return r.ResponseWriter.Write(p)
  171. }
  172. // WriteHeader sends an HTTP response header with status code,
  173. // and sets `started` to true.
  174. func (r *Response) WriteHeader(code int) {
  175. if r.Status > 0 {
  176. //prevent multiple response.WriteHeader calls
  177. return
  178. }
  179. r.Status = code
  180. r.Started = true
  181. r.ResponseWriter.WriteHeader(code)
  182. }
  183. // Hijack hijacker for http
  184. func (r *Response) Hijack() (net.Conn, *bufio.ReadWriter, error) {
  185. hj, ok := r.ResponseWriter.(http.Hijacker)
  186. if !ok {
  187. return nil, nil, errors.New("webserver doesn't support hijacking")
  188. }
  189. return hj.Hijack()
  190. }
  191. // Flush http.Flusher
  192. func (r *Response) Flush() {
  193. if f, ok := r.ResponseWriter.(http.Flusher); ok {
  194. f.Flush()
  195. }
  196. }
  197. // CloseNotify http.CloseNotifier
  198. func (r *Response) CloseNotify() <-chan bool {
  199. if cn, ok := r.ResponseWriter.(http.CloseNotifier); ok {
  200. return cn.CloseNotify()
  201. }
  202. return nil
  203. }