1 files changed, 123 insertions, 0 deletions
diff --git a/infra/base-images/base-runner/profraw_update.py b/infra/base-images/base-runner/profraw_update.py
new file mode 100644
index 000000000..408b5fb93
--- /dev/null
+++ b/infra/base-images/base-runner/profraw_update.py
@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+################################################################################
+"""Helper script for upgrading a profraw file to latest version."""
+
+from collections import namedtuple
+import struct
+import subprocess
+import sys
+
+HeaderGeneric = namedtuple('HeaderGeneric', 'magic version')
+HeaderVersion7 = namedtuple(
+    'HeaderVersion7',
+    'BinaryIdsSize DataSize PaddingBytesBeforeCounters CountersSize \
+    PaddingBytesAfterCounters NamesSize CountersDelta NamesDelta ValueKindLast')
+
+PROFRAW_MAGIC = 0xff6c70726f667281
+
+
+def relativize_address(data, offset, databegin, sect_prf_cnts, sect_prf_data):
+  """Turns an absolute offset into a relative one."""
+  value = struct.unpack('Q', data[offset:offset + 8])[0]
+  if sect_prf_cnts <= value < sect_prf_data:
+    # If the value is an address in the right section, make it relative.
+    value = (value - databegin) & 0xffffffffffffffff
+    value = struct.pack('Q', value)
+    for i in range(8):
+      data[offset + i] = value[i]
+
+
+def upgrade(data, sect_prf_cnts, sect_prf_data):
+  """Upgrades profraw data, knowing the sections addresses."""
+  generic_header = HeaderGeneric._make(struct.unpack('QQ', data[:16]))
+  if generic_header.magic != PROFRAW_MAGIC:
+    raise Exception('Bad magic.')
+  if generic_header.version == 5:
+    generic_header = generic_header._replace(version=7)
+    # Upgrade from version 5 to 7 by adding binaryids field.
+    data = struct.pack('QQ', generic_header) + struct.pack('Q', 0) + data[16:]
+  if generic_header.version < 7:
+    raise Exception('Unhandled version.')
+  v7_header = HeaderVersion7._make(struct.unpack('QQQQQQQQQ', data[16:88]))
+
+  if v7_header.BinaryIdsSize % 8 != 0:
+    # Adds padding for binary ids.
+    # cf commit b9f547e8e51182d32f1912f97a3e53f4899ea6be
+    # cf https://reviews.llvm.org/D110365
+    padlen = 8 - (v7_header.BinaryIdsSize % 8)
+    v7_header = v7_header._replace(BinaryIdsSize=v7_header.BinaryIdsSize +
+                                   padlen)
+    data = data[:16] + struct.pack('Q', v7_header.BinaryIdsSize) + data[24:]
+    data = data[:88 + v7_header.BinaryIdsSize] + bytes(
+        padlen) + data[88 + v7_header.BinaryIdsSize:]
+
+  if v7_header.CountersDelta != sect_prf_cnts - sect_prf_data:
+    # Rust linking seems to add an offset...
+    sect_prf_data = v7_header.CountersDelta - sect_prf_cnts + sect_prf_data
+    sect_prf_cnts = v7_header.CountersDelta
+
+  dataref = sect_prf_data
+  relativize_address(data, 64, dataref, sect_prf_cnts, sect_prf_data)
+
+  offset = 88 + v7_header.BinaryIdsSize
+  # This also works for C+Rust binaries compiled with
+  # clang-14/rust-nightly-clang-13.
+  for _ in range(v7_header.DataSize):
+    # 16 is the offset of CounterPtr in ProfrawData structure.
+    relativize_address(data, offset + 16, dataref, sect_prf_cnts, sect_prf_data)
+    # We need this because of CountersDelta -= sizeof(*SrcData);
+    # seen in __llvm_profile_merge_from_buffer.
+    dataref += 44 + 2 * (v7_header.ValueKindLast + 1)
+    # This is the size of one ProfrawData structure.
+    offset += 44 + 2 * (v7_header.ValueKindLast + 1)
+
+  return data
+
+
+def main():
+  """Helper script for upgrading a profraw file to latest version."""
+  if len(sys.argv) != 4:
+    sys.stderr.write('Usage: %s <binary> <profraw> <output>\n' % sys.argv[0])
+    return 1
+
+  # First find llvm profile sections addresses in the elf, quick and dirty.
+  process = subprocess.Popen(['readelf', '-S', sys.argv[1]],
+                             stdout=subprocess.PIPE)
+  output, err = process.communicate()
+  if err:
+    print('readelf failed')
+    return 2
+  for line in iter(output.split(b'\n')):
+    if b'__llvm_prf_cnts' in line:
+      sect_prf_cnts = int(line.split()[4], 16)
+    elif b'__llvm_prf_data' in line:
+      sect_prf_data = int(line.split()[4], 16)
+
+  # Then open and read the input profraw file.
+  with open(sys.argv[2], 'rb') as input_file:
+    profraw_base = bytearray(input_file.read())
+  # Do the upgrade, returning a bytes object.
+  profraw_latest = upgrade(profraw_base, sect_prf_cnts, sect_prf_data)
+  # Write the output to the file given to the command line.
+  with open(sys.argv[3], 'wb') as output_file:
+    output_file.write(profraw_latest)
+
+  return 0
+
+
+if __name__ == '__main__':
+  sys.exit(main())