1// Copyright 2015 The etcd Authors 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// http://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14 15package rafthttp 16 17import ( 18 "encoding/binary" 19 "errors" 20 "io" 21 22 "github.com/coreos/etcd/pkg/pbutil" 23 "github.com/coreos/etcd/raft/raftpb" 24) 25 26// messageEncoder is a encoder that can encode all kinds of messages. 27// It MUST be used with a paired messageDecoder. 28type messageEncoder struct { 29 w io.Writer 30} 31 32func (enc *messageEncoder) encode(m *raftpb.Message) error { 33 if err := binary.Write(enc.w, binary.BigEndian, uint64(m.Size())); err != nil { 34 return err 35 } 36 _, err := enc.w.Write(pbutil.MustMarshal(m)) 37 return err 38} 39 40// messageDecoder is a decoder that can decode all kinds of messages. 41type messageDecoder struct { 42 r io.Reader 43} 44 45var ( 46 readBytesLimit uint64 = 512 * 1024 * 1024 // 512 MB 47 ErrExceedSizeLimit = errors.New("rafthttp: error limit exceeded") 48) 49 50func (dec *messageDecoder) decode() (raftpb.Message, error) { 51 return dec.decodeLimit(readBytesLimit) 52} 53 54func (dec *messageDecoder) decodeLimit(numBytes uint64) (raftpb.Message, error) { 55 var m raftpb.Message 56 var l uint64 57 if err := binary.Read(dec.r, binary.BigEndian, &l); err != nil { 58 return m, err 59 } 60 if l > numBytes { 61 return m, ErrExceedSizeLimit 62 } 63 buf := make([]byte, int(l)) 64 if _, err := io.ReadFull(dec.r, buf); err != nil { 65 return m, err 66 } 67 return m, m.Unmarshal(buf) 68} 69