// Copyright 2018 Google Inc. All rights reserved.
//
// 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.

package symbol_inject

import (
 "debug/elf"
 "fmt"
 "io"
)

type mockableElfFile interface {
 Symbols() ([]elf.Symbol, error)
 Sections() []elf.SectionHeader
 Type() elf.Type
}

var _ mockableElfFile = elfFileWrapper{}

type elfFileWrapper struct {
 *elf.File
}

func (f elfFileWrapper) Sections() []elf.SectionHeader {
 ret := make([]elf.SectionHeader, len(f.File.Sections))
 for i, section := range f.File.Sections {
  ret[i] = section.SectionHeader
 }

 return ret
}

func (f elfFileWrapper) Type() elf.Type {
 return f.File.Type
}

type mockElfFile struct {
 symbols  []elf.Symbol
 sections []elf.SectionHeader
 t        elf.Type
}

func (f mockElfFile) Sections() []elf.SectionHeader  { return f.sections }
func (f mockElfFile) Symbols() ([]elf.Symbol, error) { return f.symbols, nil }
func (f mockElfFile) Type() elf.Type                 { return f.t }

func elfSymbolsFromFile(r io.ReaderAt) (*File, error) {
 elfFile, err := elf.NewFile(r)
 if err != nil {
  return nil, cantParseError{err}
 }
 return extractElfSymbols(elfFileWrapper{elfFile})
}

func extractElfSymbols(elfFile mockableElfFile) (*File, error) {
 symbols, err := elfFile.Symbols()
 if err != nil {
  return nil, err
 }

 file := &File{}

 for _, section := range elfFile.Sections() {
  file.Sections = append(file.Sections, &Section{
   Name:   section.Name,
   Addr:   section.Addr,
   Offset: section.Offset,
   Size:   section.Size,
  })
 }

 _ = elf.Section{}

 for _, symbol := range symbols {
  if elf.ST_TYPE(symbol.Info) != elf.STT_OBJECT {
   continue
  }
  if symbol.Section == elf.SHN_UNDEF || symbol.Section >= elf.SHN_LORESERVE {
   continue
  }
  if int(symbol.Section) >= len(file.Sections) {
   return nil, fmt.Errorf("invalid section index %d", symbol.Section)
  }

  section := file.Sections[symbol.Section]

  var addr uint64
  switch elfFile.Type() {
  case elf.ET_REL:
   // "In relocatable files, st_value holds a section offset for a defined symbol.
   // That is, st_value is an offset from the beginning of the section that st_shndx identifies."
   addr = symbol.Value
  case elf.ET_EXEC, elf.ET_DYN:
   // "In executable and shared object files, st_value holds a virtual address. To make these
   // files’ symbols more useful for the dynamic linker, the section offset (file interpretation)
   // gives way to a virtual address (memory interpretation) for which the section number is
   // irrelevant."
   if symbol.Value < section.Addr {
    return nil, fmt.Errorf("symbol starts before the start of its section")
   }
   addr = symbol.Value - section.Addr
   if addr+symbol.Size > section.Size {
    return nil, fmt.Errorf("symbol extends past the end of its section")
   }
  default:
   return nil, fmt.Errorf("unsupported elf file type %d", elfFile.Type())
  }

  file.Symbols = append(file.Symbols, &Symbol{
   Name:    symbol.Name,
   Addr:    addr,
   Size:    symbol.Size,
   Section: section,
  })
 }

 return file, nil
}

func dumpElfSymbols(r io.ReaderAt) error {
 elfFile, err := elf.NewFile(r)
 if err != nil {
  return cantParseError{err}
 }

 symbols, err := elfFile.Symbols()
 if err != nil {
  return err
 }

 fmt.Println("mockElfFile{")
 fmt.Printf("\tt: %#v,\n", elfFile.Type)

 fmt.Println("\tsections: []elf.SectionHeader{")
 for _, section := range elfFile.Sections {
  fmt.Printf("\t\t%#v,\n", section.SectionHeader)
 }
 fmt.Println("\t},")

 fmt.Println("\tsymbols: []elf.Symbol{")
 for _, symbol := range symbols {
  fmt.Printf("\t\t%#v,\n", symbol)
 }
 fmt.Println("\t},")

 fmt.Println("}")

 return nil
}