forked from st3v/go-plugins
-
Notifications
You must be signed in to change notification settings - Fork 0
/
bsonrpc.go
89 lines (78 loc) · 1.66 KB
/
bsonrpc.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
// Package bsonrpc provides a bson-rpc codec
package bsonrpc
import (
"bytes"
"fmt"
"io"
"github.com/asim/go-bson"
"github.com/micro/go-micro/codec"
)
type bsonCodec struct {
buf *bytes.Buffer
mt codec.MessageType
rwc io.ReadWriteCloser
c *clientCodec
s *serverCodec
}
func (b *bsonCodec) Close() error {
b.buf.Reset()
return b.rwc.Close()
}
func (b *bsonCodec) String() string {
return "bson-rpc"
}
func (b *bsonCodec) Write(m *codec.Message, body interface{}) error {
switch m.Type {
case codec.Request:
return b.c.Write(m, body)
case codec.Response:
return b.s.Write(m, body)
case codec.Publication:
data, err := bson.Marshal(body)
if err != nil {
return err
}
_, err = b.rwc.Write(data)
return err
default:
return fmt.Errorf("Unrecognised message type: %v", m.Type)
}
}
func (b *bsonCodec) ReadHeader(m *codec.Message, mt codec.MessageType) error {
b.buf.Reset()
b.mt = mt
switch mt {
case codec.Request:
return b.s.ReadHeader(m)
case codec.Response:
return b.c.ReadHeader(m)
case codec.Publication:
io.Copy(b.buf, b.rwc)
default:
return fmt.Errorf("Unrecognised message type: %v", mt)
}
return nil
}
func (b *bsonCodec) ReadBody(body interface{}) error {
switch b.mt {
case codec.Request:
return b.s.ReadBody(body)
case codec.Response:
return b.c.ReadBody(body)
case codec.Publication:
if body != nil {
return bson.Unmarshal(b.buf.Bytes(), body)
}
default:
return fmt.Errorf("Unrecognised message type: %v", b.mt)
}
return nil
}
func NewCodec(rwc io.ReadWriteCloser) codec.Codec {
return &bsonCodec{
buf: bytes.NewBuffer(nil),
rwc: rwc,
c: newClientCodec(rwc),
s: newServerCodec(rwc),
}
}