/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License (the "License").
 * You may not use this file except in compliance with the License.
 *
 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
 * or http://www.opensolaris.org/os/licensing.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
 * If applicable, add the following below this CDDL HEADER, with the
 * fields enclosed by brackets "[]" replaced with your own identifying
 * information: Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */
/*
 * Copyright 2006 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */
/*
 * Copyright (c) 2015, Joyent, Inc.
 */

/*
 * Routines for writing ctf data to elf files.
 */

#include <libctf_impl.h>
#include <libctf.h>
#include <gelf.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <errno.h>
#include <unistd.h>
#include <libelf.h>

static int
ctf_write_elf(ctf_file_t *fp, Elf *src, Elf *dst, int flags)
{
        GElf_Ehdr sehdr, dehdr;
        Elf_Scn *sscn, *dscn;
        Elf_Data *sdata, *ddata;
        GElf_Shdr shdr;
        int symtab_idx = -1;
        off_t new_offset = 0;
        off_t ctfnameoff = 0;
        int compress = (flags & CTF_ELFWRITE_F_COMPRESS);
        int *secxlate = NULL;
        int srcidx, dstidx, pad, i;
        int curnmoff = 0;
        int changing = 0;
        int ret;
        size_t nshdr, nphdr, strndx;
        void *strdatabuf = NULL, *symdatabuf = NULL;
        size_t strdatasz = 0, symdatasz = 0;

        void *cdata = NULL;
        size_t elfsize, asize;

        if ((flags & ~(CTF_ELFWRITE_F_COMPRESS)) != 0) {
                ret = ctf_set_errno(fp, EINVAL);
                goto out;
        }

        if (gelf_newehdr(dst, gelf_getclass(src)) == 0) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }
        if (gelf_getehdr(src, &sehdr) == NULL) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }
        (void) memcpy(&dehdr, &sehdr, sizeof (GElf_Ehdr));
        if (gelf_update_ehdr(dst, &dehdr) == 0) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }

        /*
         * Use libelf to get the number of sections and the string section to
         * deal with ELF files that may have a large number of sections. We just
         * always use this to make our live easier.
         */
        if (elf_getphdrnum(src, &nphdr) != 0) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }
        if (elf_getshdrnum(src, &nshdr) != 0) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }
        if (elf_getshdrstrndx(src, &strndx) != 0) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }

        /*
         * Neither the existing debug sections nor the SUNW_ctf sections (new or
         * existing) are SHF_ALLOC'd, so they won't be in areas referenced by
         * program headers.  As such, we can just blindly copy the program
         * headers from the existing file to the new file.
         */
        if (nphdr != 0) {
                (void) elf_flagelf(dst, ELF_C_SET, ELF_F_LAYOUT);
                if (gelf_newphdr(dst, nphdr) == 0) {
                        ret = ctf_set_errno(fp, ECTF_ELF);
                        goto out;
                }

                for (i = 0; i < nphdr; i++) {
                        GElf_Phdr phdr;

                        if (gelf_getphdr(src, i, &phdr) == NULL) {
                                ret = ctf_set_errno(fp, ECTF_ELF);
                                goto out;
                        }
                        if (gelf_update_phdr(dst, i, &phdr) == 0) {
                                ret = ctf_set_errno(fp, ECTF_ELF);
                                goto out;
                        }
                }
        }

        secxlate = ctf_alloc(sizeof (int) * nshdr);
        for (srcidx = dstidx = 0; srcidx < nshdr; srcidx++) {
                Elf_Scn *scn = elf_getscn(src, srcidx);
                GElf_Shdr shdr;
                char *sname;

                if (gelf_getshdr(scn, &shdr) == NULL) {
                        ret = ctf_set_errno(fp, ECTF_ELF);
                        goto out;
                }
                sname = elf_strptr(src, strndx, shdr.sh_name);
                if (sname == NULL) {
                        ret = ctf_set_errno(fp, ECTF_ELF);
                        goto out;
                }

                if (strcmp(sname, CTF_ELF_SCN_NAME) == 0) {
                        secxlate[srcidx] = -1;
                } else {
                        secxlate[srcidx] = dstidx++;
                        curnmoff += strlen(sname) + 1;
                }

                new_offset = (off_t)dehdr.e_phoff;
        }

        for (srcidx = 1; srcidx < nshdr; srcidx++) {
                char *sname;

                sscn = elf_getscn(src, srcidx);
                if (gelf_getshdr(sscn, &shdr) == NULL) {
                        ret = ctf_set_errno(fp, ECTF_ELF);
                        goto out;
                }

                if (secxlate[srcidx] == -1) {
                        changing = 1;
                        continue;
                }

                dscn = elf_newscn(dst);
                if (dscn == NULL) {
                        ret = ctf_set_errno(fp, ECTF_ELF);
                        goto out;
                }

                /*
                 * If this file has program headers, we need to explicitly lay
                 * out sections.  If none of the sections prior to this one have
                 * been removed, then we can just use the existing location.  If
                 * one or more sections have been changed, then we need to
                 * adjust this one to avoid holes.
                 */
                if (changing && nphdr != 0) {
                        pad = new_offset % shdr.sh_addralign;

                        if (pad != 0)
                                new_offset += shdr.sh_addralign - pad;
                        shdr.sh_offset = new_offset;
                }

                shdr.sh_link = secxlate[shdr.sh_link];

                if (shdr.sh_type == SHT_REL || shdr.sh_type == SHT_RELA)
                        shdr.sh_info = secxlate[shdr.sh_info];

                sname = elf_strptr(src, strndx, shdr.sh_name);
                if (sname == NULL) {
                        ret = ctf_set_errno(fp, ECTF_ELF);
                        goto out;
                }
                if ((sdata = elf_getdata(sscn, NULL)) == NULL) {
                        ret = ctf_set_errno(fp, ECTF_ELF);
                        goto out;
                }
                if ((ddata = elf_newdata(dscn)) == NULL) {
                        ret = ctf_set_errno(fp, ECTF_ELF);
                        goto out;
                }
                bcopy(sdata, ddata, sizeof (Elf_Data));

                if (srcidx == strndx) {
                        char seclen = strlen(CTF_ELF_SCN_NAME);

                        strdatasz = ddata->d_size + shdr.sh_size +
                            seclen + 1;
                        ddata->d_buf = strdatabuf = ctf_alloc(strdatasz);
                        if (ddata->d_buf == NULL) {
                                ret = ctf_set_errno(fp, ECTF_ELF);
                                goto out;
                        }
                        bcopy(sdata->d_buf, ddata->d_buf, shdr.sh_size);
                        (void) strcpy((caddr_t)ddata->d_buf + shdr.sh_size,
                            CTF_ELF_SCN_NAME);
                        ctfnameoff = (off_t)shdr.sh_size;
                        shdr.sh_size += seclen + 1;
                        ddata->d_size += seclen + 1;

                        if (nphdr != 0)
                                changing = 1;
                }

                if (shdr.sh_type == SHT_SYMTAB && shdr.sh_entsize != 0) {
                        int nsym = shdr.sh_size / shdr.sh_entsize;

                        symtab_idx = secxlate[srcidx];

                        symdatasz = shdr.sh_size;
                        ddata->d_buf = symdatabuf = ctf_alloc(symdatasz);
                        if (ddata->d_buf == NULL) {
                                ret = ctf_set_errno(fp, ECTF_ELF);
                                goto out;
                        }
                        (void) bcopy(sdata->d_buf, ddata->d_buf, shdr.sh_size);

                        for (i = 0; i < nsym; i++) {
                                GElf_Sym sym;
                                short newscn;

                                (void) gelf_getsym(ddata, i, &sym);

                                if (sym.st_shndx >= SHN_LORESERVE)
                                        continue;

                                if ((newscn = secxlate[sym.st_shndx]) !=
                                    sym.st_shndx) {
                                        sym.st_shndx =
                                            (newscn == -1 ? 1 : newscn);

                                        if (gelf_update_sym(ddata, i, &sym) ==
                                            0) {
                                                ret = ctf_set_errno(fp,
                                                    ECTF_ELF);
                                                goto out;
                                        }
                                }
                        }
                }

                if (gelf_update_shdr(dscn, &shdr) == 0) {
                        ret = ctf_set_errno(fp, ECTF_ELF);
                        goto out;
                }

                new_offset = (off_t)shdr.sh_offset;
                if (shdr.sh_type != SHT_NOBITS)
                        new_offset += shdr.sh_size;
        }

        if (symtab_idx == -1) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }

        /* Add the ctf section */
        if ((dscn = elf_newscn(dst)) == NULL) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }
        if (gelf_getshdr(dscn, &shdr) == NULL) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }
        shdr.sh_name = ctfnameoff;
        shdr.sh_type = SHT_PROGBITS;
        shdr.sh_size = fp->ctf_size;
        shdr.sh_link = symtab_idx;
        shdr.sh_addralign = 4;
        if (changing && nphdr != 0) {
                pad = new_offset % shdr.sh_addralign;

                if (pad)
                        new_offset += shdr.sh_addralign - pad;

                shdr.sh_offset = new_offset;
                new_offset += shdr.sh_size;
        }

        if ((ddata = elf_newdata(dscn)) == NULL) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }

        if (compress != 0) {
                int err;

                if (ctf_zopen(&err) == NULL) {
                        ret = ctf_set_errno(fp, err);
                        goto out;
                }

                if ((err = ctf_compress(fp, &cdata, &asize, &elfsize)) != 0) {
                        ret = ctf_set_errno(fp, err);
                        goto out;
                }
                ddata->d_buf = cdata;
                ddata->d_size = elfsize;
        } else {
                ddata->d_buf = (void *)fp->ctf_base;
                ddata->d_size = fp->ctf_size;
        }
        ddata->d_align = shdr.sh_addralign;

        if (gelf_update_shdr(dscn, &shdr) == 0) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }

        /* update the section header location */
        if (nphdr != 0) {
                size_t align = gelf_fsize(dst, ELF_T_ADDR, 1, EV_CURRENT);
                size_t r = new_offset % align;

                if (r)
                        new_offset += align - r;

                dehdr.e_shoff = new_offset;
        }

        /* commit to disk */
        if (sehdr.e_shstrndx == SHN_XINDEX)
                dehdr.e_shstrndx = SHN_XINDEX;
        else
                dehdr.e_shstrndx = secxlate[sehdr.e_shstrndx];
        if (gelf_update_ehdr(dst, &dehdr) == 0) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }
        if (elf_update(dst, ELF_C_WRITE) < 0) {
                ret = ctf_set_errno(fp, ECTF_ELF);
                goto out;
        }

        ret = 0;

out:
        if (strdatabuf != NULL)
                ctf_free(strdatabuf, strdatasz);
        if (symdatabuf != NULL)
                ctf_free(symdatabuf, symdatasz);
        if (cdata != NULL)
                ctf_data_free(cdata, fp->ctf_size);
        if (secxlate != NULL)
                ctf_free(secxlate, sizeof (int) * nshdr);

        return (ret);
}

int
ctf_elffdwrite(ctf_file_t *fp, int ifd, int ofd, int flags)
{
        int ret;
        Elf *ielf, *oelf;

        (void) elf_version(EV_CURRENT);
        if ((ielf = elf_begin(ifd, ELF_C_READ, NULL)) == NULL)
                return (ctf_set_errno(fp, ECTF_ELF));

        if ((oelf = elf_begin(ofd, ELF_C_WRITE, NULL)) == NULL)
                return (ctf_set_errno(fp, ECTF_ELF));

        ret = ctf_write_elf(fp, ielf, oelf, flags);

        (void) elf_end(ielf);
        (void) elf_end(oelf);

        return (ret);
}

int
ctf_elfwrite(ctf_file_t *fp, const char *input, const char *output, int flags)
{
        struct stat st;
        int ifd, ofd, ret;

        if ((ifd = open(input, O_RDONLY)) < 0)
                return (ctf_set_errno(fp, errno));

        if (fstat(ifd, &st) < 0)
                return (ctf_set_errno(fp, errno));

        if ((ofd = open(output, O_RDWR | O_CREAT | O_TRUNC, st.st_mode)) < 0)
                return (ctf_set_errno(fp, errno));

        ret = ctf_elffdwrite(fp, ifd, ofd, flags);

        if (close(ifd) != 0 && ret == 0)
                ret = ctf_set_errno(fp, errno);
        if (close(ofd) != 0 && ret == 0)
                ret = ctf_set_errno(fp, errno);

        return (ret);
}