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

     1  package httpx
     2  
     3  import (
     4  	"encoding/json"
     5  	"net/http"
     6  	"sync"
     7  
     8  	"github.com/lingyao2333/mo-zero/core/logx"
     9  	"github.com/lingyao2333/mo-zero/rest/internal/errcode"
    10  	"github.com/lingyao2333/mo-zero/rest/internal/header"
    11  )
    12  
    13  var (
    14  	errorHandler func(error) (int, interface{})
    15  	lock         sync.RWMutex
    16  )
    17  
    18  // Error writes err into w.
    19  func Error(w http.ResponseWriter, err error, fns ...func(w http.ResponseWriter, err error)) {
    20  	lock.RLock()
    21  	handler := errorHandler
    22  	lock.RUnlock()
    23  
    24  	if handler == nil {
    25  		if len(fns) > 0 {
    26  			fns[0](w, err)
    27  		} else if errcode.IsGrpcError(err) {
    28  			// don't unwrap error and get status.Message(),
    29  			// it hides the rpc error headers.
    30  			http.Error(w, err.Error(), errcode.CodeFromGrpcError(err))
    31  		} else {
    32  			http.Error(w, err.Error(), http.StatusBadRequest)
    33  		}
    34  
    35  		return
    36  	}
    37  
    38  	code, body := handler(err)
    39  	if body == nil {
    40  		w.WriteHeader(code)
    41  		return
    42  	}
    43  
    44  	e, ok := body.(error)
    45  	if ok {
    46  		http.Error(w, e.Error(), code)
    47  	} else {
    48  		WriteJson(w, code, body)
    49  	}
    50  }
    51  
    52  // Ok writes HTTP 200 OK into w.
    53  func Ok(w http.ResponseWriter) {
    54  	w.WriteHeader(http.StatusOK)
    55  }
    56  
    57  // OkJson writes v into w with 200 OK.
    58  func OkJson(w http.ResponseWriter, v interface{}) {
    59  	WriteJson(w, http.StatusOK, v)
    60  }
    61  
    62  // SetErrorHandler sets the error handler, which is called on calling Error.
    63  func SetErrorHandler(handler func(error) (int, interface{})) {
    64  	lock.Lock()
    65  	defer lock.Unlock()
    66  	errorHandler = handler
    67  }
    68  
    69  // WriteJson writes v as json string into w with code.
    70  func WriteJson(w http.ResponseWriter, code int, v interface{}) {
    71  	bs, err := json.Marshal(v)
    72  	if err != nil {
    73  		http.Error(w, err.Error(), http.StatusInternalServerError)
    74  		return
    75  	}
    76  
    77  	w.Header().Set(ContentType, header.JsonContentType)
    78  	w.WriteHeader(code)
    79  
    80  	if n, err := w.Write(bs); err != nil {
    81  		// http.ErrHandlerTimeout has been handled by http.TimeoutHandler,
    82  		// so it's ignored here.
    83  		if err != http.ErrHandlerTimeout {
    84  			logx.Errorf("write response failed, error: %s", err)
    85  		}
    86  	} else if n < len(bs) {
    87  		logx.Errorf("actual bytes: %d, written bytes: %d", len(bs), n)
    88  	}
    89  }