libutil: add SizedSource

Introduce the SizeSource which allows to bound how much data is being
read from a source. It also contains a drainAll() function to discard
the rest of the source, useful to keep the nix protocol in sync.
This commit is contained in:
Jonas Chevalier 2019-08-12 21:03:48 +02:00
parent b7ea98bf34
commit 91b00b145f
No known key found for this signature in database
GPG key ID: 71BAF6D40C1D63D7

View file

@ -179,6 +179,36 @@ struct TeeSource : Source
}
};
/* A reader that consumes the original Source until 'size'. */
struct SizedSource : Source
{
Source & orig;
size_t remain;
SizedSource(Source & orig, size_t size)
: orig(orig), remain(size) { }
size_t read(unsigned char * data, size_t len)
{
if (this->remain <= 0) {
throw EndOfFile("sized: unexpected end-of-file");
}
len = std::min(len, this->remain);
size_t n = this->orig.read(data, len);
this->remain -= n;
return n;
}
/* Consume the original source until no remain data is left to consume. */
size_t drainAll()
{
std::vector<unsigned char> buf(8192);
size_t sum = 0;
while (this->remain > 0) {
size_t n = read(buf.data(), buf.size());
sum += n;
}
return sum;
}
};
/* Convert a function into a sink. */
struct LambdaSink : Sink