From e7ed9ae0c711c4efd83756b16379549ecff52355 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Na=C3=AFm=20Favier?= <n@monade.li>
Date: Fri, 4 Nov 2022 14:19:31 +0100
Subject: [PATCH] Restrict `readFile` context to references that appear in the
 string

When calling `builtins.readFile` on a store path, the references of that
path are currently added to the resulting string's context.

This change makes those references the *possible* context of the string,
but filters them to keep only the references whose hash actually appears
in the string, similarly to what is done for determining the runtime
references of a path.
---
 src/libexpr/primops.cc     |  5 ++++
 src/libstore/references.cc | 57 +++++++++++++++++++++++++-------------
 src/libstore/references.hh | 13 +++++++++
 3 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/src/libexpr/primops.cc b/src/libexpr/primops.cc
index 28b998474..ff620ca63 100644
--- a/src/libexpr/primops.cc
+++ b/src/libexpr/primops.cc
@@ -5,6 +5,7 @@
 #include "globals.hh"
 #include "json-to-value.hh"
 #include "names.hh"
+#include "references.hh"
 #include "store-api.hh"
 #include "util.hh"
 #include "json.hh"
@@ -1542,6 +1543,10 @@ static void prim_readFile(EvalState & state, const PosIdx pos, Value * * args, V
             refs = state.store->queryPathInfo(state.store->toStorePath(path).first)->references;
         } catch (Error &) { // FIXME: should be InvalidPathError
         }
+        // Re-scan references to filter down to just the ones that actually occur in the file.
+        auto refsSink = PathRefScanSink::fromPaths(refs);
+        refsSink << s;
+        refs = refsSink.getResultPaths();
     }
     auto context = state.store->printStorePathSet(refs);
     v.mkString(s, context);
diff --git a/src/libstore/references.cc b/src/libstore/references.cc
index 34dce092c..3bb297fc8 100644
--- a/src/libstore/references.cc
+++ b/src/libstore/references.cc
@@ -67,6 +67,40 @@ void RefScanSink::operator () (std::string_view data)
 }
 
 
+PathRefScanSink::PathRefScanSink(StringSet && hashes, std::map<std::string, StorePath> && backMap)
+    : RefScanSink(std::move(hashes))
+    , backMap(std::move(backMap))
+{ }
+
+PathRefScanSink PathRefScanSink::fromPaths(const StorePathSet & refs)
+{
+    StringSet hashes;
+    std::map<std::string, StorePath> backMap;
+
+    for (auto & i : refs) {
+        std::string hashPart(i.hashPart());
+        auto inserted = backMap.emplace(hashPart, i).second;
+        assert(inserted);
+        hashes.insert(hashPart);
+    }
+
+    return PathRefScanSink(std::move(hashes), std::move(backMap));
+}
+
+StorePathSet PathRefScanSink::getResultPaths()
+{
+    /* Map the hashes found back to their store paths. */
+    StorePathSet found;
+    for (auto & i : getResult()) {
+        auto j = backMap.find(i);
+        assert(j != backMap.end());
+        found.insert(j->second);
+    }
+
+    return found;
+}
+
+
 std::pair<StorePathSet, HashResult> scanForReferences(
     const std::string & path,
     const StorePathSet & refs)
@@ -82,30 +116,13 @@ StorePathSet scanForReferences(
     const Path & path,
     const StorePathSet & refs)
 {
-    StringSet hashes;
-    std::map<std::string, StorePath> backMap;
-
-    for (auto & i : refs) {
-        std::string hashPart(i.hashPart());
-        auto inserted = backMap.emplace(hashPart, i).second;
-        assert(inserted);
-        hashes.insert(hashPart);
-    }
+    PathRefScanSink refsSink = PathRefScanSink::fromPaths(refs);
+    TeeSink sink { refsSink, toTee };
 
     /* Look for the hashes in the NAR dump of the path. */
-    RefScanSink refsSink(std::move(hashes));
-    TeeSink sink { refsSink, toTee };
     dumpPath(path, sink);
 
-    /* Map the hashes found back to their store paths. */
-    StorePathSet found;
-    for (auto & i : refsSink.getResult()) {
-        auto j = backMap.find(i);
-        assert(j != backMap.end());
-        found.insert(j->second);
-    }
-
-    return found;
+    return refsSink.getResultPaths();
 }
 
 
diff --git a/src/libstore/references.hh b/src/libstore/references.hh
index a6119c861..6f381f96c 100644
--- a/src/libstore/references.hh
+++ b/src/libstore/references.hh
@@ -27,6 +27,19 @@ public:
     void operator () (std::string_view data) override;
 };
 
+class PathRefScanSink : public RefScanSink
+{
+    std::map<std::string, StorePath> backMap;
+
+    PathRefScanSink(StringSet && hashes, std::map<std::string, StorePath> && backMap);
+
+public:
+
+    static PathRefScanSink fromPaths(const StorePathSet & refs);
+
+    StorePathSet getResultPaths();
+};
+
 struct RewritingSink : Sink
 {
     std::string from, to, prev;