cryptionhandler.go 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121
  1. package handler
  2. import (
  3. "bytes"
  4. "encoding/base64"
  5. "errors"
  6. "io"
  7. "io/ioutil"
  8. "net/http"
  9. "github.com/tal-tech/go-zero/core/codec"
  10. "github.com/tal-tech/go-zero/core/logx"
  11. )
  12. const maxBytes = 1 << 20 // 1 MiB
  13. var errContentLengthExceeded = errors.New("content length exceeded")
  14. func CryptionHandler(key []byte) func(http.Handler) http.Handler {
  15. return func(next http.Handler) http.Handler {
  16. return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  17. cw := newCryptionResponseWriter(w)
  18. defer cw.flush(key)
  19. if r.ContentLength <= 0 {
  20. next.ServeHTTP(cw, r)
  21. return
  22. }
  23. if err := decryptBody(key, r); err != nil {
  24. w.WriteHeader(http.StatusBadRequest)
  25. return
  26. }
  27. next.ServeHTTP(cw, r)
  28. })
  29. }
  30. }
  31. func decryptBody(key []byte, r *http.Request) error {
  32. if r.ContentLength > maxBytes {
  33. return errContentLengthExceeded
  34. }
  35. var content []byte
  36. var err error
  37. if r.ContentLength > 0 {
  38. content = make([]byte, r.ContentLength)
  39. _, err = io.ReadFull(r.Body, content)
  40. } else {
  41. content, err = ioutil.ReadAll(io.LimitReader(r.Body, maxBytes))
  42. }
  43. if err != nil {
  44. return err
  45. }
  46. content, err = base64.StdEncoding.DecodeString(string(content))
  47. if err != nil {
  48. return err
  49. }
  50. output, err := codec.EcbDecrypt(key, content)
  51. if err != nil {
  52. return err
  53. }
  54. var buf bytes.Buffer
  55. buf.Write(output)
  56. r.Body = ioutil.NopCloser(&buf)
  57. return nil
  58. }
  59. type cryptionResponseWriter struct {
  60. http.ResponseWriter
  61. buf *bytes.Buffer
  62. }
  63. func newCryptionResponseWriter(w http.ResponseWriter) *cryptionResponseWriter {
  64. return &cryptionResponseWriter{
  65. ResponseWriter: w,
  66. buf: new(bytes.Buffer),
  67. }
  68. }
  69. func (w *cryptionResponseWriter) Flush() {
  70. if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
  71. flusher.Flush()
  72. }
  73. }
  74. func (w *cryptionResponseWriter) Header() http.Header {
  75. return w.ResponseWriter.Header()
  76. }
  77. func (w *cryptionResponseWriter) Write(p []byte) (int, error) {
  78. return w.buf.Write(p)
  79. }
  80. func (w *cryptionResponseWriter) WriteHeader(statusCode int) {
  81. w.ResponseWriter.WriteHeader(statusCode)
  82. }
  83. func (w *cryptionResponseWriter) flush(key []byte) {
  84. if w.buf.Len() == 0 {
  85. return
  86. }
  87. content, err := codec.EcbEncrypt(key, w.buf.Bytes())
  88. if err != nil {
  89. w.WriteHeader(http.StatusInternalServerError)
  90. return
  91. }
  92. body := base64.StdEncoding.EncodeToString(content)
  93. if n, err := io.WriteString(w.ResponseWriter, body); err != nil {
  94. logx.Errorf("write response failed, error: %s", err)
  95. } else if n < len(content) {
  96. logx.Errorf("actual bytes: %d, written bytes: %d", len(content), n)
  97. }
  98. }