github.com/unionj-cloud/go-doudou@v1.3.8-0.20221011095552-0088008e5b31/framework/http/model/model.go (about)

     1  package model
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"github.com/pkg/errors"
     7  	logger "github.com/unionj-cloud/go-doudou/toolkit/zlogger"
     8  	"io"
     9  	"io/ioutil"
    10  	"net/http"
    11  	"net/http/httptest"
    12  	"net/url"
    13  	"strings"
    14  )
    15  
    16  type Annotation struct {
    17  	Name   string
    18  	Params []string
    19  }
    20  
    21  type AnnotationStore map[string][]Annotation
    22  
    23  func (receiver AnnotationStore) HasAnnotation(key string, annotationName string) bool {
    24  	for _, item := range receiver[key] {
    25  		if item.Name == annotationName {
    26  			return true
    27  		}
    28  	}
    29  	return false
    30  }
    31  
    32  func (receiver AnnotationStore) GetParams(key string, annotationName string) []string {
    33  	for _, item := range receiver[key] {
    34  		if item.Name == annotationName {
    35  			return item.Params
    36  		}
    37  	}
    38  	return nil
    39  }
    40  
    41  // Route wraps config for route
    42  type Route struct {
    43  	Name        string
    44  	Method      string
    45  	Pattern     string
    46  	HandlerFunc http.HandlerFunc
    47  }
    48  
    49  // borrowed from httputil unexported function drainBody
    50  func CopyReqBody(b io.ReadCloser) (r1, r2 io.ReadCloser, err error) {
    51  	if b == nil || b == http.NoBody {
    52  		// No copying needed. Preserve the magic sentinel meaning of NoBody.
    53  		return http.NoBody, http.NoBody, nil
    54  	}
    55  	var buf bytes.Buffer
    56  	if _, err = buf.ReadFrom(b); err != nil {
    57  		return nil, b, err
    58  	}
    59  	if err = b.Close(); err != nil {
    60  		return nil, b, err
    61  	}
    62  	return ioutil.NopCloser(&buf), ioutil.NopCloser(bytes.NewReader(buf.Bytes())), nil
    63  }
    64  
    65  func CopyRespBody(b *bytes.Buffer) (b1, b2 *bytes.Buffer, err error) {
    66  	if b == nil {
    67  		return
    68  	}
    69  	var buf bytes.Buffer
    70  	if _, err = buf.ReadFrom(b); err != nil {
    71  		return nil, b, err
    72  	}
    73  	return &buf, bytes.NewBuffer(buf.Bytes()), nil
    74  }
    75  
    76  func JsonMarshalIndent(data interface{}, prefix, indent string, disableHTMLEscape bool) (string, error) {
    77  	b := &bytes.Buffer{}
    78  	encoder := json.NewEncoder(b)
    79  	encoder.SetEscapeHTML(!disableHTMLEscape)
    80  	encoder.SetIndent(prefix, indent)
    81  	if err := encoder.Encode(data); err != nil {
    82  		return "", errors.Errorf("failed to marshal data to JSON, %s", err)
    83  	}
    84  	return b.String(), nil
    85  }
    86  
    87  func GetReqBody(cp io.ReadCloser, r *http.Request) string {
    88  	var contentType string
    89  	if len(r.Header["Content-Type"]) > 0 {
    90  		contentType = r.Header["Content-Type"][0]
    91  	}
    92  	var reqBody string
    93  	if cp != nil {
    94  		if strings.Contains(contentType, "multipart/form-data") {
    95  			r.Body = cp
    96  			if err := r.ParseMultipartForm(32 << 20); err == nil {
    97  				reqBody = r.Form.Encode()
    98  				if unescape, err := url.QueryUnescape(reqBody); err == nil {
    99  					reqBody = unescape
   100  				}
   101  			} else {
   102  				logger.Error().Err(err).Msg("call r.ParseMultipartForm(32 << 20) error")
   103  			}
   104  		} else if strings.Contains(contentType, "application/json") {
   105  			data := make(map[string]interface{})
   106  			if err := json.NewDecoder(cp).Decode(&data); err == nil {
   107  				b, _ := json.MarshalIndent(data, "", "    ")
   108  				reqBody = string(b)
   109  			} else {
   110  				logger.Error().Err(err).Msg("call json.NewDecoder(reqBodyCopy).Decode(&data) error")
   111  			}
   112  		} else {
   113  			var buf bytes.Buffer
   114  			if _, err := buf.ReadFrom(cp); err == nil {
   115  				data := []rune(buf.String())
   116  				end := len(data)
   117  				if end > 1000 {
   118  					end = 1000
   119  				}
   120  				reqBody = string(data[:end])
   121  				if strings.Contains(contentType, "application/x-www-form-urlencoded") {
   122  					if unescape, err := url.QueryUnescape(reqBody); err == nil {
   123  						reqBody = unescape
   124  					}
   125  				}
   126  			} else {
   127  				logger.Error().Err(err).Msg("call buf.ReadFrom(reqBodyCopy) error")
   128  			}
   129  		}
   130  	}
   131  	return reqBody
   132  }
   133  
   134  func GetRespBody(rec *httptest.ResponseRecorder) string {
   135  	var (
   136  		respBody string
   137  		err      error
   138  	)
   139  	if strings.Contains(rec.Result().Header.Get("Content-Type"), "application/json") {
   140  		var respBodyCopy *bytes.Buffer
   141  		if respBodyCopy, rec.Body, err = CopyRespBody(rec.Body); err == nil {
   142  			data := make(map[string]interface{})
   143  			if err := json.NewDecoder(rec.Body).Decode(&data); err == nil {
   144  				b, _ := json.MarshalIndent(data, "", "    ")
   145  				respBody = string(b)
   146  			} else {
   147  				logger.Error().Err(err).Msg("call json.NewDecoder(rec.Body).Decode(&data) error")
   148  			}
   149  		} else {
   150  			logger.Error().Err(err).Msg("call respBodyCopy.ReadFrom(rec.Body) error")
   151  		}
   152  		rec.Body = respBodyCopy
   153  	} else {
   154  		data := []rune(rec.Body.String())
   155  		end := len(data)
   156  		if end > 1000 {
   157  			end = 1000
   158  		}
   159  		respBody = string(data[:end])
   160  	}
   161  	return respBody
   162  }