Fix use after free in BinaryStream library.
[lldb.git] / llvm / unittests / Support / BinaryStreamTest.cpp
1 //===- llvm/unittest/Support/BinaryStreamTest.cpp -------------------------===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9
10 #include "llvm/Support/BinaryByteStream.h"
11 #include "llvm/Support/BinaryItemStream.h"
12 #include "llvm/Support/BinaryStreamArray.h"
13 #include "llvm/Support/BinaryStreamReader.h"
14 #include "llvm/Support/BinaryStreamRef.h"
15 #include "llvm/Support/BinaryStreamWriter.h"
16 #include "gtest/gtest.h"
17
18 #include <unordered_map>
19
20 using namespace llvm;
21 using namespace llvm::support;
22
23 #define EXPECT_NO_ERROR(Err)                                                   \
24   {                                                                            \
25     auto E = Err;                                                              \
26     EXPECT_FALSE(static_cast<bool>(E));                                        \
27     if (E)                                                                     \
28       consumeError(std::move(E));                                              \
29   }
30
31 #define ASSERT_NO_ERROR(Err)                                                   \
32   {                                                                            \
33     auto E = Err;                                                              \
34     ASSERT_FALSE(static_cast<bool>(E));                                        \
35     if (E)                                                                     \
36       consumeError(std::move(E));                                              \
37   }
38
39 #define EXPECT_ERROR(Err)                                                      \
40   {                                                                            \
41     auto E = Err;                                                              \
42     EXPECT_TRUE(static_cast<bool>(E));                                         \
43     if (E)                                                                     \
44       consumeError(std::move(E));                                              \
45   }
46
47 namespace {
48
49 class BrokenStream : public WritableBinaryStream {
50 public:
51   BrokenStream(MutableArrayRef<uint8_t> Data, endianness Endian,
52                       uint32_t Align)
53       : Data(Data), PartitionIndex(alignDown(Data.size() / 2, Align)),
54         Endian(Endian) {}
55
56   endianness getEndian() const override { return Endian; }
57
58   Error readBytes(uint32_t Offset, uint32_t Size,
59                   ArrayRef<uint8_t> &Buffer) override {
60     if (auto EC = checkOffset(Offset, Size))
61       return EC;
62     uint32_t S = startIndex(Offset);
63     auto Ref = Data.drop_front(S);
64     if (Ref.size() >= Size) {
65       Buffer = Ref.take_front(Size);
66       return Error::success();
67     }
68
69     uint32_t BytesLeft = Size - Ref.size();
70     uint8_t *Ptr = Allocator.Allocate<uint8_t>(Size);
71     ::memcpy(Ptr, Ref.data(), Ref.size());
72     ::memcpy(Ptr + Ref.size(), Data.data(), BytesLeft);
73     Buffer = makeArrayRef<uint8_t>(Ptr, Size);
74     return Error::success();
75   }
76
77   Error readLongestContiguousChunk(uint32_t Offset,
78                                    ArrayRef<uint8_t> &Buffer) override {
79     if (auto EC = checkOffset(Offset, 1))
80       return EC;
81     uint32_t S = startIndex(Offset);
82     Buffer = Data.drop_front(S);
83     return Error::success();
84   }
85
86   uint32_t getLength() override { return Data.size(); }
87
88   Error writeBytes(uint32_t Offset, ArrayRef<uint8_t> SrcData) override {
89     if (auto EC = checkOffset(Offset, SrcData.size()))
90       return EC;
91     if (SrcData.empty())
92       return Error::success();
93
94     uint32_t S = startIndex(Offset);
95     MutableArrayRef<uint8_t> Ref(Data);
96     Ref = Ref.drop_front(S);
97     if (Ref.size() >= SrcData.size()) {
98       ::memcpy(Ref.data(), SrcData.data(), SrcData.size());
99       return Error::success();
100     }
101
102     uint32_t BytesLeft = SrcData.size() - Ref.size();
103     ::memcpy(Ref.data(), SrcData.data(), Ref.size());
104     ::memcpy(&Data[0], SrcData.data() + Ref.size(), BytesLeft);
105     return Error::success();
106   }
107   Error commit() override { return Error::success(); }
108
109 private:
110   uint32_t startIndex(uint32_t Offset) const {
111     return (Offset + PartitionIndex) % Data.size();
112   }
113
114   uint32_t endIndex(uint32_t Offset, uint32_t Size) const {
115     return (startIndex(Offset) + Size - 1) % Data.size();
116   }
117
118   // Buffer is organized like this:
119   // -------------------------------------------------
120   // | N/2 | N/2+1 | ... | N-1 | 0 | 1 | ... | N-2-1 |
121   // -------------------------------------------------
122   // So reads from the beginning actually come from the middle.
123   MutableArrayRef<uint8_t> Data;
124   uint32_t PartitionIndex = 0;
125   endianness Endian;
126   BumpPtrAllocator Allocator;
127 };
128
129 constexpr endianness Endians[] = {big, little, native};
130 constexpr uint32_t NumEndians = llvm::array_lengthof(Endians);
131 constexpr uint32_t NumStreams = 2 * NumEndians;
132
133 class BinaryStreamTest : public testing::Test {
134
135 public:
136   BinaryStreamTest() {}
137
138   void SetUp() override {
139     Streams.clear();
140     Streams.resize(NumStreams);
141     for (uint32_t I = 0; I < NumStreams; ++I)
142       Streams[I].IsContiguous = (I % 2 == 0);
143
144     InputData.clear();
145     OutputData.clear();
146   }
147
148 protected:
149   struct StreamPair {
150     bool IsContiguous;
151     std::unique_ptr<BinaryStream> Input;
152     std::unique_ptr<WritableBinaryStream> Output;
153   };
154
155   void initializeInput(ArrayRef<uint8_t> Input, uint32_t Align) {
156     InputData = Input;
157
158     BrokenInputData.resize(InputData.size());
159     if (!Input.empty()) {
160       uint32_t PartitionIndex = alignDown(InputData.size() / 2, Align);
161       uint32_t RightBytes = InputData.size() - PartitionIndex;
162       uint32_t LeftBytes = PartitionIndex;
163       if (RightBytes > 0)
164         ::memcpy(&BrokenInputData[PartitionIndex], Input.data(), RightBytes);
165       if (LeftBytes > 0)
166         ::memcpy(&BrokenInputData[0], Input.data() + RightBytes, LeftBytes);
167     }
168
169     for (uint32_t I = 0; I < NumEndians; ++I) {
170       auto InByteStream =
171           llvm::make_unique<BinaryByteStream>(InputData, Endians[I]);
172       auto InBrokenStream = llvm::make_unique<BrokenStream>(
173           BrokenInputData, Endians[I], Align);
174
175       Streams[I * 2].Input = std::move(InByteStream);
176       Streams[I * 2 + 1].Input = std::move(InBrokenStream);
177     }
178   }
179
180   void initializeOutput(uint32_t Size, uint32_t Align) {
181     OutputData.resize(Size);
182     BrokenOutputData.resize(Size);
183
184     for (uint32_t I = 0; I < NumEndians; ++I) {
185       Streams[I * 2].Output =
186           llvm::make_unique<MutableBinaryByteStream>(OutputData, Endians[I]);
187       Streams[I * 2 + 1].Output = llvm::make_unique<BrokenStream>(
188           BrokenOutputData, Endians[I], Align);
189     }
190   }
191
192   void initializeOutputFromInput(uint32_t Align) {
193     for (uint32_t I = 0; I < NumEndians; ++I) {
194       Streams[I * 2].Output =
195           llvm::make_unique<MutableBinaryByteStream>(InputData, Endians[I]);
196       Streams[I * 2 + 1].Output = llvm::make_unique<BrokenStream>(
197           BrokenInputData, Endians[I], Align);
198     }
199   }
200
201   void initializeInputFromOutput(uint32_t Align) {
202     for (uint32_t I = 0; I < NumEndians; ++I) {
203       Streams[I * 2].Input =
204           llvm::make_unique<BinaryByteStream>(OutputData, Endians[I]);
205       Streams[I * 2 + 1].Input = llvm::make_unique<BrokenStream>(
206           BrokenOutputData, Endians[I], Align);
207     }
208   }
209
210   std::vector<uint8_t> InputData;
211   std::vector<uint8_t> BrokenInputData;
212
213   std::vector<uint8_t> OutputData;
214   std::vector<uint8_t> BrokenOutputData;
215
216   std::vector<StreamPair> Streams;
217 };
218
219 // Tests that a we can read from a BinaryByteStream without a StreamReader.
220 TEST_F(BinaryStreamTest, BinaryByteStreamBounds) {
221   std::vector<uint8_t> InputData = {1, 2, 3, 4, 5};
222   initializeInput(InputData, 1);
223
224   for (auto &Stream : Streams) {
225     ArrayRef<uint8_t> Buffer;
226
227     // 1. If the read fits it should work.
228     ASSERT_EQ(InputData.size(), Stream.Input->getLength());
229     ASSERT_NO_ERROR(Stream.Input->readBytes(2, 1, Buffer));
230     EXPECT_EQ(makeArrayRef(InputData).slice(2, 1), Buffer);
231     ASSERT_NO_ERROR(Stream.Input->readBytes(0, 4, Buffer));
232     EXPECT_EQ(makeArrayRef(InputData).slice(0, 4), Buffer);
233
234     // 2. Reading past the bounds of the input should fail.
235     EXPECT_ERROR(Stream.Input->readBytes(4, 2, Buffer));
236   }
237 }
238
239 TEST_F(BinaryStreamTest, StreamRefBounds) {
240   std::vector<uint8_t> InputData = {1, 2, 3, 4, 5};
241   initializeInput(InputData, 1);
242
243   for (const auto &Stream : Streams) {
244     ArrayRef<uint8_t> Buffer;
245     BinaryStreamRef Ref(*Stream.Input);
246
247     // Read 1 byte from offset 2 should work
248     ASSERT_EQ(InputData.size(), Ref.getLength());
249     ASSERT_NO_ERROR(Ref.readBytes(2, 1, Buffer));
250     EXPECT_EQ(makeArrayRef(InputData).slice(2, 1), Buffer);
251
252     // Reading everything from offset 2 on.
253     ASSERT_NO_ERROR(Ref.readLongestContiguousChunk(2, Buffer));
254     if (Stream.IsContiguous)
255       EXPECT_EQ(makeArrayRef(InputData).slice(2), Buffer);
256     else
257       EXPECT_FALSE(Buffer.empty());
258
259     // Reading 6 bytes from offset 0 is too big.
260     EXPECT_ERROR(Ref.readBytes(0, 6, Buffer));
261     EXPECT_ERROR(Ref.readLongestContiguousChunk(6, Buffer));
262
263     // Reading 1 byte from offset 2 after dropping 1 byte is the same as reading
264     // 1 byte from offset 3.
265     Ref = Ref.drop_front(1);
266     ASSERT_NO_ERROR(Ref.readBytes(2, 1, Buffer));
267     if (Stream.IsContiguous)
268       EXPECT_EQ(makeArrayRef(InputData).slice(3, 1), Buffer);
269     else
270       EXPECT_FALSE(Buffer.empty());
271
272     // Reading everything from offset 2 on after dropping 1 byte.
273     ASSERT_NO_ERROR(Ref.readLongestContiguousChunk(2, Buffer));
274     if (Stream.IsContiguous)
275       EXPECT_EQ(makeArrayRef(InputData).slice(3), Buffer);
276     else
277       EXPECT_FALSE(Buffer.empty());
278
279     // Reading 2 bytes from offset 2 after dropping 2 bytes is the same as
280     // reading 2 bytes from offset 4, and should fail.
281     Ref = Ref.drop_front(1);
282     EXPECT_ERROR(Ref.readBytes(2, 2, Buffer));
283
284     // But if we read the longest contiguous chunk instead, we should still
285     // get the 1 byte at the end.
286     ASSERT_NO_ERROR(Ref.readLongestContiguousChunk(2, Buffer));
287     EXPECT_EQ(makeArrayRef(InputData).take_back(), Buffer);
288   }
289 }
290
291 // Test that we can write to a BinaryStream without a StreamWriter.
292 TEST_F(BinaryStreamTest, MutableBinaryByteStreamBounds) {
293   std::vector<uint8_t> InputData = {'T', 'e', 's', 't', '\0'};
294   initializeInput(InputData, 1);
295   initializeOutput(InputData.size(), 1);
296
297   // For every combination of input stream and output stream.
298   for (auto &Stream : Streams) {
299     MutableArrayRef<uint8_t> Buffer;
300     ASSERT_EQ(InputData.size(), Stream.Input->getLength());
301
302     // 1. Try two reads that are supposed to work.  One from offset 0, and one
303     // from the middle.
304     uint32_t Offsets[] = {0, 3};
305     for (auto Offset : Offsets) {
306       uint32_t ExpectedSize = Stream.Input->getLength() - Offset;
307
308       // Read everything from Offset until the end of the input data.
309       ArrayRef<uint8_t> Data;
310       ASSERT_NO_ERROR(Stream.Input->readBytes(Offset, ExpectedSize, Data));
311       ASSERT_EQ(ExpectedSize, Data.size());
312
313       // Then write it to the destination.
314       ASSERT_NO_ERROR(Stream.Output->writeBytes(0, Data));
315
316       // Then we read back what we wrote, it should match the corresponding
317       // slice of the original input data.
318       ArrayRef<uint8_t> Data2;
319       ASSERT_NO_ERROR(Stream.Output->readBytes(Offset, ExpectedSize, Data2));
320       EXPECT_EQ(makeArrayRef(InputData).drop_front(Offset), Data2);
321     }
322
323     std::vector<uint8_t> BigData = {0, 1, 2, 3, 4};
324     // 2. If the write is too big, it should fail.
325     EXPECT_ERROR(Stream.Output->writeBytes(3, BigData));
326   }
327 }
328
329 // Test that FixedStreamArray works correctly.
330 TEST_F(BinaryStreamTest, FixedStreamArray) {
331   std::vector<uint32_t> Ints = {90823, 12908, 109823, 209823};
332   ArrayRef<uint8_t> IntBytes(reinterpret_cast<uint8_t *>(Ints.data()),
333                              Ints.size() * sizeof(uint32_t));
334
335   initializeInput(IntBytes, alignof(uint32_t));
336
337   for (auto &Stream : Streams) {
338     MutableArrayRef<uint8_t> Buffer;
339     ASSERT_EQ(InputData.size(), Stream.Input->getLength());
340
341     FixedStreamArray<uint32_t> Array(*Stream.Input);
342     auto Iter = Array.begin();
343     ASSERT_EQ(Ints[0], *Iter++);
344     ASSERT_EQ(Ints[1], *Iter++);
345     ASSERT_EQ(Ints[2], *Iter++);
346     ASSERT_EQ(Ints[3], *Iter++);
347     ASSERT_EQ(Array.end(), Iter);
348   }
349 }
350
351 // Test that VarStreamArray works correctly.
352 TEST_F(BinaryStreamTest, VarStreamArray) {
353   StringLiteral Strings("1. Test2. Longer Test3. Really Long Test4. Super "
354                         "Extra Longest Test Of All");
355   ArrayRef<uint8_t> StringBytes(
356       reinterpret_cast<const uint8_t *>(Strings.data()), Strings.size());
357   initializeInput(StringBytes, 1);
358
359   struct StringExtractor {
360   public:
361     typedef uint32_t &ContextType;
362     static Error extract(BinaryStreamRef Stream, uint32_t &Len, StringRef &Item,
363                          uint32_t &Index) {
364       if (Index == 0)
365         Len = strlen("1. Test");
366       else if (Index == 1)
367         Len = strlen("2. Longer Test");
368       else if (Index == 2)
369         Len = strlen("3. Really Long Test");
370       else
371         Len = strlen("4. Super Extra Longest Test Of All");
372       ArrayRef<uint8_t> Bytes;
373       if (auto EC = Stream.readBytes(0, Len, Bytes))
374         return EC;
375       Item =
376           StringRef(reinterpret_cast<const char *>(Bytes.data()), Bytes.size());
377       ++Index;
378       return Error::success();
379     }
380   };
381
382   for (auto &Stream : Streams) {
383     uint32_t Context = 0;
384     VarStreamArray<StringRef, StringExtractor> Array(*Stream.Input, Context);
385     auto Iter = Array.begin();
386     ASSERT_EQ("1. Test", *Iter++);
387     ASSERT_EQ("2. Longer Test", *Iter++);
388     ASSERT_EQ("3. Really Long Test", *Iter++);
389     ASSERT_EQ("4. Super Extra Longest Test Of All", *Iter++);
390     ASSERT_EQ(Array.end(), Iter);
391   }
392 }
393
394 TEST_F(BinaryStreamTest, StreamReaderBounds) {
395   std::vector<uint8_t> Bytes;
396
397   initializeInput(Bytes, 1);
398   for (auto &Stream : Streams) {
399     StringRef S;
400     BinaryStreamReader Reader(*Stream.Input);
401     EXPECT_EQ(0U, Reader.bytesRemaining());
402     EXPECT_ERROR(Reader.readFixedString(S, 1));
403   }
404
405   Bytes.resize(5);
406   initializeInput(Bytes, 1);
407   for (auto &Stream : Streams) {
408     StringRef S;
409     BinaryStreamReader Reader(*Stream.Input);
410     EXPECT_EQ(Bytes.size(), Reader.bytesRemaining());
411     EXPECT_NO_ERROR(Reader.readFixedString(S, 5));
412     EXPECT_ERROR(Reader.readFixedString(S, 6));
413   }
414 }
415
416 TEST_F(BinaryStreamTest, StreamReaderIntegers) {
417   support::ulittle64_t Little{908234};
418   support::ubig32_t Big{28907823};
419   short NS = 2897;
420   int NI = -89723;
421   unsigned long NUL = 902309023UL;
422   constexpr uint32_t Size =
423       sizeof(Little) + sizeof(Big) + sizeof(NS) + sizeof(NI) + sizeof(NUL);
424
425   initializeOutput(Size, alignof(support::ulittle64_t));
426   initializeInputFromOutput(alignof(support::ulittle64_t));
427
428   for (auto &Stream : Streams) {
429     BinaryStreamWriter Writer(*Stream.Output);
430     ASSERT_NO_ERROR(Writer.writeObject(Little));
431     ASSERT_NO_ERROR(Writer.writeObject(Big));
432     ASSERT_NO_ERROR(Writer.writeInteger(NS));
433     ASSERT_NO_ERROR(Writer.writeInteger(NI));
434     ASSERT_NO_ERROR(Writer.writeInteger(NUL));
435
436     const support::ulittle64_t *Little2;
437     const support::ubig32_t *Big2;
438     short NS2;
439     int NI2;
440     unsigned long NUL2;
441
442     // 1. Reading fields individually.
443     BinaryStreamReader Reader(*Stream.Input);
444     ASSERT_NO_ERROR(Reader.readObject(Little2));
445     ASSERT_NO_ERROR(Reader.readObject(Big2));
446     ASSERT_NO_ERROR(Reader.readInteger(NS2));
447     ASSERT_NO_ERROR(Reader.readInteger(NI2));
448     ASSERT_NO_ERROR(Reader.readInteger(NUL2));
449     ASSERT_EQ(0U, Reader.bytesRemaining());
450
451     EXPECT_EQ(Little, *Little2);
452     EXPECT_EQ(Big, *Big2);
453     EXPECT_EQ(NS, NS2);
454     EXPECT_EQ(NI, NI2);
455     EXPECT_EQ(NUL, NUL2);
456   }
457 }
458
459 TEST_F(BinaryStreamTest, StreamReaderIntegerArray) {
460   // 1. Arrays of integers
461   std::vector<int> Ints = {1, 2, 3, 4, 5};
462   ArrayRef<uint8_t> IntBytes(reinterpret_cast<uint8_t *>(&Ints[0]),
463                              Ints.size() * sizeof(int));
464
465   initializeInput(IntBytes, alignof(int));
466   for (auto &Stream : Streams) {
467     BinaryStreamReader Reader(*Stream.Input);
468     ArrayRef<int> IntsRef;
469     ASSERT_NO_ERROR(Reader.readArray(IntsRef, Ints.size()));
470     ASSERT_EQ(0U, Reader.bytesRemaining());
471     EXPECT_EQ(makeArrayRef(Ints), IntsRef);
472
473     Reader.setOffset(0);
474     FixedStreamArray<int> FixedIntsRef;
475     ASSERT_NO_ERROR(Reader.readArray(FixedIntsRef, Ints.size()));
476     ASSERT_EQ(0U, Reader.bytesRemaining());
477     ASSERT_EQ(Ints, std::vector<int>(FixedIntsRef.begin(), FixedIntsRef.end()));
478   }
479 }
480
481 TEST_F(BinaryStreamTest, StreamReaderEnum) {
482   enum class MyEnum : int64_t { Foo = -10, Bar = 0, Baz = 10 };
483
484   std::vector<MyEnum> Enums = {MyEnum::Bar, MyEnum::Baz, MyEnum::Foo};
485
486   initializeOutput(Enums.size() * sizeof(MyEnum), alignof(MyEnum));
487   initializeInputFromOutput(alignof(MyEnum));
488   for (auto &Stream : Streams) {
489     BinaryStreamWriter Writer(*Stream.Output);
490     for (auto Value : Enums)
491       ASSERT_NO_ERROR(Writer.writeEnum(Value));
492
493     BinaryStreamReader Reader(*Stream.Input);
494
495     ArrayRef<MyEnum> Array;
496     FixedStreamArray<MyEnum> FSA;
497
498     for (size_t I = 0; I < Enums.size(); ++I) {
499       MyEnum Value;
500       ASSERT_NO_ERROR(Reader.readEnum(Value));
501       EXPECT_EQ(Enums[I], Value);
502     }
503     ASSERT_EQ(0U, Reader.bytesRemaining());
504   }
505 }
506
507 TEST_F(BinaryStreamTest, StreamReaderObject) {
508   struct Foo {
509     int X;
510     double Y;
511     char Z;
512
513     bool operator==(const Foo &Other) const {
514       return X == Other.X && Y == Other.Y && Z == Other.Z;
515     }
516   };
517
518   std::vector<Foo> Foos;
519   Foos.push_back({-42, 42.42, 42});
520   Foos.push_back({100, 3.1415, static_cast<char>(-89)});
521   Foos.push_back({200, 2.718, static_cast<char>(-12) });
522
523   const uint8_t *Bytes = reinterpret_cast<const uint8_t *>(&Foos[0]);
524
525   initializeInput(makeArrayRef(Bytes, 3 * sizeof(Foo)), alignof(Foo));
526
527   for (auto &Stream : Streams) {
528     // 1. Reading object pointers.
529     BinaryStreamReader Reader(*Stream.Input);
530     const Foo *FPtrOut = nullptr;
531     const Foo *GPtrOut = nullptr;
532     const Foo *HPtrOut = nullptr;
533     ASSERT_NO_ERROR(Reader.readObject(FPtrOut));
534     ASSERT_NO_ERROR(Reader.readObject(GPtrOut));
535     ASSERT_NO_ERROR(Reader.readObject(HPtrOut));
536     EXPECT_EQ(0U, Reader.bytesRemaining());
537     EXPECT_EQ(Foos[0], *FPtrOut);
538     EXPECT_EQ(Foos[1], *GPtrOut);
539     EXPECT_EQ(Foos[2], *HPtrOut);
540   }
541 }
542
543 TEST_F(BinaryStreamTest, StreamReaderStrings) {
544   std::vector<uint8_t> Bytes = {'O',  'n', 'e', '\0', 'T', 'w', 'o',
545                                 '\0', 'T', 'h', 'r',  'e', 'e', '\0',
546                                 'F',  'o', 'u', 'r',  '\0'};
547   initializeInput(Bytes, 1);
548
549   for (auto &Stream : Streams) {
550     BinaryStreamReader Reader(*Stream.Input);
551
552     StringRef S1;
553     StringRef S2;
554     StringRef S3;
555     StringRef S4;
556     ASSERT_NO_ERROR(Reader.readCString(S1));
557     ASSERT_NO_ERROR(Reader.readCString(S2));
558     ASSERT_NO_ERROR(Reader.readCString(S3));
559     ASSERT_NO_ERROR(Reader.readCString(S4));
560     ASSERT_EQ(0U, Reader.bytesRemaining());
561
562     EXPECT_EQ("One", S1);
563     EXPECT_EQ("Two", S2);
564     EXPECT_EQ("Three", S3);
565     EXPECT_EQ("Four", S4);
566
567     S1 = S2 = S3 = S4 = "";
568     Reader.setOffset(0);
569     ASSERT_NO_ERROR(Reader.readFixedString(S1, 3));
570     ASSERT_NO_ERROR(Reader.skip(1));
571     ASSERT_NO_ERROR(Reader.readFixedString(S2, 3));
572     ASSERT_NO_ERROR(Reader.skip(1));
573     ASSERT_NO_ERROR(Reader.readFixedString(S3, 5));
574     ASSERT_NO_ERROR(Reader.skip(1));
575     ASSERT_NO_ERROR(Reader.readFixedString(S4, 4));
576     ASSERT_NO_ERROR(Reader.skip(1));
577     ASSERT_EQ(0U, Reader.bytesRemaining());
578
579     EXPECT_EQ("One", S1);
580     EXPECT_EQ("Two", S2);
581     EXPECT_EQ("Three", S3);
582     EXPECT_EQ("Four", S4);
583   }
584 }
585
586 TEST_F(BinaryStreamTest, StreamWriterBounds) {
587   initializeOutput(5, 1);
588
589   for (auto &Stream : Streams) {
590     BinaryStreamWriter Writer(*Stream.Output);
591
592     // 1. Can write a string that exactly fills the buffer.
593     EXPECT_EQ(5U, Writer.bytesRemaining());
594     EXPECT_NO_ERROR(Writer.writeFixedString("abcde"));
595     EXPECT_EQ(0U, Writer.bytesRemaining());
596
597     // 2. Can write an empty string even when you're full
598     EXPECT_NO_ERROR(Writer.writeFixedString(""));
599     EXPECT_ERROR(Writer.writeFixedString("a"));
600
601     // 3. Can't write a string that is one character too long.
602     Writer.setOffset(0);
603     EXPECT_ERROR(Writer.writeFixedString("abcdef"));
604   }
605 }
606
607 TEST_F(BinaryStreamTest, StreamWriterIntegerArrays) {
608   // 3. Arrays of integers
609   std::vector<int> SourceInts = {1, 2, 3, 4, 5};
610   ArrayRef<uint8_t> SourceBytes(reinterpret_cast<uint8_t *>(&SourceInts[0]),
611                                 SourceInts.size() * sizeof(int));
612
613   initializeInput(SourceBytes, alignof(int));
614   initializeOutputFromInput(alignof(int));
615
616   for (auto &Stream : Streams) {
617     BinaryStreamReader Reader(*Stream.Input);
618     BinaryStreamWriter Writer(*Stream.Output);
619     ArrayRef<int> Ints;
620     ArrayRef<int> Ints2;
621     // First read them, then write them, then read them back.
622     ASSERT_NO_ERROR(Reader.readArray(Ints, SourceInts.size()));
623     ASSERT_NO_ERROR(Writer.writeArray(Ints));
624
625     BinaryStreamReader ReaderBacker(*Stream.Output);
626     ASSERT_NO_ERROR(ReaderBacker.readArray(Ints2, SourceInts.size()));
627
628     EXPECT_EQ(makeArrayRef(SourceInts), Ints2);
629   }
630 }
631
632 TEST_F(BinaryStreamTest, StringWriterStrings) {
633   StringRef Strings[] = {"First", "Second", "Third", "Fourth"};
634
635   size_t Length = 0;
636   for (auto S : Strings)
637     Length += S.size() + 1;
638   initializeOutput(Length, 1);
639   initializeInputFromOutput(1);
640
641   for (auto &Stream : Streams) {
642     BinaryStreamWriter Writer(*Stream.Output);
643     for (auto S : Strings)
644       ASSERT_NO_ERROR(Writer.writeCString(S));
645     std::vector<StringRef> InStrings;
646     BinaryStreamReader Reader(*Stream.Input);
647     while (!Reader.empty()) {
648       StringRef S;
649       ASSERT_NO_ERROR(Reader.readCString(S));
650       InStrings.push_back(S);
651     }
652     EXPECT_EQ(makeArrayRef(Strings), makeArrayRef(InStrings));
653   }
654 }
655 }
656
657 namespace {
658 struct BinaryItemStreamObject {
659   explicit BinaryItemStreamObject(ArrayRef<uint8_t> Bytes) : Bytes(Bytes) {}
660
661   ArrayRef<uint8_t> Bytes;
662 };
663 }
664
665 namespace llvm {
666 template <> struct BinaryItemTraits<BinaryItemStreamObject> {
667   static size_t length(const BinaryItemStreamObject &Item) {
668     return Item.Bytes.size();
669   }
670
671   static ArrayRef<uint8_t> bytes(const BinaryItemStreamObject &Item) {
672     return Item.Bytes;
673   }
674 };
675 }
676
677 namespace {
678
679 TEST_F(BinaryStreamTest, BinaryItemStream) {
680   std::vector<BinaryItemStreamObject> Objects;
681
682   struct Foo {
683     int X;
684     double Y;
685   };
686   std::vector<Foo> Foos = {{1, 1.0}, {2, 2.0}, {3, 3.0}};
687   BumpPtrAllocator Allocator;
688   for (const auto &F : Foos) {
689     uint8_t *Ptr = static_cast<uint8_t *>(Allocator.Allocate(sizeof(Foo), 
690                                                              alignof(Foo)));
691     MutableArrayRef<uint8_t> Buffer(Ptr, sizeof(Foo));
692     MutableBinaryByteStream Stream(Buffer, llvm::support::big);
693     BinaryStreamWriter Writer(Stream);
694     ASSERT_NO_ERROR(Writer.writeObject(F));
695     Objects.push_back(BinaryItemStreamObject(Buffer));
696   }
697
698   BinaryItemStream<BinaryItemStreamObject> ItemStream(big);
699   ItemStream.setItems(Objects);
700   BinaryStreamReader Reader(ItemStream);
701
702   for (const auto &F : Foos) {
703     const Foo *F2;
704     ASSERT_NO_ERROR(Reader.readObject(F2));
705
706     EXPECT_EQ(F.X, F2->X);
707     EXPECT_DOUBLE_EQ(F.Y, F2->Y);
708   }
709 }
710
711 } // end anonymous namespace