kratos-utils/http/handle.go

package http

import (
	"net/http"
	"strings"

	"github.com/go-kratos/kratos/v2/encoding"
	"github.com/go-kratos/kratos/v2/encoding/json"
	"github.com/go-kratos/kratos/v2/errors"

	_ "github.com/go-kratos/kratos/v2/transport"
)

func ErrHandle(w http.ResponseWriter, r *http.Request, err error) {
	se, ok := errors.FromError(err)
	if !ok {
		se = &errors.StatusError{
			Code:    2,
			Message: err.Error(),
		}
	}
	codec := codecForRequest(r)
	data, _ := codec.Marshal(se)
	w.Header().Set(ContentTypeHeader, contentType(codec.Name()))
	if se.Code == 0 {
		w.WriteHeader(200)
	} else {
		if se.Code < 10000 {
			se.Code = 10000 + se.Code
		}
		w.WriteHeader(400)
	}
	_, _ = w.Write(data)
}

const baseContentType = "application"

var (
	acceptHeader      = http.CanonicalHeaderKey("Accept")
	ContentTypeHeader = http.CanonicalHeaderKey("Content-Type")
)

func contentType(subtype string) string {
	return strings.Join([]string{baseContentType, subtype}, "/")
}

// codecForRequest get encoding.Codec via http.Request
func codecForRequest(r *http.Request) encoding.Codec {
	var codec encoding.Codec
	for _, accept := range r.Header[acceptHeader] {
		if codec = encoding.GetCodec(contentSubtype(accept)); codec != nil {
			break
		}
	}
	if codec == nil {
		codec = encoding.GetCodec(json.Name)
	}
	return codec
}

func contentSubtype(contentType string) string {
	if contentType == baseContentType {
		return ""
	}
	if !strings.HasPrefix(contentType, baseContentType) {
		return ""
	}
	switch contentType[len(baseContentType)] {
	case '/', ';':
		if i := strings.Index(contentType, ";"); i != -1 {
			return contentType[len(baseContentType)+1 : i]
		}
		return contentType[len(baseContentType)+1:]
	default:
		return ""
	}
}