// Copyright 2022 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 elf

import (
 "debug/elf"
 "encoding/binary"
 "encoding/hex"
 "errors"
 "fmt"
 "io"
 "os"
)

const gnuBuildID = "GNU\x00"

// Identifier extracts the elf build ID from an elf file.  If allowMissing is true it returns
// an empty identifier if the file exists but the build ID note does not.
func Identifier(filename string, allowMissing bool) (string, error) {
 f, err := os.Open(filename)
 if err != nil {
  return "", fmt.Errorf("failed to open %s: %w", filename, err)
 }
 defer f.Close()

 return elfIdentifierFromReaderAt(f, filename, allowMissing)
}

// elfIdentifierFromReaderAt extracts the elf build ID from a ReaderAt.  If allowMissing is true it
// returns an empty identifier if the file exists but the build ID note does not.
func elfIdentifierFromReaderAt(r io.ReaderAt, filename string, allowMissing bool) (string, error) {
 f, err := elf.NewFile(r)
 if err != nil {
  if allowMissing {
   if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) {
    return "", nil
   }
   if _, ok := err.(*elf.FormatError); ok {
    // The file was not an elf file.
    return "", nil
   }
  }
  return "", fmt.Errorf("failed to parse elf file %s: %w", filename, err)
 }
 defer f.Close()

 buildIDNote := f.Section(".note.gnu.build-id")
 if buildIDNote == nil {
  if allowMissing {
   return "", nil
  }
  return "", fmt.Errorf("failed to find .note.gnu.build-id in  %s", filename)
 }

 buildIDs, err := readNote(buildIDNote.Open(), f.ByteOrder)
 if err != nil {
  return "", fmt.Errorf("failed to read .note.gnu.build-id: %w", err)
 }

 for name, desc := range buildIDs {
  if name == gnuBuildID {
   return hex.EncodeToString(desc), nil
  }
 }

 return "", nil
}

// readNote reads the contents of a note section, returning it as a map from name to descriptor.
func readNote(note io.Reader, byteOrder binary.ByteOrder) (map[string][]byte, error) {
 var noteHeader struct {
  Namesz uint32
  Descsz uint32
  Type   uint32
 }

 notes := make(map[string][]byte)
 for {
  err := binary.Read(note, byteOrder, ¬eHeader)
  if err != nil {
   if err == io.EOF {
    return notes, nil
   }
   return nil, fmt.Errorf("failed to read note header: %w", err)
  }

  nameBuf := make([]byte, align4(noteHeader.Namesz))
  err = binary.Read(note, byteOrder, &nameBuf)
  if err != nil {
   return nil, fmt.Errorf("failed to read note name: %w", err)
  }
  name := string(nameBuf[:noteHeader.Namesz])

  descBuf := make([]byte, align4(noteHeader.Descsz))
  err = binary.Read(note, byteOrder, &descBuf)
  if err != nil {
   return nil, fmt.Errorf("failed to read note desc: %w", err)
  }
  notes[name] = descBuf[:noteHeader.Descsz]
 }
}

// align4 rounds the input up to the next multiple of 4.
func align4(i uint32) uint32 {
 return (i + 3) &^ 3
}