Skip to content

Commit c0aafe8

Browse files
xokdviumxokdvium
committed
libstore: Make uploads with filetransfer.cc consume a RestartableSource
Make uploads run in constant memory. Also change the callbacks to be noexcept, since we really don't want to be unwinding the stack in the curl thread. That will definitely corrupt that stack and make nix/curl crash in very bad ways.
1 parent 3c891ea commit c0aafe8

File tree

9 files changed

+85
-41
lines changed

9 files changed

+85
-41
lines changed

‎src/libfetchers/git-lfs-fetch.cc‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,7 +219,9 @@ std::vector<nlohmann::json> Fetch::fetchUrls(const std::vector<Pointer> & pointe
219219
nlohmann::json oidList = pointerToPayload(pointers);
220220
nlohmann::json data = {{"operation", "download"}};
221221
data["objects"] = oidList;
222-
request.data = data.dump();
222+
auto payload = data.dump();
223+
StringSource source{payload};
224+
request.data = {source};
223225

224226
FileTransferResult result = getFileTransfer()->upload(request);
225227
auto responseString = result.data;

‎src/libstore/binary-cache-store.cc‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -79,8 +79,8 @@ std::optional<std::string> BinaryCacheStore::getNixCacheInfo()
7979
void BinaryCacheStore::upsertFile(
8080
const std::string & path, std::string && data, const std::string & mimeType, uint64_t sizeHint)
8181
{
82-
StringSource source{data};
83-
upsertFile(path, source, mimeType, sizeHint);
82+
auto source = restartableSourceFromFactory([data = std::move(data)]() { return make_unique<StringSource>(data); });
83+
upsertFile(path, *source, mimeType, sizeHint);
8484
}
8585

8686
void BinaryCacheStore::getFile(const std::string & path, Callback<std::optional<std::string>> callback) noexcept
@@ -272,10 +272,19 @@ ref<const ValidPathInfo> BinaryCacheStore::addToStoreCommon(
272272

273273
/* Atomically write the NAR file. */
274274
if (repair || !fileExists(narInfo->url)) {
275-
AutoCloseFD fd = toDescriptor(open(fnTemp.c_str(), O_RDONLY));
276-
FdSource source{fd.get()};
275+
auto source = restartableSourceFromFactory([fnTemp]() {
276+
struct AutoCloseFDSource : AutoCloseFD, FdSource
277+
{
278+
AutoCloseFDSource(AutoCloseFD fd)
279+
: AutoCloseFD(std::move(fd))
280+
, FdSource(get())
281+
{
282+
}
283+
};
284+
return std::make_unique<AutoCloseFDSource>(toDescriptor(open(fnTemp.c_str(), O_RDONLY)));
285+
});
277286
stats.narWrite++;
278-
upsertFile(narInfo->url, source, "application/x-nix-nar", narInfo->fileSize);
287+
upsertFile(narInfo->url, *source, "application/x-nix-nar", narInfo->fileSize);
279288
} else
280289
stats.narWriteAverted++;
281290

‎src/libstore/filetransfer.cc‎

Lines changed: 22 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -295,20 +295,17 @@ struct curlFileTransfer : public FileTransfer
295295
return 0;
296296
}
297297

298-
size_t readOffset = 0;
299-
300-
size_t readCallback(char * buffer, size_t size, size_t nitems)
301-
{
302-
if (readOffset == request.data->length())
303-
return 0;
304-
auto count = std::min(size * nitems, request.data->length() - readOffset);
305-
assert(count);
306-
memcpy(buffer, request.data->data() + readOffset, count);
307-
readOffset += count;
308-
return count;
298+
size_t readCallback(char * buffer, size_t size, size_t nitems) noexcept
299+
try {
300+
auto data = request.data;
301+
return data->source->read(buffer, nitems * size);
302+
} catch (EndOfFile &) {
303+
return 0;
304+
} catch (...) {
305+
return CURL_READFUNC_ABORT;
309306
}
310307

311-
static size_t readCallbackWrapper(char * buffer, size_t size, size_t nitems, void * userp)
308+
static size_t readCallbackWrapper(char * buffer, size_t size, size_t nitems, void * userp) noexcept
312309
{
313310
return ((TransferItem *) userp)->readCallback(buffer, size, nitems);
314311
}
@@ -322,19 +319,24 @@ struct curlFileTransfer : public FileTransfer
322319
}
323320
#endif
324321

325-
size_t seekCallback(curl_off_t offset, int origin)
326-
{
322+
size_t seekCallback(curl_off_t offset, int origin) noexcept
323+
try {
324+
auto source = request.data->source;
327325
if (origin == SEEK_SET) {
328-
readOffset = offset;
326+
source->restart();
327+
source->skip(offset);
329328
} else if (origin == SEEK_CUR) {
330-
readOffset += offset;
329+
source->skip(offset);
331330
} else if (origin == SEEK_END) {
332-
readOffset = request.data->length() + offset;
331+
NullSink sink{};
332+
source->drainInto(sink);
333333
}
334334
return CURL_SEEKFUNC_OK;
335+
} catch (...) {
336+
return CURL_SEEKFUNC_FAIL;
335337
}
336338

337-
static size_t seekCallbackWrapper(void * clientp, curl_off_t offset, int origin)
339+
static size_t seekCallbackWrapper(void * clientp, curl_off_t offset, int origin) noexcept
338340
{
339341
return ((TransferItem *) clientp)->seekCallback(offset, origin);
340342
}
@@ -393,10 +395,10 @@ struct curlFileTransfer : public FileTransfer
393395
if (request.data) {
394396
if (request.method == HttpMethod::POST) {
395397
curl_easy_setopt(req, CURLOPT_POST, 1L);
396-
curl_easy_setopt(req, CURLOPT_POSTFIELDSIZE_LARGE, (curl_off_t) request.data->length());
398+
curl_easy_setopt(req, CURLOPT_POSTFIELDSIZE_LARGE, (curl_off_t) request.data->sizeHint);
397399
} else {
398400
curl_easy_setopt(req, CURLOPT_UPLOAD, 1L);
399-
curl_easy_setopt(req, CURLOPT_INFILESIZE_LARGE, (curl_off_t) request.data->length());
401+
curl_easy_setopt(req, CURLOPT_INFILESIZE_LARGE, (curl_off_t) request.data->sizeHint);
400402
}
401403
curl_easy_setopt(req, CURLOPT_READFUNCTION, readCallbackWrapper);
402404
curl_easy_setopt(req, CURLOPT_READDATA, this);

‎src/libstore/http-binary-cache-store.cc‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -135,18 +135,27 @@ bool HttpBinaryCacheStore::fileExists(const std::string & path)
135135
}
136136

137137
void HttpBinaryCacheStore::upsertFile(
138-
const std::string & path, Source & source, const std::string & mimeType, uint64_t sizeHint)
138+
const std::string & path, RestartableSource & source, const std::string & mimeType, uint64_t sizeHint)
139139
{
140140
auto req = makeRequest(path);
141-
auto data = source.drain();
142141
auto compressionMethod = getCompressionMethod(path);
143142

143+
std::string data;
144+
std::optional<StringSource> stringSource{};
145+
144146
if (compressionMethod) {
145-
data = compress(*compressionMethod, data);
147+
StringSink sink{};
148+
auto compressionSink = makeCompressionSink(*compressionMethod, sink);
149+
source.drainInto(*compressionSink);
150+
compressionSink->finish();
151+
data = std::move(sink.s);
146152
req.headers.emplace_back("Content-Encoding", *compressionMethod);
153+
stringSource = StringSource{data};
154+
req.data = {*stringSource};
155+
} else {
156+
req.data = {sizeHint, source};
147157
}
148158

149-
req.data = std::move(data);
150159
req.mimeType = mimeType;
151160

152161
try {

‎src/libstore/include/nix/store/binary-cache-store.hh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,8 @@ public:
100100

101101
virtual bool fileExists(const std::string & path) = 0;
102102

103-
virtual void
104-
upsertFile(const std::string & path, Source & source, const std::string & mimeType, uint64_t sizeHint) = 0;
103+
virtual void upsertFile(
104+
const std::string & path, RestartableSource & source, const std::string & mimeType, uint64_t sizeHint) = 0;
105105

106106
void upsertFile(
107107
const std::string & path,

‎src/libstore/include/nix/store/filetransfer.hh‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,26 @@ struct FileTransferRequest
114114
unsigned int baseRetryTimeMs = RETRY_TIME_MS_DEFAULT;
115115
ActivityId parentAct;
116116
bool decompress = true;
117-
std::optional<std::string> data;
117+
118+
struct UploadData
119+
{
120+
UploadData(StringSource & s)
121+
: sizeHint(s.s.length())
122+
, source(&s)
123+
{
124+
}
125+
126+
UploadData(std::size_t sizeHint, RestartableSource & source)
127+
: sizeHint(sizeHint)
128+
, source(&source)
129+
{
130+
}
131+
132+
std::size_t sizeHint = 0;
133+
RestartableSource * source = nullptr;
134+
};
135+
136+
std::optional<UploadData> data;
118137
std::string mimeType;
119138
std::function<void(std::string_view data)> dataCallback;
120139
/**

‎src/libstore/include/nix/store/http-binary-cache-store.hh‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,8 +80,8 @@ protected:
8080

8181
bool fileExists(const std::string & path) override;
8282

83-
void
84-
upsertFile(const std::string & path, Source & source, const std::string & mimeType, uint64_t sizeHint) override;
83+
void upsertFile(
84+
const std::string & path, RestartableSource & source, const std::string & mimeType, uint64_t sizeHint) override;
8585

8686
FileTransferRequest makeRequest(std::string_view path);
8787

‎src/libstore/local-binary-cache-store.cc‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,8 @@ struct LocalBinaryCacheStore : virtual BinaryCacheStore
5353

5454
bool fileExists(const std::string & path) override;
5555

56-
void upsertFile(const std::string & path, Source & source, const std::string & mimeType, uint64_t sizeHint) override
56+
void upsertFile(
57+
const std::string & path, RestartableSource & source, const std::string & mimeType, uint64_t sizeHint) override
5758
{
5859
auto path2 = config->binaryCacheDir + "/" + path;
5960
static std::atomic<int> counter{0};

‎src/libstore/s3-binary-cache-store.cc‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@ class S3BinaryCacheStore : public virtual HttpBinaryCacheStore
1919
{
2020
}
2121

22-
void
23-
upsertFile(const std::string & path, Source & source, const std::string & mimeType, uint64_t sizeHint) override;
22+
void upsertFile(
23+
const std::string & path, RestartableSource & source, const std::string & mimeType, uint64_t sizeHint) override;
2424

2525
private:
2626
ref<S3BinaryCacheStoreConfig> s3Config;
@@ -53,7 +53,7 @@ class S3BinaryCacheStore : public virtual HttpBinaryCacheStore
5353
};
5454

5555
void S3BinaryCacheStore::upsertFile(
56-
const std::string & path, Source & source, const std::string & mimeType, uint64_t sizeHint)
56+
const std::string & path, RestartableSource & source, const std::string & mimeType, uint64_t sizeHint)
5757
{
5858
HttpBinaryCacheStore::upsertFile(path, source, mimeType, sizeHint);
5959
}
@@ -72,7 +72,8 @@ std::string S3BinaryCacheStore::createMultipartUpload(
7272
req.uri = VerbatimURL(url);
7373

7474
req.method = HttpMethod::POST;
75-
req.data = "";
75+
StringSource payload{std::string_view("")};
76+
req.data = {payload};
7677
req.mimeType = mimeType;
7778

7879
if (contentEncoding) {
@@ -101,7 +102,8 @@ S3BinaryCacheStore::uploadPart(std::string_view key, std::string_view uploadId,
101102
url.query["partNumber"] = std::to_string(partNumber);
102103
url.query["uploadId"] = uploadId;
103104
req.uri = VerbatimURL(url);
104-
req.data = std::move(data);
105+
StringSource payload{data};
106+
req.data = {payload};
105107
req.mimeType = "application/octet-stream";
106108

107109
auto result = getFileTransfer()->enqueueFileTransfer(req).get();

0 commit comments

Comments
 (0)