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