micro/gateway/rpc/response.go

package rpc

import (
	"bytes"
	"encoding/binary"
	"errors"
	"git.nspix.com/golang/micro/gateway/rpc/codec"
	"git.nspix.com/golang/micro/helper/unsafestr"
	"io"
)

const (
	SuccessCode = 0
)

type Response struct {
	codec   codec.Codec
	code    uint16 //状态码
	message string //错误消息
	error   error  //错误实体
	Body    []byte //内容
}

//获取状态码
func (r *Response) StatusCode() int {
	return int(r.code)
}

//获取错误
func (r *Response) Error() (err error) {
	if r.error != nil {
		return r.error
	}
	if r.code != SuccessCode {
		r.error = errors.New(r.message)
		err = r.error
	}
	return
}

func (r *Response) SetBody(i interface{}) (err error) {
	r.Body, err = r.codec.Encode(i)
	return
}

//解码内容
func (r *Response) Decode(i interface{}) (err error) {
	return r.codec.Decode(r.Body, i)
}

func (r *Response) Bytes() []byte {
	buf := &bytes.Buffer{}
	if _, err := r.WriteTo(buf); err == nil {
		return buf.Bytes()
	}
	return nil
}

//  |--code--|--msg length--|--body length--|--msg--|--body--|
//  |-- 2 --|-- 2 ----------| --2-----------|
func (r *Response) WriteTo(w io.Writer) (n int64, err error) {
	var (
		m int
	)
	if err = binary.Write(w, binary.LittleEndian, r.code); err != nil {
		return
	}
	n += 2
	msgLen := len(r.message)
	bodyLen := len(r.Body)
	if err = binary.Write(w, binary.LittleEndian, uint16(msgLen)); err != nil {
		return
	}
	n += 2
	if err = binary.Write(w, binary.LittleEndian, uint16(bodyLen)); err != nil {
		return
	}
	n += 2
	if m, err = w.Write(unsafestr.StringToBytes(r.message)); err == nil {
		n += int64(m)
	}
	if m, err = w.Write(r.Body); err == nil {
		n += int64(m)
	}
	return
}

func ReadResponse(b []byte) (resp *Response, err error) {
	resp = &Response{
		codec: codec.DefaultCodec,
	}
	if b == nil || len(b) <= 6 {
		err = io.ErrShortBuffer
		return
	}
	var (
		p       []byte
		msgLen  uint16
		bodyLen uint16
	)
	resp.code = binary.LittleEndian.Uint16(b[:2])
	msgLen = binary.LittleEndian.Uint16(b[2:4])
	bodyLen = binary.LittleEndian.Uint16(b[4:6])
	p = b[6:]
	if msgLen > 0 {
		if len(p) < int(msgLen) {
			err = io.ErrShortBuffer
			return
		}
		resp.message = unsafestr.BytesToString(p[:msgLen])
		p = p[msgLen:]
	}
	if bodyLen > 0 {
		if len(p) < int(bodyLen) {
			err = io.ErrShortBuffer
			return
		}
		resp.Body = p[:bodyLen]
	}
	return
}

func NewResponse() *Response {
	return &Response{
		codec: codec.DefaultCodec,
	}
}