github.com/lingyao2333/mo-zero@v1.4.1/rest/handler/cryptionhandler.go (about)

     1  package handler
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"encoding/base64"
     7  	"errors"
     8  	"io"
     9  	"net"
    10  	"net/http"
    11  
    12  	"github.com/lingyao2333/mo-zero/core/codec"
    13  	"github.com/lingyao2333/mo-zero/core/logx"
    14  )
    15  
    16  const maxBytes = 1 << 20 // 1 MiB
    17  
    18  var errContentLengthExceeded = errors.New("content length exceeded")
    19  
    20  // CryptionHandler returns a middleware to handle cryption.
    21  func CryptionHandler(key []byte) func(http.Handler) http.Handler {
    22  	return func(next http.Handler) http.Handler {
    23  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    24  			cw := newCryptionResponseWriter(w)
    25  			defer cw.flush(key)
    26  
    27  			if r.ContentLength <= 0 {
    28  				next.ServeHTTP(cw, r)
    29  				return
    30  			}
    31  
    32  			if err := decryptBody(key, r); err != nil {
    33  				w.WriteHeader(http.StatusBadRequest)
    34  				return
    35  			}
    36  
    37  			next.ServeHTTP(cw, r)
    38  		})
    39  	}
    40  }
    41  
    42  func decryptBody(key []byte, r *http.Request) error {
    43  	if r.ContentLength > maxBytes {
    44  		return errContentLengthExceeded
    45  	}
    46  
    47  	var content []byte
    48  	var err error
    49  	if r.ContentLength > 0 {
    50  		content = make([]byte, r.ContentLength)
    51  		_, err = io.ReadFull(r.Body, content)
    52  	} else {
    53  		content, err = io.ReadAll(io.LimitReader(r.Body, maxBytes))
    54  	}
    55  	if err != nil {
    56  		return err
    57  	}
    58  
    59  	content, err = base64.StdEncoding.DecodeString(string(content))
    60  	if err != nil {
    61  		return err
    62  	}
    63  
    64  	output, err := codec.EcbDecrypt(key, content)
    65  	if err != nil {
    66  		return err
    67  	}
    68  
    69  	var buf bytes.Buffer
    70  	buf.Write(output)
    71  	r.Body = io.NopCloser(&buf)
    72  
    73  	return nil
    74  }
    75  
    76  type cryptionResponseWriter struct {
    77  	http.ResponseWriter
    78  	buf *bytes.Buffer
    79  }
    80  
    81  func newCryptionResponseWriter(w http.ResponseWriter) *cryptionResponseWriter {
    82  	return &cryptionResponseWriter{
    83  		ResponseWriter: w,
    84  		buf:            new(bytes.Buffer),
    85  	}
    86  }
    87  
    88  func (w *cryptionResponseWriter) Flush() {
    89  	if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
    90  		flusher.Flush()
    91  	}
    92  }
    93  
    94  func (w *cryptionResponseWriter) Header() http.Header {
    95  	return w.ResponseWriter.Header()
    96  }
    97  
    98  // Hijack implements the http.Hijacker interface.
    99  // This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
   100  func (w *cryptionResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
   101  	if hijacked, ok := w.ResponseWriter.(http.Hijacker); ok {
   102  		return hijacked.Hijack()
   103  	}
   104  
   105  	return nil, nil, errors.New("server doesn't support hijacking")
   106  }
   107  
   108  func (w *cryptionResponseWriter) Write(p []byte) (int, error) {
   109  	return w.buf.Write(p)
   110  }
   111  
   112  func (w *cryptionResponseWriter) WriteHeader(statusCode int) {
   113  	w.ResponseWriter.WriteHeader(statusCode)
   114  }
   115  
   116  func (w *cryptionResponseWriter) flush(key []byte) {
   117  	if w.buf.Len() == 0 {
   118  		return
   119  	}
   120  
   121  	content, err := codec.EcbEncrypt(key, w.buf.Bytes())
   122  	if err != nil {
   123  		w.WriteHeader(http.StatusInternalServerError)
   124  		return
   125  	}
   126  
   127  	body := base64.StdEncoding.EncodeToString(content)
   128  	if n, err := io.WriteString(w.ResponseWriter, body); err != nil {
   129  		logx.Errorf("write response failed, error: %s", err)
   130  	} else if n < len(content) {
   131  		logx.Errorf("actual bytes: %d, written bytes: %d", len(content), n)
   132  	}
   133  }