Add pluggable codec support
This commit is contained in:
		| @@ -2,11 +2,9 @@ package server | ||||
|  | ||||
| import ( | ||||
| 	"bytes" | ||||
| 	"fmt" | ||||
|  | ||||
| 	"github.com/micro/go-micro/transport" | ||||
| 	rpc "github.com/youtube/vitess/go/rpcplus" | ||||
| 	js "github.com/youtube/vitess/go/rpcplus/jsonrpc" | ||||
| 	pb "github.com/youtube/vitess/go/rpcplus/pbrpc" | ||||
| ) | ||||
|  | ||||
| type rpcPlusCodec struct { | ||||
| @@ -36,30 +34,21 @@ func (rwc *readWriteCloser) Close() error { | ||||
| 	return nil | ||||
| } | ||||
|  | ||||
| func newRpcPlusCodec(req *transport.Message, socket transport.Socket) rpc.ServerCodec { | ||||
| func newRpcPlusCodec(req *transport.Message, socket transport.Socket, cf codecFunc) rpc.ServerCodec { | ||||
| 	rwc := &readWriteCloser{ | ||||
| 		rbuf: bytes.NewBuffer(req.Body), | ||||
| 		wbuf: bytes.NewBuffer(nil), | ||||
| 	} | ||||
| 	r := &rpcPlusCodec{ | ||||
| 		socket: socket, | ||||
| 		buf:    rwc, | ||||
| 		codec:  cf(rwc), | ||||
| 		req:    req, | ||||
| 		buf: &readWriteCloser{ | ||||
| 			rbuf: bytes.NewBuffer(req.Body), | ||||
| 			wbuf: bytes.NewBuffer(nil), | ||||
| 		}, | ||||
| 		socket: socket, | ||||
| 	} | ||||
|  | ||||
| 	switch req.Header["Content-Type"] { | ||||
| 	case "application/octet-stream": | ||||
| 		r.codec = pb.NewServerCodec(r.buf) | ||||
| 	case "application/json": | ||||
| 		r.codec = js.NewServerCodec(r.buf) | ||||
| 	} | ||||
|  | ||||
| 	return r | ||||
| } | ||||
|  | ||||
| func (c *rpcPlusCodec) ReadRequestHeader(r *rpc.Request) error { | ||||
| 	if c.codec == nil { | ||||
| 		return fmt.Errorf("unsupported content type %s", c.req.Header["Content-Type"]) | ||||
| 	} | ||||
| 	return c.codec.ReadRequestHeader(r) | ||||
| } | ||||
|  | ||||
| @@ -68,19 +57,14 @@ func (c *rpcPlusCodec) ReadRequestBody(r interface{}) error { | ||||
| } | ||||
|  | ||||
| func (c *rpcPlusCodec) WriteResponse(r *rpc.Response, body interface{}, last bool) error { | ||||
| 	if c.codec == nil { | ||||
| 		return fmt.Errorf("unsupported request type: %s", c.req.Header["Content-Type"]) | ||||
| 	} | ||||
| 	c.buf.wbuf.Reset() | ||||
| 	if err := c.codec.WriteResponse(r, body, last); err != nil { | ||||
| 		return err | ||||
| 	} | ||||
|  | ||||
| 	return c.socket.Send(&transport.Message{ | ||||
| 		Header: map[string]string{"Content-Type": c.req.Header["Content-Type"]}, | ||||
| 		Body:   c.buf.wbuf.Bytes(), | ||||
| 	}) | ||||
|  | ||||
| } | ||||
|  | ||||
| func (c *rpcPlusCodec) Close() error { | ||||
|   | ||||
		Reference in New Issue
	
	Block a user