// 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 (
 "bytes"
 "encoding/binary"
 "fmt"
 "io"
 "math"
)

var maxUint64 uint64 = math.MaxUint64

type cantParseError struct {
 error
}

func OpenFile(r io.ReaderAt) (*File, error) {
 file, err := elfSymbolsFromFile(r)
 if elfError, ok := err.(cantParseError); ok {
  // Try as a mach-o file
  file, err = machoSymbolsFromFile(r)
  if _, ok := err.(cantParseError); ok {
   // Try as a windows PE file
   file, err = peSymbolsFromFile(r)
   if _, ok := err.(cantParseError); ok {
    // Can't parse as elf, macho, or PE, return the elf error
    return nil, elfError
   }
  }
 }
 if err != nil {
  return nil, err
 }

 file.r = r

 return file, err
}

func InjectStringSymbol(file *File, w io.Writer, symbol, value, from string) error {
 offset, size, err := findSymbol(file, symbol)
 if err != nil {
  return err
 }

 if uint64(len(value))+1 > size {
  return fmt.Errorf("value length %d overflows symbol size %d", len(value), size)
 }

 if from != "" {
  // Read the exsting symbol contents and verify they match the expected value
  expected := make([]byte, size)
  existing := make([]byte, size)
  copy(expected, from)
  _, err := file.r.ReadAt(existing, int64(offset))
  if err != nil {
   return err
  }
  if bytes.Compare(existing, expected) != 0 {
   return fmt.Errorf("existing symbol contents %q did not match expected value %q",
    string(existing), string(expected))
  }
 }

 buf := make([]byte, size)
 copy(buf, value)

 return copyAndInject(file.r, w, offset, buf)
}

func InjectUint64Symbol(file *File, w io.Writer, symbol string, value uint64) error {
 offset, size, err := findSymbol(file, symbol)
 if err != nil {
  return err
 }

 if size != 8 {
  return fmt.Errorf("symbol %q is not a uint64, it is %d bytes long", symbol, size)
 }

 buf := make([]byte, 8)
 binary.LittleEndian.PutUint64(buf, value)

 return copyAndInject(file.r, w, offset, buf)
}

func copyAndInject(r io.ReaderAt, w io.Writer, offset uint64, buf []byte) (err error) {
 // Copy the first bytes up to the symbol offset
 _, err = io.Copy(w, io.NewSectionReader(r, 0, int64(offset)))

 // Write the injected value in the output file
 if err == nil {
  _, err = w.Write(buf)
 }

 // Write the remainder of the file
 pos := int64(offset) + int64(len(buf))
 if err == nil {
  _, err = io.Copy(w, io.NewSectionReader(r, pos, 1<<63-1-pos))
 }

 if err == io.EOF {
  err = io.ErrUnexpectedEOF
 }

 return err
}

func findSymbol(file *File, symbolName string) (uint64, uint64, error) {
 for i, symbol := range file.Symbols {
  if symbol.Name == symbolName {
   // Find the next symbol (n the same section with a higher address
   var n int
   for n = i; n < len(file.Symbols); n++ {
    if file.Symbols[n].Section != symbol.Section {
     n = len(file.Symbols)
     break
    }
    if file.Symbols[n].Addr > symbol.Addr {
     break
    }
   }

   size := symbol.Size
   if size == 0 {
    var end uint64
    if n < len(file.Symbols) {
     end = file.Symbols[n].Addr
    } else {
     end = symbol.Section.Size
    }

    if end <= symbol.Addr || end > symbol.Addr+4096 {
     return maxUint64, maxUint64, fmt.Errorf("symbol end address does not seem valid, %x:%x", symbol.Addr, end)
    }

    size = end - symbol.Addr
   }

   offset := symbol.Section.Offset + symbol.Addr

   return uint64(offset), uint64(size), nil
  }
 }

 return maxUint64, maxUint64, fmt.Errorf("symbol not found")
}

type File struct {
 r           io.ReaderAt
 Symbols     []*Symbol
 Sections    []*Section
 IsMachoFile bool
}

type Symbol struct {
 Name    string
 Addr    uint64 // Address of the symbol inside the section.
 Size    uint64 // Size of the symbol, if known.
 Section *Section
}

type Section struct {
 Name   string
 Addr   uint64 // Virtual address of the start of the section.
 Offset uint64 // Offset into the file of the start of the section.
 Size   uint64
}

func DumpSymbols(r io.ReaderAt) error {
 err := dumpElfSymbols(r)
 if elfError, ok := err.(cantParseError); ok {
  // Try as a mach-o file
  err = dumpMachoSymbols(r)
  if _, ok := err.(cantParseError); ok {
   // Try as a windows PE file
   err = dumpPESymbols(r)
   if _, ok := err.(cantParseError); ok {
    // Can't parse as elf, macho, or PE, return the elf error
    return elfError
   }
  }
 }
 return err
}