Redact object data in heap snapshots, with option to opt-out (#55326)

The contents of strings can contain user data which may be proprietary and emitting them in the heap snapshot makes the heap snapshot a potential vulnerability rather than a useful debugging artifact. There are likely other tweaks necessary to make heap snapshots "safe", but this is one less. --------- Co-authored-by: Nathan Daly <[email protected]> Co-authored-by: Ian Butterworth <[email protected]>
JuliaLang · Aug 25, 2024 · 383c8ef · 383c8ef
1 parent 0c8641a
commit 383c8ef
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 18 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -135,6 +135,10 @@ Standard library changes
 
 #### Profile
 
+* `Profile.take_heap_snapshot` takes a new keyword argument, `redact_data::Bool`,
+  that is `true` by default. When set, the contents of Julia objects are not emitted
+  in the heap snapshot. This currently only applies to strings. ([#55326])
+
 #### Random
 
 #### REPL

diff --git a/src/gc-heap-snapshot.cpp b/src/gc-heap-snapshot.cpp
@@ -182,6 +182,7 @@ struct HeapSnapshot {
 // global heap snapshot, mutated by garbage collector
 // when snapshotting is on.
 int gc_heap_snapshot_enabled = 0;
+int gc_heap_snapshot_redact_data = 0;
 HeapSnapshot *g_snapshot = nullptr;
 // mutex for gc-heap-snapshot.
 jl_mutex_t heapsnapshot_lock;
@@ -195,7 +196,7 @@ void _add_synthetic_root_entries(HeapSnapshot *snapshot) JL_NOTSAFEPOINT;
 
 
 JL_DLLEXPORT void jl_gc_take_heap_snapshot(ios_t *nodes, ios_t *edges,
-    ios_t *strings, ios_t *json, char all_one)
+    ios_t *strings, ios_t *json, char all_one, char redact_data)
 {
     HeapSnapshot snapshot;
     snapshot.nodes = nodes;
@@ -207,6 +208,7 @@ JL_DLLEXPORT void jl_gc_take_heap_snapshot(ios_t *nodes, ios_t *edges,
 
     // Enable snapshotting
     g_snapshot = &snapshot;
+    gc_heap_snapshot_redact_data = redact_data;
     gc_heap_snapshot_enabled = true;
 
     _add_synthetic_root_entries(&snapshot);
@@ -216,6 +218,7 @@ JL_DLLEXPORT void jl_gc_take_heap_snapshot(ios_t *nodes, ios_t *edges,
 
     // Disable snapshotting
     gc_heap_snapshot_enabled = false;
+    gc_heap_snapshot_redact_data = 0;
     g_snapshot = nullptr;
 
     jl_mutex_unlock(&heapsnapshot_lock);
@@ -328,7 +331,7 @@ size_t record_node_to_gc_snapshot(jl_value_t *a) JL_NOTSAFEPOINT
 
     if (jl_is_string(a)) {
         node_type = "String";
-        name = jl_string_data(a);
+        name = gc_heap_snapshot_redact_data ? "<redacted>" : jl_string_data(a);
         self_size = jl_string_len(a);
     }
     else if (jl_is_symbol(a)) {

diff --git a/src/gc-heap-snapshot.h b/src/gc-heap-snapshot.h
@@ -122,7 +122,7 @@ static inline void gc_heap_snapshot_record_finlist(jl_value_t *finlist, size_t i
 // Functions to call from Julia to take heap snapshot
 // ---------------------------------------------------------------------
 JL_DLLEXPORT void jl_gc_take_heap_snapshot(ios_t *nodes, ios_t *edges,
-    ios_t *strings, ios_t *json, char all_one);
+    ios_t *strings, ios_t *json, char all_one, char redact_data);
 
 
 #ifdef __cplusplus

diff --git a/stdlib/Profile/src/Profile.jl b/stdlib/Profile/src/Profile.jl
@@ -1250,8 +1250,10 @@ end
 
 
 """
-    Profile.take_heap_snapshot(filepath::String, all_one::Bool=false, streaming=false)
-    Profile.take_heap_snapshot(all_one::Bool=false; dir::String, streaming=false)
+    Profile.take_heap_snapshot(filepath::String, all_one::Bool=false;
+                               redact_data::Bool=true, streaming::Bool=false)
+    Profile.take_heap_snapshot(all_one::Bool=false; redact_data:Bool=true,
+                               dir::String=nothing, streaming::Bool=false)
 
 Write a snapshot of the heap, in the JSON format expected by the Chrome
 Devtools Heap Snapshot viewer (.heapsnapshot extension) to a file
@@ -1262,6 +1264,8 @@ full file path, or IO stream.
 If `all_one` is true, then report the size of every object as one so they can be easily
 counted. Otherwise, report the actual size.
 
+If `redact_data` is true (default), then do not emit the contents of any object.
+
 If `streaming` is true, we will stream the snapshot data out into four files, using filepath
 as the prefix, to avoid having to hold the entire snapshot in memory. This option should be
 used for any setting where your memory is constrained. These files can then be reassembled
@@ -1277,28 +1281,28 @@ backwards-compatibility) and your process is killed, note that this will always
 parts in the same directory as your provided filepath, so you can still reconstruct the
 snapshot after the fact, via `assemble_snapshot()`.
 """
-function take_heap_snapshot(filepath::AbstractString, all_one::Bool=false; streaming::Bool=false)
+function take_heap_snapshot(filepath::AbstractString, all_one::Bool=false; redact_data::Bool=true, streaming::Bool=false)
     if streaming
-        _stream_heap_snapshot(filepath, all_one)
+        _stream_heap_snapshot(filepath, all_one, redact_data)
     else
         # Support the legacy, non-streaming mode, by first streaming the parts, then
         # reassembling it after we're done.
         prefix = filepath
-        _stream_heap_snapshot(prefix, all_one)
+        _stream_heap_snapshot(prefix, all_one, redact_data)
         Profile.HeapSnapshot.assemble_snapshot(prefix, filepath)
         Profile.HeapSnapshot.cleanup_streamed_files(prefix)
     end
     return filepath
 end
-function take_heap_snapshot(io::IO, all_one::Bool=false)
+function take_heap_snapshot(io::IO, all_one::Bool=false; redact_data::Bool=true)
     # Support the legacy, non-streaming mode, by first streaming the parts to a tempdir,
     # then reassembling it after we're done.
     dir = tempdir()
     prefix = joinpath(dir, "snapshot")
-    _stream_heap_snapshot(prefix, all_one)
+    _stream_heap_snapshot(prefix, all_one, redact_data)
     Profile.HeapSnapshot.assemble_snapshot(prefix, io)
 end
-function _stream_heap_snapshot(prefix::AbstractString, all_one::Bool)
+function _stream_heap_snapshot(prefix::AbstractString, all_one::Bool, redact_data::Bool)
     # Nodes and edges are binary files
     open("$prefix.nodes", "w") do nodes
         open("$prefix.edges", "w") do edges
@@ -1311,9 +1315,9 @@ function _stream_heap_snapshot(prefix::AbstractString, all_one::Bool)
                     Base.@_lock_ios(json,
                         ccall(:jl_gc_take_heap_snapshot,
                             Cvoid,
-                            (Ptr{Cvoid},Ptr{Cvoid},Ptr{Cvoid},Ptr{Cvoid}, Cchar),
+                            (Ptr{Cvoid},Ptr{Cvoid},Ptr{Cvoid},Ptr{Cvoid}, Cchar, Cchar),
                             nodes.handle, edges.handle, strings.handle, json.handle,
-                            Cchar(all_one))
+                            Cchar(all_one), Cchar(redact_data))
                     )
                     )
                     )
@@ -1323,7 +1327,7 @@ function _stream_heap_snapshot(prefix::AbstractString, all_one::Bool)
         end
     end
 end
-function take_heap_snapshot(all_one::Bool=false; dir::Union{Nothing,S}=nothing) where {S <: AbstractString}
+function take_heap_snapshot(all_one::Bool=false; dir::Union{Nothing,S}=nothing, kwargs...) where {S <: AbstractString}
     fname = "$(getpid())_$(time_ns()).heapsnapshot"
     if isnothing(dir)
         wd = pwd()
@@ -1338,7 +1342,7 @@ function take_heap_snapshot(all_one::Bool=false; dir::Union{Nothing,S}=nothing)
     else
         fpath = joinpath(expanduser(dir), fname)
     end
-    return take_heap_snapshot(fpath, all_one)
+    return take_heap_snapshot(fpath, all_one; kwargs...)
 end
 
 """

diff --git a/stdlib/Profile/test/runtests.jl b/stdlib/Profile/test/runtests.jl
@@ -279,16 +279,31 @@ end
 
 @testset "HeapSnapshot" begin
     tmpdir = mktempdir()
+
+    # ensure that we can prevent redacting data
     fname = cd(tmpdir) do
-        read(`$(Base.julia_cmd()) --startup-file=no -e "using Profile; print(Profile.take_heap_snapshot())"`, String)
+        read(`$(Base.julia_cmd()) --startup-file=no -e "using Profile; const x = \"redact_this\"; print(Profile.take_heap_snapshot(; redact_data=false))"`, String)
     end
 
     @test isfile(fname)
 
-    open(fname) do fs
-        @test readline(fs) != ""
+    sshot = read(fname, String)
+    @test sshot != ""
+    @test contains(sshot, "redact_this")
+
+    rm(fname)
+
+    # ensure that string data is redacted by default
+    fname = cd(tmpdir) do
+        read(`$(Base.julia_cmd()) --startup-file=no -e "using Profile; const x = \"redact_this\"; print(Profile.take_heap_snapshot())"`, String)
     end
 
+    @test isfile(fname)
+
+    sshot = read(fname, String)
+    @test sshot != ""
+    @test !contains(sshot, "redact_this")
+
     rm(fname)
     rm(tmpdir, force = true, recursive = true)
 end