micro/gateway/rpc/request.go

package rpc

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

type Request struct {
	codec         codec.Codec
	conn          net.Conn
	Method        string
	Header        map[string]string
	ContentLength int
	Sequence      uint16
	Body          interface{}
	body          []byte
}

func (r *Request) reset(seq uint16, conn net.Conn) {
	r.Sequence = seq
	r.conn = conn
}

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

func (r *Request) RemoteAddr() net.Addr {
	if r.conn != nil {
		return r.conn.RemoteAddr()
	}
	return nil
}

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

func (r *Request) WriteTo(w io.Writer) (n int64, err error) {
	var (
		headerBuffer []byte
		bodyBuffer   []byte
		m            int
		headerLen    int
		bodyLen      int
	)
	if r.Header == nil {
		r.Header = make(map[string]string)
	}
	r.Header["_METHOD"] = r.Method
	if headerBuffer, err = r.codec.Encode(r.Header); err != nil {
		return
	}
	if bodyBuffer, err = r.codec.Encode(r.Body); err != nil {
		return
	}
	headerLen = len(headerBuffer)
	bodyLen = len(bodyBuffer)
	if err = binary.Write(w, binary.LittleEndian, uint16(headerLen)); err != nil {
		return
	}
	n += 2
	if err = binary.Write(w, binary.LittleEndian, uint16(bodyLen)); err != nil {
		return
	}
	n += 2
	if m, err = w.Write(headerBuffer); err == nil {
		n += int64(m)
	}
	if m, err = w.Write(bodyBuffer); err == nil {
		n += int64(m)
	}
	return
}

func ReadRequest(b []byte) (req *Request, err error) {
	var (
		p            []byte
		headerBuffer []byte
		bodyBuffer   []byte
		headerLen    uint16
		bodyLen      uint16
	)
	req = &Request{
		codec:  codec.DefaultCodec,
		Header: make(map[string]string),
	}
	if b == nil || len(b) <= 4 {
		err = io.ErrShortBuffer
		return
	}
	headerLen = binary.LittleEndian.Uint16(b[0:2])
	bodyLen = binary.LittleEndian.Uint16(b[2:4])
	p = b[4:]
	if headerLen > 0 {
		if len(p) < int(headerLen) {
			err = io.ErrShortBuffer
			return
		}
		headerBuffer = p[:headerLen]
		if err = req.codec.Decode(headerBuffer, &req.Header); err != nil {
			return
		}
		req.Method = req.Header["_METHOD"]
		p = p[headerLen:]
	}
	if bodyLen > 0 {
		if len(p) < int(bodyLen) {
			err = io.ErrShortBuffer
			return
		}
		bodyBuffer = p[:bodyLen]
		req.body = make([]byte, bodyLen)
		copy(req.body[:], bodyBuffer[:])
		req.ContentLength = int(bodyLen)
	}
	return
}

func NewRequest(method string, i interface{}) *Request {
	return &Request{
		Method: method,
		Body:   i,
		codec:  codec.DefaultCodec,
		Header: make(map[string]string),
	}
}