1#region Copyright notice and license
2// Protocol Buffers - Google's data interchange format
3// Copyright 2008 Google Inc.  All rights reserved.
4// https://developers.google.com/protocol-buffers/
5//
6// Redistribution and use in source and binary forms, with or without
7// modification, are permitted provided that the following conditions are
8// met:
9//
10//     * Redistributions of source code must retain the above copyright
11// notice, this list of conditions and the following disclaimer.
12//     * Redistributions in binary form must reproduce the above
13// copyright notice, this list of conditions and the following disclaimer
14// in the documentation and/or other materials provided with the
15// distribution.
16//     * Neither the name of Google Inc. nor the names of its
17// contributors may be used to endorse or promote products derived from
18// this software without specific prior written permission.
19//
20// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
21// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
22// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
23// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
24// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
25// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
26// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
27// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
28// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
29// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
30// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
31#endregion
32
33using System;
34using System.Buffers;
35using System.IO;
36using System.Runtime.CompilerServices;
37using System.Security;
38
39namespace Google.Protobuf
40{
41    /// <summary>
42    /// Abstraction for reading from a stream / read only sequence.
43    /// Parsing from the buffer is a loop of reading from current buffer / refreshing the buffer once done.
44    /// </summary>
45    [SecuritySafeCritical]
46    internal struct SegmentedBufferHelper
47    {
48        private int? totalLength;
49        private ReadOnlySequence<byte>.Enumerator readOnlySequenceEnumerator;
50        private CodedInputStream codedInputStream;
51
52        /// <summary>
53        /// Initialize an instance with a coded input stream.
54        /// This approach is faster than using a constructor because the instance to initialize is passed by reference
55        /// and we can write directly into it without copying.
56        /// </summary>
57        [MethodImpl(MethodImplOptions.AggressiveInlining)]
58        public static void Initialize(CodedInputStream codedInputStream, out SegmentedBufferHelper instance)
59        {
60            instance.totalLength = codedInputStream.InternalInputStream == null ? (int?)codedInputStream.InternalBuffer.Length : null;
61            instance.readOnlySequenceEnumerator = default;
62            instance.codedInputStream = codedInputStream;
63        }
64
65        /// <summary>
66        /// Initialize an instance with a read only sequence.
67        /// This approach is faster than using a constructor because the instance to initialize is passed by reference
68        /// and we can write directly into it without copying.
69        /// </summary>
70        [MethodImpl(MethodImplOptions.AggressiveInlining)]
71        public static void Initialize(ReadOnlySequence<byte> sequence, out SegmentedBufferHelper instance, out ReadOnlySpan<byte> firstSpan)
72        {
73            instance.codedInputStream = null;
74            if (sequence.IsSingleSegment)
75            {
76                firstSpan = sequence.First.Span;
77                instance.totalLength = firstSpan.Length;
78                instance.readOnlySequenceEnumerator = default;
79            }
80            else
81            {
82                instance.readOnlySequenceEnumerator = sequence.GetEnumerator();
83                instance.totalLength = (int) sequence.Length;
84
85                // set firstSpan to the first segment
86                instance.readOnlySequenceEnumerator.MoveNext();
87                firstSpan = instance.readOnlySequenceEnumerator.Current.Span;
88            }
89        }
90
91        public bool RefillBuffer(ref ReadOnlySpan<byte> buffer, ref ParserInternalState state, bool mustSucceed)
92        {
93            if (codedInputStream != null)
94            {
95                return RefillFromCodedInputStream(ref buffer, ref state, mustSucceed);
96            }
97            else
98            {
99                return RefillFromReadOnlySequence(ref buffer, ref state, mustSucceed);
100            }
101        }
102
103        public int? TotalLength => totalLength;
104
105        public CodedInputStream CodedInputStream => codedInputStream;
106
107        /// <summary>
108        /// Sets currentLimit to (current position) + byteLimit. This is called
109        /// when descending into a length-delimited embedded message. The previous
110        /// limit is returned.
111        /// </summary>
112        /// <returns>The old limit.</returns>
113        public static int PushLimit(ref ParserInternalState state, int byteLimit)
114        {
115            if (byteLimit < 0)
116            {
117                throw InvalidProtocolBufferException.NegativeSize();
118            }
119            byteLimit += state.totalBytesRetired + state.bufferPos;
120            int oldLimit = state.currentLimit;
121            if (byteLimit > oldLimit)
122            {
123                throw InvalidProtocolBufferException.TruncatedMessage();
124            }
125            state.currentLimit = byteLimit;
126
127            RecomputeBufferSizeAfterLimit(ref state);
128
129            return oldLimit;
130        }
131
132        /// <summary>
133        /// Discards the current limit, returning the previous limit.
134        /// </summary>
135        public static void PopLimit(ref ParserInternalState state, int oldLimit)
136        {
137            state.currentLimit = oldLimit;
138            RecomputeBufferSizeAfterLimit(ref state);
139        }
140
141        /// <summary>
142        /// Returns whether or not all the data before the limit has been read.
143        /// </summary>
144        /// <returns></returns>
145        public static bool IsReachedLimit(ref ParserInternalState state)
146        {
147            if (state.currentLimit == int.MaxValue)
148            {
149                return false;
150            }
151            int currentAbsolutePosition = state.totalBytesRetired + state.bufferPos;
152            return currentAbsolutePosition >= state.currentLimit;
153        }
154
155        /// <summary>
156        /// Returns true if the stream has reached the end of the input. This is the
157        /// case if either the end of the underlying input source has been reached or
158        /// the stream has reached a limit created using PushLimit.
159        /// </summary>
160        [MethodImpl(MethodImplOptions.AggressiveInlining)]
161        public static bool IsAtEnd(ref ReadOnlySpan<byte> buffer, ref ParserInternalState state)
162        {
163            return state.bufferPos == state.bufferSize && !state.segmentedBufferHelper.RefillBuffer(ref buffer, ref state, false);
164        }
165
166        private bool RefillFromReadOnlySequence(ref ReadOnlySpan<byte> buffer, ref ParserInternalState state, bool mustSucceed)
167        {
168            CheckCurrentBufferIsEmpty(ref state);
169
170            if (state.totalBytesRetired + state.bufferSize == state.currentLimit)
171            {
172                // Oops, we hit a limit.
173                if (mustSucceed)
174                {
175                    throw InvalidProtocolBufferException.TruncatedMessage();
176                }
177                else
178                {
179                    return false;
180                }
181            }
182
183            state.totalBytesRetired += state.bufferSize;
184
185            state.bufferPos = 0;
186            state.bufferSize = 0;
187            while (readOnlySequenceEnumerator.MoveNext())
188            {
189                buffer = readOnlySequenceEnumerator.Current.Span;
190                state.bufferSize = buffer.Length;
191                if (buffer.Length != 0)
192                {
193                    break;
194                }
195            }
196
197            if (state.bufferSize == 0)
198            {
199                if (mustSucceed)
200                {
201                    throw InvalidProtocolBufferException.TruncatedMessage();
202                }
203                else
204                {
205                    return false;
206                }
207            }
208            else
209            {
210                RecomputeBufferSizeAfterLimit(ref state);
211                int totalBytesRead =
212                    state.totalBytesRetired + state.bufferSize + state.bufferSizeAfterLimit;
213                if (totalBytesRead < 0 || totalBytesRead > state.sizeLimit)
214                {
215                    throw InvalidProtocolBufferException.SizeLimitExceeded();
216                }
217                return true;
218            }
219        }
220
221        private bool RefillFromCodedInputStream(ref ReadOnlySpan<byte> buffer, ref ParserInternalState state, bool mustSucceed)
222        {
223            CheckCurrentBufferIsEmpty(ref state);
224
225            if (state.totalBytesRetired + state.bufferSize == state.currentLimit)
226            {
227                // Oops, we hit a limit.
228                if (mustSucceed)
229                {
230                    throw InvalidProtocolBufferException.TruncatedMessage();
231                }
232                else
233                {
234                    return false;
235                }
236            }
237
238            Stream input = codedInputStream.InternalInputStream;
239
240            state.totalBytesRetired += state.bufferSize;
241
242            state.bufferPos = 0;
243            state.bufferSize = (input == null) ? 0 : input.Read(codedInputStream.InternalBuffer, 0, buffer.Length);
244            if (state.bufferSize < 0)
245            {
246                throw new InvalidOperationException("Stream.Read returned a negative count");
247            }
248            if (state.bufferSize == 0)
249            {
250                if (mustSucceed)
251                {
252                    throw InvalidProtocolBufferException.TruncatedMessage();
253                }
254                else
255                {
256                    return false;
257                }
258            }
259            else
260            {
261                RecomputeBufferSizeAfterLimit(ref state);
262                int totalBytesRead =
263                    state.totalBytesRetired + state.bufferSize + state.bufferSizeAfterLimit;
264                if (totalBytesRead < 0 || totalBytesRead > state.sizeLimit)
265                {
266                    throw InvalidProtocolBufferException.SizeLimitExceeded();
267                }
268                return true;
269            }
270        }
271
272        private static void RecomputeBufferSizeAfterLimit(ref ParserInternalState state)
273        {
274            state.bufferSize += state.bufferSizeAfterLimit;
275            int bufferEnd = state.totalBytesRetired + state.bufferSize;
276            if (bufferEnd > state.currentLimit)
277            {
278                // Limit is in current buffer.
279                state.bufferSizeAfterLimit = bufferEnd - state.currentLimit;
280                state.bufferSize -= state.bufferSizeAfterLimit;
281            }
282            else
283            {
284                state.bufferSizeAfterLimit = 0;
285            }
286        }
287
288        private static void CheckCurrentBufferIsEmpty(ref ParserInternalState state)
289        {
290            if (state.bufferPos < state.bufferSize)
291            {
292                throw new InvalidOperationException("RefillBuffer() called when buffer wasn't empty.");
293            }
294        }
295    }
296}