1 files changed, 113 insertions, 24 deletions

diff --git a/scripts/recordmcount.c b/scripts/recordmcount.c
index 698768bdc581..301d70b0174f 100644
--- a/scripts/recordmcount.c
+++ b/scripts/recordmcount.c
@@ -48,12 +48,17 @@
 
 static int fd_map;      /* File descriptor for file being modified. */
 static int mmap_failed; /* Boolean flag. */
-static void *ehdr_curr; /* current ElfXX_Ehdr *  for resource cleanup */
 static char gpfx;       /* prefix for global symbol name (sometimes '_') */
 static struct stat sb;  /* Remember .st_size, etc. */
 static jmp_buf jmpenv;  /* setjmp/longjmp per-file error escape */
 static const char *altmcount;   /* alternate mcount symbol name */
 static int warn_on_notrace_sect; /* warn when section has mcount not being recorded */
+static void *file_map;  /* pointer of the mapped file */
+static void *file_end;  /* pointer to the end of the mapped file */
+static int file_updated; /* flag to state file was changed */
+static void *file_ptr;  /* current file pointer location */
+static void *file_append; /* added to the end of the file */
+static size_t file_append_size; /* how much is added to end of file */
 
 /* setjmp() return values */
 enum {
@@ -67,10 +72,14 @@ static void
 cleanup(void)
 {
         if (!mmap_failed)
-                munmap(ehdr_curr, sb.st_size);
+                munmap(file_map, sb.st_size);
         else
-                free(ehdr_curr);
-        close(fd_map);
+                free(file_map);
+        file_map = NULL;
+        free(file_append);
+        file_append = NULL;
+        file_append_size = 0;
+        file_updated = 0;
 }
 
 static void __attribute__((noreturn))
@@ -92,12 +101,22 @@ succeed_file(void)
 static off_t
 ulseek(int const fd, off_t const offset, int const whence)
 {
-        off_t const w = lseek(fd, offset, whence);
-        if (w == (off_t)-1) {
-                perror("lseek");
+        switch (whence) {
+        case SEEK_SET:
+                file_ptr = file_map + offset;
+                break;
+        case SEEK_CUR:
+                file_ptr += offset;
+                break;
+        case SEEK_END:
+                file_ptr = file_map + (sb.st_size - offset);
+                break;
+        }
+        if (file_ptr < file_map) {
+                fprintf(stderr, "lseek: seek before file\n");
                 fail_file();
         }
-        return w;
+        return file_ptr - file_map;
 }
 
 static size_t
@@ -114,12 +133,38 @@ uread(int const fd, void *const buf, size_t const count)
 static size_t
 uwrite(int const fd, void const *const buf, size_t const count)
 {
-        size_t const n = write(fd, buf, count);
-        if (n != count) {
-                perror("write");
-                fail_file();
+        size_t cnt = count;
+        off_t idx = 0;
+
+        file_updated = 1;
+
+        if (file_ptr + count >= file_end) {
+                off_t aoffset = (file_ptr + count) - file_end;
+
+                if (aoffset > file_append_size) {
+                        file_append = realloc(file_append, aoffset);
+                        file_append_size = aoffset;
+                }
+                if (!file_append) {
+                        perror("write");
+                        fail_file();
+                }
+                if (file_ptr < file_end) {
+                        cnt = file_end - file_ptr;
+                } else {
+                        cnt = 0;
+                        idx = aoffset - count;
+                }
         }
-        return n;
+
+        if (cnt)
+                memcpy(file_ptr, buf, cnt);
+
+        if (cnt < count)
+                memcpy(file_append + idx, buf + cnt, count - cnt);
+
+        file_ptr += count;
+        return count;
 }
 
 static void *
@@ -192,9 +237,7 @@ static int make_nop_arm64(void *map, size_t const offset)
  */
 static void *mmap_file(char const *fname)
 {
-        void *addr;
-
-        fd_map = open(fname, O_RDWR);
+        fd_map = open(fname, O_RDONLY);
         if (fd_map < 0 || fstat(fd_map, &sb) < 0) {
                 perror(fname);
                 fail_file();
@@ -203,15 +246,58 @@ static void *mmap_file(char const *fname)
                 fprintf(stderr, "not a regular file: %s\n", fname);
                 fail_file();
         }
-        addr = mmap(0, sb.st_size, PROT_READ|PROT_WRITE, MAP_PRIVATE,
-                    fd_map, 0);
+        file_map = mmap(0, sb.st_size, PROT_READ|PROT_WRITE, MAP_PRIVATE,
+                        fd_map, 0);
         mmap_failed = 0;
-        if (addr == MAP_FAILED) {
+        if (file_map == MAP_FAILED) {
                 mmap_failed = 1;
-                addr = umalloc(sb.st_size);
-                uread(fd_map, addr, sb.st_size);
+                file_map = umalloc(sb.st_size);
+                uread(fd_map, file_map, sb.st_size);
+        }
+        close(fd_map);
+
+        file_end = file_map + sb.st_size;
+
+        return file_map;
+}
+
+static void write_file(const char *fname)
+{
+        char tmp_file[strlen(fname) + 4];
+        size_t n;
+
+        if (!file_updated)
+                return;
+
+        sprintf(tmp_file, "%s.rc", fname);
+
+        /*
+         * After reading the entire file into memory, delete it
+         * and write it back, to prevent weird side effects of modifying
+         * an object file in place.
+         */
+        fd_map = open(tmp_file, O_WRONLY | O_TRUNC | O_CREAT, sb.st_mode);
+        if (fd_map < 0) {
+                perror(fname);
+                fail_file();
+        }
+        n = write(fd_map, file_map, sb.st_size);
+        if (n != sb.st_size) {
+                perror("write");
+                fail_file();
+        }
+        if (file_append_size) {
+                n = write(fd_map, file_append, file_append_size);
+                if (n != file_append_size) {
+                        perror("write");
+                        fail_file();
+                }
+        }
+        close(fd_map);
+        if (rename(tmp_file, fname) < 0) {
+                perror(fname);
+                fail_file();
         }
-        return addr;
 }
 
 /* w8rev, w8nat, ...: Handle endianness. */
@@ -318,7 +404,6 @@ do_file(char const *const fname)
         Elf32_Ehdr *const ehdr = mmap_file(fname);
         unsigned int reltype = 0;
 
-        ehdr_curr = ehdr;
         w = w4nat;
         w2 = w2nat;
         w8 = w8nat;
@@ -441,6 +526,7 @@ do_file(char const *const fname)
         }
         }  /* end switch */
 
+        write_file(fname);
         cleanup();
 }
 
@@ -493,11 +579,14 @@ main(int argc, char *argv[])
                 case SJ_SETJMP:    /* normal sequence */
                         /* Avoid problems if early cleanup() */
                         fd_map = -1;
-                        ehdr_curr = NULL;
                         mmap_failed = 1;
+                        file_map = NULL;
+                        file_ptr = NULL;
+                        file_updated = 0;
                         do_file(file);
                         break;
                 case SJ_FAIL:    /* error in do_file or below */
+                        sprintf("%s: failed\n", file);
                         ++n_error;
                         break;
                 case SJ_SUCCEED:    /* premature success */