...
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 package io
30
31 import (
32 "github.com/gogo/protobuf/proto"
33 "io"
34 )
35
36 func NewFullWriter(w io.Writer) WriteCloser {
37 return &fullWriter{w, nil}
38 }
39
40 type fullWriter struct {
41 w io.Writer
42 buffer []byte
43 }
44
45 func (this *fullWriter) WriteMsg(msg proto.Message) (err error) {
46 var data []byte
47 if m, ok := msg.(marshaler); ok {
48 n, ok := getSize(m)
49 if !ok {
50 data, err = proto.Marshal(msg)
51 if err != nil {
52 return err
53 }
54 }
55 if n >= len(this.buffer) {
56 this.buffer = make([]byte, n)
57 }
58 _, err = m.MarshalTo(this.buffer)
59 if err != nil {
60 return err
61 }
62 data = this.buffer[:n]
63 } else {
64 data, err = proto.Marshal(msg)
65 if err != nil {
66 return err
67 }
68 }
69 _, err = this.w.Write(data)
70 return err
71 }
72
73 func (this *fullWriter) Close() error {
74 if closer, ok := this.w.(io.Closer); ok {
75 return closer.Close()
76 }
77 return nil
78 }
79
80 type fullReader struct {
81 r io.Reader
82 buf []byte
83 }
84
85 func NewFullReader(r io.Reader, maxSize int) ReadCloser {
86 return &fullReader{r, make([]byte, maxSize)}
87 }
88
89 func (this *fullReader) ReadMsg(msg proto.Message) error {
90 length, err := this.r.Read(this.buf)
91 if err != nil {
92 return err
93 }
94 return proto.Unmarshal(this.buf[:length], msg)
95 }
96
97 func (this *fullReader) Close() error {
98 if closer, ok := this.r.(io.Closer); ok {
99 return closer.Close()
100 }
101 return nil
102 }
103
View as plain text