Add pluggable codec support

server/codec.go

@@ -0,0 +1,73 @@
+package server
+
+import (
+	"io"
+	"net/rpc"
+
+	"github.com/youtube/vitess/go/rpcplus"
+	"github.com/youtube/vitess/go/rpcplus/jsonrpc"
+	"github.com/youtube/vitess/go/rpcplus/pbrpc"
+)
+
+type CodecFunc func(io.ReadWriteCloser) rpc.ServerCodec
+
+// for internal use only
+type codecFunc func(io.ReadWriteCloser) rpcplus.ServerCodec
+
+var (
+	defaultCodecs = map[string]codecFunc{
+		"application/json":         jsonrpc.NewServerCodec,
+		"application/json-rpc":     jsonrpc.NewServerCodec,
+		"application/protobuf":     pbrpc.NewServerCodec,
+		"application/proto-rpc":    pbrpc.NewServerCodec,
+		"application/octet-stream": pbrpc.NewServerCodec,
+	}
+)
+
+// wraps an net/rpc ServerCodec to provide an rpcplus.ServerCodec
+// temporary until we strip out use of rpcplus
+type rpcCodecWrap struct {
+	r rpc.ServerCodec
+}
+
+func (cw *rpcCodecWrap) ReadRequestHeader(r *rpcplus.Request) error {
+	rc := &rpc.Request{
+		ServiceMethod: r.ServiceMethod,
+		Seq:           r.Seq,
+	}
+	err := cw.r.ReadRequestHeader(rc)
+	r.ServiceMethod = rc.ServiceMethod
+	r.Seq = rc.Seq
+	return err
+}
+
+func (cw *rpcCodecWrap) ReadRequestBody(b interface{}) error {
+	return cw.r.ReadRequestBody(b)
+}
+
+func (cw *rpcCodecWrap) WriteResponse(r *rpcplus.Response, b interface{}, l bool) error {
+	rc := &rpc.Response{
+		ServiceMethod: r.ServiceMethod,
+		Seq:           r.Seq,
+		Error:         r.Error,
+	}
+	err := cw.r.WriteResponse(rc, b)
+	r.ServiceMethod = rc.ServiceMethod
+	r.Seq = rc.Seq
+	r.Error = r.Error
+	return err
+}
+
+func (cw *rpcCodecWrap) Close() error {
+	return cw.r.Close()
+}
+
+// wraps a CodecFunc to provide an internal codecFunc
+// temporary until we strip rpcplus out
+func codecWrap(cf CodecFunc) codecFunc {
+	return func(rwc io.ReadWriteCloser) rpcplus.ServerCodec {
+		return &rpcCodecWrap{
+			r: cf(rwc),
+		}
+	}
+}

server/options.go

@@ -7,6 +7,7 @@ import (
 )
 
 type options struct {
+	codecs    map[string]CodecFunc
 	broker    broker.Broker
 	registry  registry.Registry
 	transport transport.Transport
@@ -19,7 +20,9 @@ type options struct {
 }
 
 func newOptions(opt ...Option) options {
-	var opts options
+	opts := options{
+		codecs: make(map[string]CodecFunc),
+	}
 
 	for _, o := range opt {
 		o(&opts)
@@ -116,6 +119,12 @@ func Broker(b broker.Broker) Option {
 	}
 }
 
+func Codec(contentType string, cf CodecFunc) Option {
+	return func(o *options) {
+		o.codecs[contentType] = cf
+	}
+}
+
 func Registry(r registry.Registry) Option {
 	return func(o *options) {
 		o.registry = r

server/rpc_codec.go

@@ -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 {

server/rpc_server.go

@@ -43,7 +43,20 @@ func (s *rpcServer) accept(sock transport.Socket) {
 		return
 	}
 
-	codec := newRpcPlusCodec(&msg, sock)
+	cf, err := s.codecFunc(msg.Header["Content-Type"])
+	// TODO: needs better error handling
+	if err != nil {
+		sock.Send(&transport.Message{
+			Header: map[string]string{
+				"Content-Type": "text/plain",
+			},
+			Body: []byte(err.Error()),
+		})
+		sock.Close()
+		return
+	}
+
+	codec := newRpcPlusCodec(&msg, sock, cf)
 
 	// strip our headers
 	hdr := make(map[string]string)
@@ -55,11 +68,21 @@ func (s *rpcServer) accept(sock transport.Socket) {
 	ctx := c.WithMetadata(context.Background(), hdr)
 	// TODO: needs better error handling
 	if err := s.rpc.ServeRequestWithContext(ctx, codec); err != nil {
-		log.Errorf("Unexpected error servinc request, closing socket: %v", err)
+		log.Errorf("Unexpected error serving request, closing socket: %v", err)
 		sock.Close()
 	}
 }
 
+func (s *rpcServer) codecFunc(contentType string) (codecFunc, error) {
+	if cf, ok := s.opts.codecs[contentType]; ok {
+		return codecWrap(cf), nil
+	}
+	if cf, ok := defaultCodecs[contentType]; ok {
+		return cf, nil
+	}
+	return nil, fmt.Errorf("Unsupported Content-Type: %s", contentType)
+}
+
 func (s *rpcServer) Config() options {
 	s.RLock()
 	opts := s.opts