elfwriter: add WriteSectionHeaders (#3666)

Add ability to write out section headers to elfwriter.

pkg/elfwriter/writer.go

@@ -2,7 +2,6 @@
 // contents in memory at any one time.
 // This package is incomplete, only features needed to write core files are
 // implemented, notably missing:
-// - section headers
 // - program headers at the beginning of the file
 
 package elfwriter
@@ -22,12 +21,17 @@ type WriteCloserSeeker interface {
 
 // Writer writes ELF files.
 type Writer struct {
-	w     WriteCloserSeeker
-	Err   error
-	Progs []*elf.ProgHeader
+	w        WriteCloserSeeker
+	Err      error
+	Progs    []*elf.ProgHeader
+	Sections []*elf.SectionHeader
 
 	seekProgHeader int64
 	seekProgNum    int64
+
+	seekSectionHeader int64
+	seekSectionNum    int64
+	seekShstrndx      int64
 }
 
 type Note struct {
@@ -41,6 +45,7 @@ func New(w WriteCloserSeeker, fhdr *elf.FileHeader) *Writer {
 	const (
 		ehsize    = 64
 		phentsize = 56
+		shentsize = 64
 	)
 
 	if seek, _ := w.Seek(0, io.SeekCurrent); seek != 0 {
@@ -65,15 +70,18 @@ func New(w WriteCloserSeeker, fhdr *elf.FileHeader) *Writer {
 	r.u32(uint32(fhdr.Version)) // e_version
 	r.u64(0)                    // e_entry
 	r.seekProgHeader = r.Here()
-	r.u64(0)         // e_phoff
+	r.u64(0) // e_phoff
+	r.seekSectionHeader = r.Here()
 	r.u64(0)         // e_shoff
 	r.u32(0)         // e_flags
 	r.u16(ehsize)    // e_ehsize
 	r.u16(phentsize) // e_phentsize
 	r.seekProgNum = r.Here()
-	r.u16(0)                     // e_phnum
-	r.u16(0)                     // e_shentsize
-	r.u16(0)                     // e_shnum
+	r.u16(0)         // e_phnum
+	r.u16(shentsize) // e_shentsize
+	r.seekSectionNum = r.Here()
+	r.u16(0) // e_shnum
+	r.seekShstrndx = r.Here()
 	r.u16(uint16(elf.SHN_UNDEF)) // e_shstrndx
 
 	// Sanity check, size of file header should be the same as ehsize
@@ -120,7 +128,7 @@ func (w *Writer) WriteProgramHeaders() {
 	w.w.Seek(w.seekProgHeader, io.SeekStart)
 	w.u64(uint64(phoff))
 	w.w.Seek(w.seekProgNum, io.SeekStart)
-	w.u64(uint64(len(w.Progs)))
+	w.u16(uint16(len(w.Progs)))
 	w.w.Seek(0, io.SeekEnd)
 
 	for _, prog := range w.Progs {
@@ -135,6 +143,60 @@ func (w *Writer) WriteProgramHeaders() {
 	}
 }
 
+// WriteSectionHeaders writes the section headers at the current location
+// and patches the file header accordingly.
+func (w *Writer) WriteSectionHeaders() {
+	shstrndx := len(w.Sections)
+	strtab := &elf.SectionHeader{
+		Name:      ".strtab",
+		Type:      elf.SHT_STRTAB,
+		Offset:    uint64(w.Here()),
+		Addralign: 1,
+	}
+	w.Sections = append(w.Sections, strtab)
+
+	strtabStart := w.Here()
+
+	strToIndex := map[string]int64{}
+
+	w.Write([]byte{0}) // first entry must be the NUL character
+
+	for _, sect := range w.Sections {
+		if _, ok := strToIndex[sect.Name]; ok {
+			continue
+		}
+		strToIndex[sect.Name] = w.Here() - strtabStart
+		w.Write([]byte(sect.Name))
+		w.Write([]byte{0})
+	}
+
+	strtab.Size = uint64(w.Here() - strtabStart)
+
+	shoff := w.Here()
+
+	// Patch File Header
+	w.w.Seek(w.seekSectionHeader, io.SeekStart)
+	w.u64(uint64(shoff))
+	w.w.Seek(w.seekSectionNum, io.SeekStart)
+	w.u16(uint16(len(w.Sections)))
+	w.w.Seek(w.seekShstrndx, io.SeekStart)
+	w.u16(uint16(shstrndx))
+	w.w.Seek(0, io.SeekEnd)
+
+	for _, sect := range w.Sections {
+		w.u32(uint32(strToIndex[sect.Name]))
+		w.u32(uint32(sect.Type))
+		w.u64(uint64(sect.Flags))
+		w.u64(sect.Addr)
+		w.u64(sect.Offset)
+		w.u64(sect.Size)
+		w.u32(sect.Link)
+		w.u32(sect.Info)
+		w.u64(sect.Addralign)
+		w.u64(sect.Entsize)
+	}
+}
+
 // Here returns the current seek offset from the start of the file.
 func (w *Writer) Here() int64 {
 	r, err := w.w.Seek(0, io.SeekCurrent)