github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/netpoll/http_client_handler.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 netpoll

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"path"
	"strconv"
	"strings"

	"github.com/cloudwego/netpoll"

	"github.com/cloudwego/kitex/pkg/kerrors"
	"github.com/cloudwego/kitex/pkg/klog"
	"github.com/cloudwego/kitex/pkg/remote"
	"github.com/cloudwego/kitex/pkg/remote/trans"
	"github.com/cloudwego/kitex/pkg/rpcinfo"
	"github.com/cloudwego/kitex/pkg/stats"
)

func newHTTPCliTransHandler(opt *remote.ClientOption, ext trans.Extension) (remote.ClientTransHandler, error) {
	return &httpCliTransHandler{
		opt:   opt,
		codec: opt.Codec,
		ext:   ext,
	}, nil
}

type httpCliTransHandler struct {
	opt       *remote.ClientOption
	codec     remote.Codec
	transPipe *remote.TransPipeline
	ext       trans.Extension
}

// Write implements the remote.ClientTransHandler interface.
func (t *httpCliTransHandler) Write(ctx context.Context, conn net.Conn, sendMsg remote.Message) (nctx context.Context, err error) {
	var bufWriter remote.ByteBuffer
	ri := sendMsg.RPCInfo()
	rpcinfo.Record(ctx, ri, stats.WriteStart, nil)
	defer func() {
		t.ext.ReleaseBuffer(bufWriter, err)
		rpcinfo.Record(ctx, ri, stats.WriteFinish, err)
	}()

	bufWriter = t.ext.NewWriteByteBuffer(ctx, conn, sendMsg)
	buffer := netpoll.NewLinkBuffer(0)
	bodyReaderWriter := NewReaderWriterByteBuffer(buffer)
	defer bodyReaderWriter.Release(err)
	if err != nil {
		return ctx, err
	}
	err = t.codec.Encode(ctx, sendMsg, bodyReaderWriter)
	if err != nil {
		return ctx, err
	}
	err = bodyReaderWriter.Flush()
	if err != nil {
		return ctx, err
	}
	var url string
	if hu, ok := ri.To().Tag(rpcinfo.HTTPURL); ok {
		url = hu
	} else {
		url = path.Join("/", ri.Invocation().MethodName())
	}
	req, err := http.NewRequest("POST", url, netpoll.NewIOReader(buffer))
	if err != nil {
		return ctx, err
	}
	err = addMetaInfo(sendMsg, req.Header)
	if err != nil {
		return ctx, err
	}
	err = req.Write(bufWriter)
	if err != nil {
		return ctx, err
	}
	return ctx, bufWriter.Flush()
}

// Read implements the remote.ClientTransHandler interface. Read is blocked.
func (t *httpCliTransHandler) Read(ctx context.Context, conn net.Conn, msg remote.Message) (nctx context.Context, err error) {
	var bufReader remote.ByteBuffer
	rpcinfo.Record(ctx, msg.RPCInfo(), stats.ReadStart, nil)
	defer func() {
		t.ext.ReleaseBuffer(bufReader, err)
		rpcinfo.Record(ctx, msg.RPCInfo(), stats.ReadFinish, err)
	}()

	t.ext.SetReadTimeout(ctx, conn, msg.RPCInfo().Config(), remote.Client)
	bufReader = t.ext.NewReadByteBuffer(ctx, conn, msg)
	bodyReader, err := getBodyBufReader(bufReader)
	if err != nil {
		return ctx, fmt.Errorf("get body bufreader error:%w", err)
	}
	err = t.codec.Decode(ctx, msg, bodyReader)
	if err != nil {
		return ctx, err
	}
	if left := bufReader.ReadableLen(); left > 0 {
		bufReader.Skip(left)
	}
	return ctx, nil
}

// OnMessage implements the remote.ClientTransHandler interface.
func (t *httpCliTransHandler) OnMessage(ctx context.Context, args, result remote.Message) (context.Context, error) {
	// do nothing
	return ctx, nil
}

// OnInactive implements the remote.ClientTransHandler interface.
// This is called when connection is closed.
func (t *httpCliTransHandler) OnInactive(ctx context.Context, conn net.Conn) {
	// ineffective now and do nothing
}

// OnError implements the remote.ClientTransHandler interface.
// This is called when panic happens.
func (t *httpCliTransHandler) OnError(ctx context.Context, err error, conn net.Conn) {
	if pe, ok := err.(*kerrors.DetailedError); ok {
		klog.CtxErrorf(ctx, "KITEX: send http request error, remote=%s, error=%s\nstack=%s", conn.RemoteAddr(), err.Error(), pe.Stack())
	} else {
		klog.CtxErrorf(ctx, "KITEX: send http request error, remote=%s, error=%s", conn.RemoteAddr(), err.Error())
	}
}

// SetPipeline implements the remote.ClientTransHandler interface.
func (t *httpCliTransHandler) SetPipeline(p *remote.TransPipeline) {
	t.transPipe = p
}

func addMetaInfo(msg remote.Message, h http.Header) error {
	meta, ok := msg.Tags()[rpcinfo.HTTPHeader]
	if !ok {
		return nil
	}
	if header, ok := meta.(http.Header); ok {
		for k, v := range header {
			h[k] = v
		}
	} else {
		return errors.New("http header in rpcinfo type assertion failed")
	}
	return nil
}

func readLine(buffer remote.ByteBuffer) ([]byte, error) {
	var buf []byte
	for {
		buf0, err := buffer.Next(1)
		if err != nil {
			return nil, err
		}
		if buf0[0] == '\r' {
			buf1, err := buffer.Peek(1)
			if err != nil {
				return nil, err
			}
			if buf1[0] == '\n' {
				err = buffer.Skip(1)
				if err != nil {
					return nil, err
				}
				return buf, nil
			}
		} else {
			buf = append(buf, buf0[0])
		}
	}
}

// return n bytes skipped ('\r\n' not included)
func skipLine(buffer remote.ByteBuffer) (n int, err error) {
	for {
		buf0, err := buffer.Next(1)
		if err != nil {
			return n, err
		}
		if buf0[0] == '\r' {
			buf1, err := buffer.Peek(1)
			if err != nil {
				return n, err
			}
			if buf1[0] == '\n' {
				err = buffer.Skip(1)
				if err != nil {
					return n, err
				}
				return n, err
			}
		} else {
			n++
		}
	}
}

func parseHTTPResponseHead(line string) (protoMajor, protoMinor, statusCodeInt int, err error) {
	var proto, status, statusCode string
	i := strings.IndexByte(line, ' ')
	if i == -1 {
		return 0, 0, 0, errors.New("malformed HTTP response: " + line)
	}
	proto = line[:i]
	status = strings.TrimLeft(line[i+1:], " ")
	statusCode = status
	if i := strings.IndexByte(status, ' '); i != -1 {
		statusCode = status[:i]
	}
	if len(statusCode) != 3 {
		return 0, 0, 0, errors.New("malformed HTTP status code: " + statusCode)
	}
	statusCodeInt, err = strconv.Atoi(statusCode)
	if err != nil || statusCodeInt < 0 {
		return 0, 0, 0, errors.New("malformed HTTP status code: " + statusCode)
	}
	var ok bool
	if protoMajor, protoMinor, ok = http.ParseHTTPVersion(proto); !ok {
		return 0, 0, 0, errors.New("malformed HTTP version: " + proto)
	}
	return
}

func skipToBody(buffer remote.ByteBuffer) error {
	head, err := readLine(buffer)
	if err != nil {
		return err
	}
	_, _, statusCode, err := parseHTTPResponseHead(string(head))
	if err != nil {
		return err
	}
	if statusCode != 200 {
		s, err := buffer.ReadString(buffer.ReadableLen())
		if err != nil {
			return fmt.Errorf("http code: %d, read request error:\n%w", statusCode, err)
		}
		return fmt.Errorf("http code: %d, error:\n%s", statusCode, s)
	}
	for {
		n, err := skipLine(buffer)
		if err != nil {
			return err
		}
		if n == 0 {
			return nil
		}
	}
}

func getBodyBufReader(buf remote.ByteBuffer) (remote.ByteBuffer, error) {
	br := bufio.NewReader(buf)
	hr, err := http.ReadResponse(br, nil)
	if err != nil {
		return nil, fmt.Errorf("read http response error:%w", err)
	}
	if hr.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("http response not OK, StatusCode: %d", hr.StatusCode)
	}
	b, err := ioutil.ReadAll(hr.Body)
	hr.Body.Close()
	if err != nil {
		return nil, fmt.Errorf("read http response body error:%w", err)
	}
	return remote.NewReaderBuffer(b), nil
}