From 100922b650cbe114f165c7269a9b6d2ad39d74f7 Mon Sep 17 00:00:00 2001
From: Rui Ueyama <ruiu@cs.stanford.edu>
Date: Fri, 15 Apr 2022 07:23:50 +0800
Subject: [PATCH] [ELF] Fix section file offsets

Previously, if a section has a very large alignment requirement,
that section and the following sections may get wrong file offsets.

Fixes https://github.com/rui314/mold/issues/405
---
 elf/mold.h                      |  9 ++---
 elf/output-chunks.cc            | 28 ++++++++-------
 elf/passes.cc                   | 62 +++++++++++++++++++++------------
 test/elf/large-alignment-dso.sh | 50 ++++++++++++++++++++++++++
 4 files changed, 108 insertions(+), 41 deletions(-)
 create mode 100755 test/elf/large-alignment-dso.sh

diff --git a/elf/mold.h b/elf/mold.h
index 6cafd301..b36446cb 100644
--- a/elf/mold.h
+++ b/elf/mold.h
@@ -336,9 +336,6 @@ void report_undef(Context<E> &ctx, InputFile<E> &file, Symbol<E> &sym);
 template <typename E>
 bool is_relro(Context<E> &ctx, Chunk<E> *chunk);
 
-template <typename E>
-bool separate_page(Context<E> &ctx, Chunk<E> *a, Chunk<E> *b);
-
 // Chunk represents a contiguous region in an output file.
 template <typename E>
 class Chunk {
@@ -353,8 +350,9 @@ public:
   virtual void update_shdr(Context<E> &ctx) {}
 
   std::string_view name;
-  i64 shndx = 0;
   ElfShdr<E> shdr = {};
+  i64 shndx = 0;
+  i64 extra_addralign = 1;
 
 protected:
   Chunk() { shdr.sh_addralign = 1; }
@@ -968,9 +966,6 @@ private:
 
 bool is_c_identifier(std::string_view name);
 
-template <typename E>
-std::vector<ElfPhdr<E>> create_phdr(Context<E> &ctx);
-
 //
 // input-files.cc
 //
diff --git a/elf/output-chunks.cc b/elf/output-chunks.cc
index 6ea7ff48..abb278ef 100644
--- a/elf/output-chunks.cc
+++ b/elf/output-chunks.cc
@@ -166,7 +166,7 @@ bool is_relro(Context<E> &ctx, Chunk<E> *chunk) {
 }
 
 template <typename E>
-bool separate_page(Context<E> &ctx, Chunk<E> *x, Chunk<E> *y) {
+static bool separate_page(Context<E> &ctx, Chunk<E> *x, Chunk<E> *y) {
   if (ctx.arg.z_relro && is_relro(ctx, x) != is_relro(ctx, y))
     return true;
 
@@ -182,7 +182,7 @@ bool separate_page(Context<E> &ctx, Chunk<E> *x, Chunk<E> *y) {
 }
 
 template <typename E>
-std::vector<ElfPhdr<E>> create_phdr(Context<E> &ctx) {
+static std::vector<ElfPhdr<E>> create_phdr(Context<E> &ctx) {
   std::vector<ElfPhdr<E>> vec;
 
   auto define = [&](u64 type, u64 flags, i64 min_align, auto &chunk) {
@@ -222,6 +222,10 @@ std::vector<ElfPhdr<E>> create_phdr(Context<E> &ctx) {
     return (shdr.sh_type == SHT_NOTE) && (shdr.sh_flags & SHF_ALLOC);
   };
 
+  // Clear previous results so that this function is idempotent.
+  for (Chunk<E> *chunk : ctx.chunks)
+    chunk->extra_addralign = 1;
+
   // Create a PT_PHDR for the program header itself.
   if (ctx.phdr)
     define(PT_PHDR, PF_R, E::word_size, ctx.phdr);
@@ -262,14 +266,14 @@ std::vector<ElfPhdr<E>> create_phdr(Context<E> &ctx) {
 
       if (!is_bss(first))
         while (i < end && !is_bss(chunks[i]) &&
-               to_phdr_flags(ctx, chunks[i]) == flags &&
-               chunks[i]->shdr.sh_addr - first->shdr.sh_addr ==
-               chunks[i]->shdr.sh_offset - first->shdr.sh_offset)
+               to_phdr_flags(ctx, chunks[i]) == flags)
           append(chunks[i++]);
 
       while (i < end && is_bss(chunks[i]) &&
              to_phdr_flags(ctx, chunks[i]) == flags)
         append(chunks[i++]);
+
+      first->extra_addralign = vec.back().p_align;
     }
   }
 
@@ -312,15 +316,17 @@ std::vector<ElfPhdr<E>> create_phdr(Context<E> &ctx) {
         continue;
 
       define(PT_GNU_RELRO, PF_R, 1, ctx.chunks[i]);
+      ctx.chunks[i]->extra_addralign = ctx.page_size;
+
       i++;
       while (i < ctx.chunks.size() && is_relro(ctx, ctx.chunks[i]))
         append(ctx.chunks[i++]);
 
-      // RELRO works on page granularity, so align it up to the next
-      // page boundary.
-      assert(i == ctx.chunks.size() ||
-             ctx.chunks[i]->shdr.sh_addr % ctx.page_size == 0);
+      // RELRO works on page granularity, so align both ends to
+      // the page size.
       vec.back().p_memsz = align_to(vec.back().p_memsz, ctx.page_size);
+      if (i < ctx.chunks.size())
+        ctx.chunks[i]->extra_addralign = ctx.page_size;
     }
   }
 
@@ -2771,9 +2777,7 @@ void RelocSection<E>::copy_buf(Context<E> &ctx) {
   template class GnuCompressedSection<E>;                               \
   template class RelocSection<E>;                                       \
   template i64 BuildId::size(Context<E> &) const;                       \
-  template bool is_relro(Context<E> &, Chunk<E> *);                     \
-  template bool separate_page(Context<E> &, Chunk<E> *, Chunk<E> *);    \
-  template std::vector<ElfPhdr<E>> create_phdr(Context<E> &)
+  template bool is_relro(Context<E> &, Chunk<E> *);
 
 INSTANTIATE(X86_64);
 INSTANTIATE(I386);
diff --git a/elf/passes.cc b/elf/passes.cc
index 7eb85fb9..78255631 100644
--- a/elf/passes.cc
+++ b/elf/passes.cc
@@ -1267,15 +1267,6 @@ i64 get_section_rank(Context<E> &ctx, Chunk<E> *chunk) {
          (!relro << 16) | (is_bss << 15);
 }
 
-// Returns the smallest number n such that
-// val <= n and n % align == skew % align.
-inline u64 align_with_skew(u64 val, u64 align, u64 skew) {
-  skew = skew % align;
-  u64 n = align_to(val + align - skew, align) - align + skew;
-  assert(val <= n && n < val + align && n % align == skew % align);
-  return n;
-}
-
 template <typename E>
 static bool is_tbss(Chunk<E> *chunk) {
   return (chunk->shdr.sh_type == SHT_NOBITS) && (chunk->shdr.sh_flags & SHF_TLS);
@@ -1286,6 +1277,10 @@ template <typename E>
 i64 do_set_osec_offsets(Context<E> &ctx) {
   std::vector<Chunk<E> *> &chunks = ctx.chunks;
 
+  auto alignment = [](Chunk<E> *chunk) {
+    return std::max<i64>(chunk->extra_addralign, chunk->shdr.sh_addralign);
+  };
+
   // Assign virtual addresses
   u64 addr = ctx.arg.image_base;
   for (i64 i = 0; i < chunks.size(); i++) {
@@ -1296,15 +1291,12 @@ i64 do_set_osec_offsets(Context<E> &ctx) {
         it != ctx.arg.section_start.end())
       addr = it->second;
 
-    if (i > 0 && separate_page(ctx, chunks[i - 1], chunks[i]))
-      addr = align_to(addr, ctx.page_size);
-
     if (is_tbss(chunks[i])) {
       chunks[i]->shdr.sh_addr = addr;
       continue;
     }
 
-    addr = align_to(addr, chunks[i]->shdr.sh_addralign);
+    addr = align_to(addr, alignment(chunks[i]));
     chunks[i]->shdr.sh_addr = addr;
     addr += chunks[i]->shdr.sh_size;
   }
@@ -1321,7 +1313,7 @@ i64 do_set_osec_offsets(Context<E> &ctx) {
     if (is_tbss(chunks[i])) {
       u64 addr = chunks[i]->shdr.sh_addr;
       for (; i < chunks.size() && is_tbss(chunks[i]); i++) {
-        addr = align_to(addr, chunks[i]->shdr.sh_addralign);
+        addr = align_to(addr, alignment(chunks[i]));
         chunks[i]->shdr.sh_addr = addr;
         addr += chunks[i]->shdr.sh_size;
       }
@@ -1330,16 +1322,42 @@ i64 do_set_osec_offsets(Context<E> &ctx) {
     }
   }
 
-  // Assign file offsets
+  // Assign file offsets to memory-allocated sections.
   u64 fileoff = 0;
-  for (Chunk<E> *chunk : chunks) {
-    if (chunk->shdr.sh_type == SHT_NOBITS) {
-      chunk->shdr.sh_offset = fileoff;
-    } else {
-      fileoff = align_with_skew(fileoff, ctx.page_size, chunk->shdr.sh_addr);
-      chunk->shdr.sh_offset = fileoff;
-      fileoff += chunk->shdr.sh_size;
+  i64 i = 0;
+
+  while (i < chunks.size() && (chunks[i]->shdr.sh_flags & SHF_ALLOC)) {
+    Chunk<E> &first = *chunks[i];
+    assert(first.shdr.sh_type != SHT_NOBITS);
+
+    fileoff = align_to(fileoff, alignment(&first));
+
+    u64 end = fileoff;
+    while (i < chunks.size() && (chunks[i]->shdr.sh_flags & SHF_ALLOC) &&
+           chunks[i]->shdr.sh_type != SHT_NOBITS) {
+      // The addresses may not increase monotonically if a user uses
+      // --start-sections.
+      if (chunks[i]->shdr.sh_addr < first.shdr.sh_addr)
+        break;
+
+      chunks[i]->shdr.sh_offset =
+        fileoff + chunks[i]->shdr.sh_addr - first.shdr.sh_addr;
+      end = chunks[i]->shdr.sh_offset + chunks[i]->shdr.sh_size;
+      i++;
     }
+
+    fileoff = end;
+
+    while (i < chunks.size() && (chunks[i]->shdr.sh_flags & SHF_ALLOC) &&
+           chunks[i]->shdr.sh_type == SHT_NOBITS)
+      i++;
+  }
+
+  // Assign file offsets to non-memory-allocated sections.
+  for (; i < chunks.size(); i++) {
+    fileoff = align_to(fileoff, chunks[i]->shdr.sh_addralign);
+    chunks[i]->shdr.sh_offset = fileoff;
+    fileoff += chunks[i]->shdr.sh_size;
   }
   return fileoff;
 }
diff --git a/test/elf/large-alignment-dso.sh b/test/elf/large-alignment-dso.sh
new file mode 100755
index 00000000..366b8b49
--- /dev/null
+++ b/test/elf/large-alignment-dso.sh
@@ -0,0 +1,50 @@
+#!/bin/bash
+export LC_ALL=C
+set -e
+CC="${CC:-cc}"
+CXX="${CXX:-c++}"
+GCC="${GCC:-gcc}"
+GXX="${GXX:-g++}"
+OBJDUMP="${OBJDUMP:-objdump}"
+MACHINE="${MACHINE:-$(uname -m)}"
+testname=$(basename "$0" .sh)
+echo -n "Testing $testname ... "
+cd "$(dirname "$0")"/../..
+mold="$(pwd)/mold"
+t=out/test/elf/$testname
+mkdir -p $t
+
+[ $MACHINE = i386 -o $MACHINE = arm ] && { echo skipped; exit; }
+
+cat <<EOF | $CC -o $t/a.o -c -xc - -ffunction-sections -fPIC
+#include <stdio.h>
+#include <stdint.h>
+
+void hello() __attribute__((aligned(32768), section(".hello")));
+void world() __attribute__((aligned(32768), section(".world")));
+
+void hello() {
+  printf("Hello");
+}
+
+void world() {
+  printf(" world");
+}
+
+void greet() {
+  hello();
+  world();
+}
+EOF
+
+$CC -B. -o $t/b.so $t/a.o -shared
+
+cat <<EOF | $CC -o $t/c.o -c -xc -
+void greet();
+int main() { greet(); }
+EOF
+
+$CC -B. -o $t/exe $t/c.o $t/b.so
+$QEMU $t/exe | grep -q 'Hello world'
+
+echo OK