streams.h raw
1 // Copyright (c) 2009-2010 Satoshi Nakamoto
2 // Copyright (c) 2009-2022 The Limenka developers
3 // Distributed under the MIT software license, see the accompanying
4 // file COPYING or http://www.opensource.org/licenses/mit-license.php.
5
6 #ifndef LIMENKA_STREAMS_H
7 #define LIMENKA_STREAMS_H
8
9 #include <logging.h>
10 #include <serialize.h>
11 #include <span.h>
12 #include <support/allocators/zeroafterfree.h>
13 #include <util/check.h>
14 #include <util/fs_helpers.h>
15 #include <util/obfuscation.h>
16 #include <util/overflow.h>
17 #include <util/syserror.h>
18
19 #include <algorithm>
20 #include <assert.h>
21 #include <cstddef>
22 #include <cstdio>
23 #include <ios>
24 #include <limits>
25 #include <optional>
26 #include <stdint.h>
27 #include <string.h>
28 #include <string>
29 #include <utility>
30 #include <vector>
31
32 /* Minimal stream for overwriting and/or appending to an existing byte vector
33 *
34 * The referenced vector will grow as necessary
35 */
36 class VectorWriter
37 {
38 public:
39 /*
40 * @param[in] vchDataIn Referenced byte vector to overwrite/append
41 * @param[in] nPosIn Starting position. Vector index where writes should start. The vector will initially
42 * grow as necessary to max(nPosIn, vec.size()). So to append, use vec.size().
43 */
44 VectorWriter(std::vector<unsigned char>& vchDataIn, size_t nPosIn) : vchData{vchDataIn}, nPos{nPosIn}
45 {
46 if(nPos > vchData.size())
47 vchData.resize(nPos);
48 }
49 /*
50 * (other params same as above)
51 * @param[in] args A list of items to serialize starting at nPosIn.
52 */
53 template <typename... Args>
54 VectorWriter(std::vector<unsigned char>& vchDataIn, size_t nPosIn, Args&&... args) : VectorWriter{vchDataIn, nPosIn}
55 {
56 ::SerializeMany(*this, std::forward<Args>(args)...);
57 }
58 void write(Span<const std::byte> src)
59 {
60 assert(nPos <= vchData.size());
61 size_t nOverwrite = std::min(src.size(), vchData.size() - nPos);
62 if (nOverwrite) {
63 memcpy(vchData.data() + nPos, src.data(), nOverwrite);
64 }
65 if (nOverwrite < src.size()) {
66 vchData.insert(vchData.end(), UCharCast(src.data()) + nOverwrite, UCharCast(src.data() + src.size()));
67 }
68 nPos += src.size();
69 }
70 template <typename T>
71 VectorWriter& operator<<(const T& obj)
72 {
73 ::Serialize(*this, obj);
74 return (*this);
75 }
76
77 private:
78 std::vector<unsigned char>& vchData;
79 size_t nPos;
80 };
81
82 /** Minimal stream for reading from an existing byte array by Span.
83 */
84 class SpanReader
85 {
86 private:
87 Span<const unsigned char> m_data;
88
89 public:
90 /**
91 * @param[in] data Referenced byte vector to overwrite/append
92 */
93 explicit SpanReader(Span<const unsigned char> data) : m_data{data} {}
94
95 template<typename T>
96 SpanReader& operator>>(T&& obj)
97 {
98 ::Unserialize(*this, obj);
99 return (*this);
100 }
101
102 size_t size() const { return m_data.size(); }
103 bool empty() const { return m_data.empty(); }
104
105 void read(Span<std::byte> dst)
106 {
107 if (dst.size() == 0) {
108 return;
109 }
110
111 // Read from the beginning of the buffer
112 if (dst.size() > m_data.size()) {
113 throw std::ios_base::failure("SpanReader::read(): end of data");
114 }
115 memcpy(dst.data(), m_data.data(), dst.size());
116 m_data = m_data.subspan(dst.size());
117 }
118
119 void ignore(size_t n)
120 {
121 m_data = m_data.subspan(n);
122 }
123 };
124
125 /** Double ended buffer combining vector and stream-like interfaces.
126 *
127 * >> and << read and write unformatted data using the above serialization templates.
128 * Fills with data in linear time; some stringstream implementations take N^2 time.
129 */
130 class DataStream
131 {
132 protected:
133 using vector_type = SerializeData;
134 vector_type vch;
135 vector_type::size_type m_read_pos{0};
136
137 public:
138 typedef vector_type::allocator_type allocator_type;
139 typedef vector_type::size_type size_type;
140 typedef vector_type::difference_type difference_type;
141 typedef vector_type::reference reference;
142 typedef vector_type::const_reference const_reference;
143 typedef vector_type::value_type value_type;
144 typedef vector_type::iterator iterator;
145 typedef vector_type::const_iterator const_iterator;
146 typedef vector_type::reverse_iterator reverse_iterator;
147
148 explicit DataStream() = default;
149 explicit DataStream(Span<const uint8_t> sp) : DataStream{AsBytes(sp)} {}
150 explicit DataStream(Span<const value_type> sp) : vch(sp.data(), sp.data() + sp.size()) {}
151
152 std::string str() const
153 {
154 return std::string{UCharCast(data()), UCharCast(data() + size())};
155 }
156
157
158 //
159 // Vector subset
160 //
161 const_iterator begin() const { return vch.begin() + m_read_pos; }
162 iterator begin() { return vch.begin() + m_read_pos; }
163 const_iterator end() const { return vch.end(); }
164 iterator end() { return vch.end(); }
165 size_type size() const { return vch.size() - m_read_pos; }
166 bool empty() const { return vch.size() == m_read_pos; }
167 void resize(size_type n, value_type c = value_type{}) { vch.resize(n + m_read_pos, c); }
168 void reserve(size_type n) { vch.reserve(n + m_read_pos); }
169 const_reference operator[](size_type pos) const { return vch[pos + m_read_pos]; }
170 reference operator[](size_type pos) { return vch[pos + m_read_pos]; }
171 void clear() { vch.clear(); m_read_pos = 0; }
172 value_type* data() { return vch.data() + m_read_pos; }
173 const value_type* data() const { return vch.data() + m_read_pos; }
174
175 inline void Compact()
176 {
177 vch.erase(vch.begin(), vch.begin() + m_read_pos);
178 m_read_pos = 0;
179 }
180
181 bool Rewind(std::optional<size_type> n = std::nullopt)
182 {
183 // Total rewind if no size is passed
184 if (!n) {
185 m_read_pos = 0;
186 return true;
187 }
188 // Rewind by n characters if the buffer hasn't been compacted yet
189 if (*n > m_read_pos)
190 return false;
191 m_read_pos -= *n;
192 return true;
193 }
194
195
196 //
197 // Stream subset
198 //
199 bool eof() const { return size() == 0; }
200 int in_avail() const { return size(); }
201
202 void read(Span<value_type> dst)
203 {
204 if (dst.size() == 0) return;
205
206 // Read from the beginning of the buffer
207 auto next_read_pos{CheckedAdd(m_read_pos, dst.size())};
208 if (!next_read_pos.has_value() || next_read_pos.value() > vch.size()) {
209 throw std::ios_base::failure("DataStream::read(): end of data");
210 }
211 memcpy(dst.data(), &vch[m_read_pos], dst.size());
212 if (next_read_pos.value() == vch.size()) {
213 m_read_pos = 0;
214 vch.clear();
215 return;
216 }
217 m_read_pos = next_read_pos.value();
218 }
219
220 void ignore(size_t num_ignore)
221 {
222 // Ignore from the beginning of the buffer
223 auto next_read_pos{CheckedAdd(m_read_pos, num_ignore)};
224 if (!next_read_pos.has_value() || next_read_pos.value() > vch.size()) {
225 throw std::ios_base::failure("DataStream::ignore(): end of data");
226 }
227 if (next_read_pos.value() == vch.size()) {
228 m_read_pos = 0;
229 vch.clear();
230 return;
231 }
232 m_read_pos = next_read_pos.value();
233 }
234
235 void write(Span<const value_type> src)
236 {
237 // Write to the end of the buffer
238 vch.insert(vch.end(), src.begin(), src.end());
239 }
240
241 template<typename T>
242 DataStream& operator<<(const T& obj)
243 {
244 ::Serialize(*this, obj);
245 return (*this);
246 }
247
248 template<typename T>
249 DataStream& operator>>(T&& obj)
250 {
251 ::Unserialize(*this, obj);
252 return (*this);
253 }
254
255 /**
256 * XOR the contents of this stream with a certain key.
257 *
258 * @param[in] key The key used to XOR the data in this stream.
259 */
260 void Xor(const Obfuscation& key)
261 {
262 key(*this);
263 }
264
265 /** Compute total memory usage of this object (own memory + any dynamic memory). */
266 size_t GetMemoryUsage() const noexcept;
267 };
268
269 template <typename IStream>
270 class BitStreamReader
271 {
272 private:
273 IStream& m_istream;
274
275 /// Buffered byte read in from the input stream. A new byte is read into the
276 /// buffer when m_offset reaches 8.
277 uint8_t m_buffer{0};
278
279 /// Number of high order bits in m_buffer already returned by previous
280 /// Read() calls. The next bit to be returned is at this offset from the
281 /// most significant bit position.
282 int m_offset{8};
283
284 public:
285 explicit BitStreamReader(IStream& istream) : m_istream(istream) {}
286
287 /** Read the specified number of bits from the stream. The data is returned
288 * in the nbits least significant bits of a 64-bit uint.
289 */
290 uint64_t Read(int nbits) {
291 if (nbits < 0 || nbits > 64) {
292 throw std::out_of_range("nbits must be between 0 and 64");
293 }
294
295 uint64_t data = 0;
296 while (nbits > 0) {
297 if (m_offset == 8) {
298 m_istream >> m_buffer;
299 m_offset = 0;
300 }
301
302 int bits = std::min(8 - m_offset, nbits);
303 data <<= bits;
304 data |= static_cast<uint8_t>(m_buffer << m_offset) >> (8 - bits);
305 m_offset += bits;
306 nbits -= bits;
307 }
308 return data;
309 }
310 };
311
312 template <typename OStream>
313 class BitStreamWriter
314 {
315 private:
316 OStream& m_ostream;
317
318 /// Buffered byte waiting to be written to the output stream. The byte is
319 /// written buffer when m_offset reaches 8 or Flush() is called.
320 uint8_t m_buffer{0};
321
322 /// Number of high order bits in m_buffer already written by previous
323 /// Write() calls and not yet flushed to the stream. The next bit to be
324 /// written to is at this offset from the most significant bit position.
325 int m_offset{0};
326
327 public:
328 explicit BitStreamWriter(OStream& ostream) : m_ostream(ostream) {}
329
330 ~BitStreamWriter()
331 {
332 Flush();
333 }
334
335 /** Write the nbits least significant bits of a 64-bit int to the output
336 * stream. Data is buffered until it completes an octet.
337 */
338 void Write(uint64_t data, int nbits) {
339 if (nbits < 0 || nbits > 64) {
340 throw std::out_of_range("nbits must be between 0 and 64");
341 }
342
343 while (nbits > 0) {
344 int bits = std::min(8 - m_offset, nbits);
345 m_buffer |= (data << (64 - nbits)) >> (64 - 8 + m_offset);
346 m_offset += bits;
347 nbits -= bits;
348
349 if (m_offset == 8) {
350 Flush();
351 }
352 }
353 }
354
355 /** Flush any unwritten bits to the output stream, padding with 0's to the
356 * next byte boundary.
357 */
358 void Flush() {
359 if (m_offset == 0) {
360 return;
361 }
362
363 m_ostream << m_buffer;
364 m_buffer = 0;
365 m_offset = 0;
366 }
367 };
368
369 /** Non-refcounted RAII wrapper for FILE*
370 *
371 * Will automatically close the file when it goes out of scope if not null.
372 * If you're returning the file pointer, return file.release().
373 * If you need to close the file early, use autofile.fclose() instead of fclose(underlying_FILE).
374 *
375 * @note If the file has been written to, then the caller must close it
376 * explicitly with the `fclose()` method, check if it returns an error and treat
377 * such an error as if the `write()` method failed. The OS's `fclose(3)` may
378 * fail to flush to disk data that has been previously written, rendering the
379 * file corrupt.
380 */
381 class AutoFile
382 {
383 protected:
384 std::FILE* m_file;
385 Obfuscation m_obfuscation;
386 std::optional<int64_t> m_position;
387 bool m_was_written{false};
388
389 public:
390 explicit AutoFile(std::FILE* file, const Obfuscation& obfuscation = {});
391
392 ~AutoFile()
393 {
394 if (m_was_written) {
395 // Callers that wrote to the file must have closed it explicitly
396 // with the fclose() method and checked that the close succeeded.
397 // This is because here in the destructor we have no way to signal
398 // errors from fclose() which, after write, could mean the file is
399 // corrupted and must be handled properly at the call site.
400 // Destructors in C++ cannot signal an error to the callers because
401 // they do not return a value and are not allowed to throw exceptions.
402 Assume(IsNull());
403 }
404
405 if (fclose() != 0) {
406 LogError("Failed to close file: %s", SysErrorString(errno));
407 }
408 }
409
410 // Disallow copies
411 AutoFile(const AutoFile&) = delete;
412 AutoFile& operator=(const AutoFile&) = delete;
413
414 bool feof() const { return std::feof(m_file); }
415
416 int fclose()
417 {
418 if (auto rel{release()}) return std::fclose(rel);
419 return 0;
420 }
421
422 /** Get wrapped FILE* with transfer of ownership.
423 * @note This will invalidate the AutoFile object, and makes it the responsibility of the caller
424 * of this function to clean up the returned FILE*.
425 */
426 std::FILE* release()
427 {
428 std::FILE* ret{m_file};
429 m_file = nullptr;
430 return ret;
431 }
432
433 /** Return true if the wrapped FILE* is nullptr, false otherwise.
434 */
435 bool IsNull() const { return m_file == nullptr; }
436
437 /** Continue with a different XOR key */
438 void SetXor(const Obfuscation& obfuscation) { m_obfuscation = obfuscation; }
439
440 /** Implementation detail, only used internally. */
441 std::size_t detail_fread(Span<std::byte> dst);
442
443 /** Wrapper around fseek(). Will throw if seeking is not possible. */
444 void seek(int64_t offset, int origin);
445
446 /** Find position within the file. Will throw if unknown. */
447 int64_t tell();
448
449 /** Wrapper around FileCommit(). */
450 bool Commit();
451
452 void SetIdlePriority();
453
454 /** Wrapper around TruncateFile(). */
455 bool Truncate(unsigned size);
456
457 void AdviseSequential()
458 {
459 ::AdviseSequential(m_file);
460 }
461
462 //! Write a mutable buffer more efficiently than write(), obfuscating the buffer in-place.
463 void write_buffer(std::span<std::byte> src);
464
465 //
466 // Stream subset
467 //
468 void read(Span<std::byte> dst);
469 void ignore(size_t nSize);
470 void write(Span<const std::byte> src);
471
472 template <typename T>
473 AutoFile& operator<<(const T& obj)
474 {
475 ::Serialize(*this, obj);
476 return *this;
477 }
478
479 template <typename T>
480 AutoFile& operator>>(T&& obj)
481 {
482 ::Unserialize(*this, obj);
483 return *this;
484 }
485 };
486
487 using DataBuffer = std::vector<std::byte>;
488
489 /** Wrapper around an AutoFile& that implements a ring buffer to
490 * deserialize from. It guarantees the ability to rewind a given number of bytes.
491 *
492 * Will automatically close the file when it goes out of scope if not null.
493 * If you need to close the file early, use file.fclose() instead of fclose(file).
494 */
495 class BufferedFile
496 {
497 private:
498 AutoFile& m_src;
499 uint64_t nSrcPos{0}; //!< how many bytes have been read from source
500 uint64_t m_read_pos{0}; //!< how many bytes have been read from this
501 uint64_t nReadLimit; //!< up to which position we're allowed to read
502 uint64_t nRewind; //!< how many bytes we guarantee to rewind
503 std::vector<std::byte> vchBuf; //!< the buffer
504
505 //! read data from the source to fill the buffer
506 bool Fill() {
507 unsigned int pos = nSrcPos % vchBuf.size();
508 unsigned int readNow = vchBuf.size() - pos;
509 unsigned int nAvail = vchBuf.size() - (nSrcPos - m_read_pos) - nRewind;
510 if (nAvail < readNow)
511 readNow = nAvail;
512 if (readNow == 0)
513 return false;
514 size_t nBytes{m_src.detail_fread(Span{vchBuf}.subspan(pos, readNow))};
515 if (nBytes == 0) {
516 throw std::ios_base::failure{m_src.feof() ? "BufferedFile::Fill: end of file" : "BufferedFile::Fill: fread failed"};
517 }
518 nSrcPos += nBytes;
519 return true;
520 }
521
522 //! Advance the stream's read pointer (m_read_pos) by up to 'length' bytes,
523 //! filling the buffer from the file so that at least one byte is available.
524 //! Return a pointer to the available buffer data and the number of bytes
525 //! (which may be less than the requested length) that may be accessed
526 //! beginning at that pointer.
527 std::pair<std::byte*, size_t> AdvanceStream(size_t length)
528 {
529 assert(m_read_pos <= nSrcPos);
530 if (m_read_pos + length > nReadLimit) {
531 throw std::ios_base::failure("Attempt to position past buffer limit");
532 }
533 // If there are no bytes available, read from the file.
534 if (m_read_pos == nSrcPos && length > 0) Fill();
535
536 size_t buffer_offset{static_cast<size_t>(m_read_pos % vchBuf.size())};
537 size_t buffer_available{static_cast<size_t>(vchBuf.size() - buffer_offset)};
538 size_t bytes_until_source_pos{static_cast<size_t>(nSrcPos - m_read_pos)};
539 size_t advance{std::min({length, buffer_available, bytes_until_source_pos})};
540 m_read_pos += advance;
541 return std::make_pair(&vchBuf[buffer_offset], advance);
542 }
543
544 public:
545 BufferedFile(AutoFile& file LIFETIMEBOUND, uint64_t nBufSize, uint64_t nRewindIn)
546 : m_src{file}, nReadLimit{std::numeric_limits<uint64_t>::max()}, nRewind{nRewindIn}, vchBuf(nBufSize, std::byte{0})
547 {
548 if (nRewindIn >= nBufSize)
549 throw std::ios_base::failure("Rewind limit must be less than buffer size");
550 m_src.AdviseSequential();
551 }
552
553 ~BufferedFile() { fclose(); }
554
555 int fclose()
556 {
557 if (auto rel{m_src.release()}) {
558 return CloseAndUncache(rel);
559 }
560 return m_src.fclose();
561 }
562
563 //! check whether we're at the end of the source file
564 bool eof() const {
565 return m_read_pos == nSrcPos && m_src.feof();
566 }
567
568 //! read a number of bytes
569 void read(Span<std::byte> dst)
570 {
571 while (dst.size() > 0) {
572 auto [buffer_pointer, length]{AdvanceStream(dst.size())};
573 memcpy(dst.data(), buffer_pointer, length);
574 dst = dst.subspan(length);
575 }
576 }
577
578 //! Move the read position ahead in the stream to the given position.
579 //! Use SetPos() to back up in the stream, not SkipTo().
580 void SkipTo(const uint64_t file_pos)
581 {
582 assert(file_pos >= m_read_pos);
583 while (m_read_pos < file_pos) AdvanceStream(file_pos - m_read_pos);
584 }
585
586 //! return the current reading position
587 uint64_t GetPos() const {
588 return m_read_pos;
589 }
590
591 //! rewind to a given reading position
592 bool SetPos(uint64_t nPos) {
593 size_t bufsize = vchBuf.size();
594 if (nPos + bufsize < nSrcPos) {
595 // rewinding too far, rewind as far as possible
596 m_read_pos = nSrcPos - bufsize;
597 return false;
598 }
599 if (nPos > nSrcPos) {
600 // can't go this far forward, go as far as possible
601 m_read_pos = nSrcPos;
602 return false;
603 }
604 m_read_pos = nPos;
605 return true;
606 }
607
608 //! prevent reading beyond a certain position
609 //! no argument removes the limit
610 bool SetLimit(uint64_t nPos = std::numeric_limits<uint64_t>::max()) {
611 if (nPos < m_read_pos)
612 return false;
613 nReadLimit = nPos;
614 return true;
615 }
616
617 template<typename T>
618 BufferedFile& operator>>(T&& obj) {
619 ::Unserialize(*this, obj);
620 return (*this);
621 }
622
623 //! search for a given byte in the stream, and remain positioned on it
624 void FindByte(std::byte byte)
625 {
626 // For best performance, avoid mod operation within the loop.
627 size_t buf_offset{size_t(m_read_pos % uint64_t(vchBuf.size()))};
628 while (true) {
629 if (m_read_pos == nSrcPos) {
630 // No more bytes available; read from the file into the buffer,
631 // setting nSrcPos to one beyond the end of the new data.
632 // Throws exception if end-of-file reached.
633 Fill();
634 }
635 const size_t len{std::min<size_t>(vchBuf.size() - buf_offset, nSrcPos - m_read_pos)};
636 const auto it_start{vchBuf.begin() + buf_offset};
637 const auto it_find{std::find(it_start, it_start + len, byte)};
638 const size_t inc{size_t(std::distance(it_start, it_find))};
639 m_read_pos += inc;
640 if (inc < len) break;
641 buf_offset += inc;
642 if (buf_offset >= vchBuf.size()) buf_offset = 0;
643 }
644 }
645 };
646
647 /**
648 * Wrapper that buffers reads from an underlying stream.
649 * Requires underlying stream to support read() and detail_fread() calls
650 * to support fixed-size and variable-sized reads, respectively.
651 */
652 template <typename S>
653 class BufferedReader
654 {
655 S& m_src;
656 DataBuffer m_buf;
657 size_t m_buf_pos;
658
659 public:
660 //! Requires stream ownership to prevent leaving the stream at an unexpected position after buffered reads.
661 explicit BufferedReader(S&& stream LIFETIMEBOUND, size_t size = 1 << 16)
662 requires std::is_rvalue_reference_v<S&&>
663 : m_src{stream}, m_buf(size), m_buf_pos{size} {}
664
665 void read(Span<std::byte> dst)
666 {
667 if (const auto available{std::min(dst.size(), m_buf.size() - m_buf_pos)}) {
668 std::copy_n(m_buf.begin() + m_buf_pos, available, dst.begin());
669 m_buf_pos += available;
670 dst = dst.subspan(available);
671 }
672 if (dst.size()) {
673 assert(m_buf_pos == m_buf.size());
674 m_src.read(dst);
675
676 m_buf_pos = 0;
677 m_buf.resize(m_src.detail_fread(m_buf));
678 }
679 }
680
681 template <typename T>
682 BufferedReader& operator>>(T&& obj)
683 {
684 Unserialize(*this, obj);
685 return *this;
686 }
687 };
688
689 /**
690 * Wrapper that buffers writes to an underlying stream.
691 * Requires underlying stream to support write_buffer() method
692 * for efficient buffer flushing and obfuscation.
693 */
694 template <typename S>
695 class BufferedWriter
696 {
697 S& m_dst;
698 DataBuffer m_buf;
699 size_t m_buf_pos{0};
700
701 public:
702 explicit BufferedWriter(S& stream LIFETIMEBOUND, size_t size = 1 << 16) : m_dst{stream}, m_buf(size) {}
703
704 ~BufferedWriter() { flush(); }
705
706 void flush()
707 {
708 if (m_buf_pos) m_dst.write_buffer(std::span{m_buf}.first(m_buf_pos));
709 m_buf_pos = 0;
710 }
711
712 void write(std::span<const std::byte> src)
713 {
714 while (const auto available{std::min(src.size(), m_buf.size() - m_buf_pos)}) {
715 std::copy_n(src.begin(), available, m_buf.begin() + m_buf_pos);
716 m_buf_pos += available;
717 if (m_buf_pos == m_buf.size()) flush();
718 src = src.subspan(available);
719 }
720 }
721
722 template <typename T>
723 BufferedWriter& operator<<(const T& obj)
724 {
725 Serialize(*this, obj);
726 return *this;
727 }
728 };
729
730 #endif // LIMENKA_STREAMS_H
731