57 lines
1006 B
Go
57 lines
1006 B
Go
package codec
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/gob"
|
|
"fmt"
|
|
"io"
|
|
)
|
|
|
|
type GobCodec struct {
|
|
conn io.ReadWriteCloser
|
|
buf *bufio.Writer
|
|
dec *gob.Decoder
|
|
enc *gob.Encoder
|
|
}
|
|
|
|
var _ Codec = (*GobCodec)(nil)
|
|
|
|
func NewGobCodec(conn io.ReadWriteCloser) Codec {
|
|
buf := bufio.NewWriter(conn)
|
|
return &GobCodec{
|
|
conn: conn,
|
|
buf: buf,
|
|
dec: gob.NewDecoder(conn),
|
|
enc: gob.NewEncoder(buf),
|
|
}
|
|
}
|
|
|
|
func (c *GobCodec) ReadHeader(h *Header) error {
|
|
return c.dec.Decode(h)
|
|
}
|
|
|
|
func (c *GobCodec) ReadBody(body interface{}) error {
|
|
return c.dec.Decode(body)
|
|
}
|
|
|
|
func (c *GobCodec) Write(h *Header, body interface{}) (err error) {
|
|
defer func() {
|
|
_ = c.buf.Flush()
|
|
if err != nil {
|
|
_ = c.Close()
|
|
}
|
|
}()
|
|
|
|
if err = c.enc.Encode(h); err != nil {
|
|
return fmt.Errorf("rpc codec: gob error encoding header: %v", err)
|
|
}
|
|
if err = c.enc.Encode(body); err != nil {
|
|
return fmt.Errorf("rpc codec: gob error encoding body: %v", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *GobCodec) Close() error {
|
|
return c.conn.Close()
|
|
}
|