github.com/cloudwego/kitex@v0.9.0/pkg/generic/descriptor/http.go

/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package descriptor

import (
	"errors"
	"net/http"
	"net/url"

	"github.com/bytedance/sonic/ast"
	dhttp "github.com/cloudwego/dynamicgo/http"

	"github.com/cloudwego/kitex/pkg/utils"
)

// MIMEType ...
type MIMEType string

const (
	MIMEApplicationJson     = "application/json"
	MIMEApplicationProtobuf = "application/x-protobuf"
)

var (
	_ dhttp.RequestGetter  = &HTTPRequest{}
	_ dhttp.ResponseSetter = &HTTPResponse{}
)

// HTTPRequest ...
type HTTPRequest struct {
	Params      *Params // path params
	Request     *http.Request
	RawBody     []byte
	Body        map[string]interface{}
	GeneralBody interface{} // body of other representation, used with ContentType
	ContentType MIMEType
	cookies     map[string]string
	query       url.Values
	bodyMap     *ast.Node
}

// GetHeader implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetHeader(key string) string {
	return req.Request.Header.Get(key)
}

// GetCookie implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetCookie(key string) string {
	if req.cookies == nil {
		req.cookies = map[string]string{}
		for _, cookie := range req.Request.Cookies() {
			req.cookies[cookie.Name] = cookie.Value
		}
	}
	return req.cookies[key]
}

// GetQuery implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetQuery(key string) string {
	if req.Request.URL == nil {
		return ""
	}
	if req.query == nil {
		req.query = req.Request.URL.Query()
	}
	return req.query.Get(key)
}

// GetBody implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetBody() []byte {
	return req.RawBody
}

// GetMethod implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetMethod() string {
	return req.Request.Method
}

// GetPath implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetPath() string {
	if req.Request.URL == nil {
		return ""
	}
	return req.Request.URL.Path
}

// GetHost implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetHost() string {
	return req.Request.Host
}

// GetParam implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetParam(key string) string {
	return req.Params.ByName(key)
}

// GetMapBody implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetMapBody(key string) string {
	if err := req.initializeBodyMap(); err != nil {
		return ""
	}

	v := req.bodyMap.Get(key)
	if v.Check() != nil {
		return ""
	}
	if _, err := v.Raw(); err != nil {
		return ""
	}
	j, err := v.String()
	if err != nil {
		return ""
	}
	return j
}

// GetPostForm implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetPostForm(key string) string {
	return req.Request.PostFormValue(key)
}

// GetUri implements http.RequestGetter of dynamicgo
func (req *HTTPRequest) GetUri() string {
	if req.Request.URL == nil {
		return ""
	}
	return req.Request.URL.String()
}

func (req *HTTPRequest) initializeBodyMap() error {
	if req.bodyMap == nil {
		if len(req.RawBody) == 0 {
			return errors.New("the length of RawBody is 0")
		}
		body := req.RawBody
		s := utils.SliceByteToString(body)
		node := ast.NewRaw(s)
		req.bodyMap = &node
	}
	return nil
}

// HTTPResponse ...
type HTTPResponse struct {
	Header      http.Header
	StatusCode  int32
	RawBody     []byte // this field is set only when generic.UseRawBodyForHTTPResp(true) is set
	Body        map[string]interface{}
	GeneralBody interface{} // body of other representation, used with ContentType
	ContentType MIMEType
	Renderer    Renderer
}

// NewHTTPResponse HTTP response for JSON body
func NewHTTPResponse() *HTTPResponse {
	return &HTTPResponse{
		Header:      http.Header{},
		ContentType: MIMEApplicationJson,
		Body:        map[string]interface{}{},
		Renderer:    JsonRenderer{},
	}
}

// SetStatusCode implements http.ResponseSetter of dynamicgo
func (resp *HTTPResponse) SetStatusCode(code int) error {
	resp.StatusCode = int32(code)
	return nil
}

// SetHeader implements http.ResponseSetter of dynamicgo
func (resp *HTTPResponse) SetHeader(key, val string) error {
	resp.Header.Set(key, val)
	return nil
}

// SetCookie implements http.ResponseSetter of dynamicgo
func (resp *HTTPResponse) SetCookie(key, val string) error {
	// kitex generic call does not care about Cookie
	return nil
}

// SetRawBody implements http.ResponseSetter of dynamicgo
func (resp *HTTPResponse) SetRawBody(body []byte) error {
	resp.RawBody = body
	return nil
}

func NewHTTPPbResponse(initBody interface{}) *HTTPResponse {
	return &HTTPResponse{
		Header:      http.Header{},
		ContentType: MIMEApplicationProtobuf,
		GeneralBody: initBody,
		Renderer:    PbRenderer{},
	}
}

// NewGeneralHTTPResponse init response with given MIMEType and body
func NewGeneralHTTPResponse(contentType MIMEType, initBody interface{}, renderer Renderer) *HTTPResponse {
	return &HTTPResponse{
		Header:      http.Header{},
		ContentType: contentType,
		GeneralBody: initBody,
		Renderer:    renderer,
	}
}

// Write to ResponseWriter
func (resp *HTTPResponse) Write(w http.ResponseWriter) error {
	w.WriteHeader(int(resp.StatusCode))
	for k := range resp.Header {
		w.Header().Set(k, resp.Header.Get(k))
	}

	resp.Renderer.WriteContentType(w)

	if resp.Body != nil {
		return resp.Renderer.Render(w, resp.Body)
	}

	return resp.Renderer.Render(w, resp.GeneralBody)
}

// Param in request path
type Param struct {
	Key   string
	Value string
}

// Params and recyclable
type Params struct {
	params   []Param
	recycle  func(*Params)
	recycled bool
}

// Recycle the Params
func (ps *Params) Recycle() {
	if ps.recycled {
		return
	}
	ps.recycled = true
	ps.recycle(ps)
}

// ByName search Param by given name
func (ps *Params) ByName(name string) string {
	for _, p := range ps.params {
		if p.Key == name {
			return p.Value
		}
	}
	return ""
}