cubefs

1
// Copyright 2013 Matt T. Proud
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

15
package pbutil
16

17
import (
18
	"encoding/binary"
19
	"errors"
20
	"io"
21

22
	"github.com/golang/protobuf/proto"
23
)
24

25
var errInvalidVarint = errors.New("invalid varint32 encountered")
26

27
// ReadDelimited decodes a message from the provided length-delimited stream,
28
// where the length is encoded as 32-bit varint prefix to the message body.
29
// It returns the total number of bytes read and any applicable error.  This is
30
// roughly equivalent to the companion Java API's
31
// MessageLite#parseDelimitedFrom.  As per the reader contract, this function
32
// calls r.Read repeatedly as required until exactly one message including its
33
// prefix is read and decoded (or an error has occurred).  The function never
34
// reads more bytes from the stream than required.  The function never returns
35
// an error if a message has been read and decoded correctly, even if the end
36
// of the stream has been reached in doing so.  In that case, any subsequent
37
// calls return (0, io.EOF).
38
func ReadDelimited(r io.Reader, m proto.Message) (n int, err error) {
39
	// Per AbstractParser#parsePartialDelimitedFrom with
40
	// CodedInputStream#readRawVarint32.
41
	var headerBuf [binary.MaxVarintLen32]byte
42
	var bytesRead, varIntBytes int
43
	var messageLength uint64
44
	for varIntBytes == 0 { // i.e. no varint has been decoded yet.
45
		if bytesRead >= len(headerBuf) {
46
			return bytesRead, errInvalidVarint
47
		}
48
		// We have to read byte by byte here to avoid reading more bytes
49
		// than required. Each read byte is appended to what we have
50
		// read before.
51
		newBytesRead, err := r.Read(headerBuf[bytesRead : bytesRead+1])
52
		if newBytesRead == 0 {
53
			if err != nil {
54
				return bytesRead, err
55
			}
56
			// A Reader should not return (0, nil), but if it does,
57
			// it should be treated as no-op (according to the
58
			// Reader contract). So let's go on...
59
			continue
60
		}
61
		bytesRead += newBytesRead
62
		// Now present everything read so far to the varint decoder and
63
		// see if a varint can be decoded already.
64
		messageLength, varIntBytes = proto.DecodeVarint(headerBuf[:bytesRead])
65
	}
66

67
	messageBuf := make([]byte, messageLength)
68
	newBytesRead, err := io.ReadFull(r, messageBuf)
69
	bytesRead += newBytesRead
70
	if err != nil {
71
		return bytesRead, err
72
	}
73

74
	return bytesRead, proto.Unmarshal(messageBuf, m)
75
}
76