/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * 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 com.facebook.buck.cxx.toolchain.objectfile;

import com.facebook.buck.util.ObjectFileCommonModificationDate;
import com.google.common.base.Charsets;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import com.google.common.primitives.Ints;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.MappedByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

public class Machos {

  // http://www.opensource.apple.com/source/xnu/xnu-1699.32.7/EXTERNAL_HEADERS/mach-o/loader.h
  // File magic
  static final byte[] MH_MAGIC = Ints.toByteArray(0xFEEDFACE);
  static final byte[] MH_MAGIC_64 = Ints.toByteArray(0xFEEDFACF);
  static final byte[] MH_CIGAM = Ints.toByteArray(0xCEFAEDFE);
  static final byte[] MH_CIGAM_64 = Ints.toByteArray(0xCFFAEDFE);
  // Map segment load command
  static final int LC_SEGMENT = 0x1;
  // Symbol table load command
  public static final int LC_SYMTAB = 0x2;
  // UUID load command
  static final int LC_UUID = 0x1B;
  // Map 64 bit segment load command
  static final int LC_SEGMENT_64 = 0x19;

  static final int LC_REQ_DYLD = 0x80000000;
  static final int LC_DYLD_INFO = 0x22;
  public static final int LC_DYLD_INFO_ONLY = (0x22 | LC_REQ_DYLD);

  // http://www.opensource.apple.com/source/xnu/xnu-1699.32.7/EXTERNAL_HEADERS/mach-o/stab.h
  // Description of object file STAB entries
  static final short N_OSO = (short) 0x66;

  static final String LINKEDIT = "__LINKEDIT";

  private static final int NO_VALUE_MARKER = -1;

  private Machos() {}

  static void setUuidIfPresent(MappedByteBuffer map, byte[] uuid) throws MachoException {
    int commandsCount = getHeader(map).getCommandsCount();

    for (int i = 0; i < commandsCount; i++) {
      int command = ObjectFileScrubbers.getLittleEndianInt(map);
      int commandSize = ObjectFileScrubbers.getLittleEndianInt(map);
      if (LC_UUID == command) {
        ObjectFileScrubbers.putBytes(map, uuid);
        return;
      } else {
        /* Command body */ ObjectFileScrubbers.getBytes(map, commandSize - 8);
      }
    }
  }

  static boolean isMacho(FileChannel file) throws IOException {
    MappedByteBuffer map = file.map(FileChannel.MapMode.READ_ONLY, 0, MH_MAGIC.length);

    byte[] magic = ObjectFileScrubbers.getBytes(map, MH_MAGIC.length);
    return Arrays.equals(MH_MAGIC, magic)
        || Arrays.equals(MH_CIGAM, magic)
        || Arrays.equals(MH_MAGIC_64, magic)
        || Arrays.equals(MH_CIGAM_64, magic);
  }

  /**
   * Relativize paths in OSO entries.
   *
   * <p>OSO entries point to other files containing debug information. These are generated by the
   * linker as absolute paths.
   */
  static void relativizeOsoSymbols(FileChannel file, ImmutableMap<Path, Path> cellRoots)
      throws IOException, MachoException {
    cellRoots.forEach(
        (from, to) -> {
          Preconditions.checkArgument(from.isAbsolute());
          Preconditions.checkArgument(!to.isAbsolute());
        });

    long size = file.size();
    MappedByteBuffer map = file.map(FileChannel.MapMode.READ_WRITE, 0, size);

    MachoHeader header = getHeader(map);

    int symbolTableOffset = 0;
    int symbolTableCount = 0;
    int stringTableOffset = 0;
    int stringTableSizePosition = 0;
    int stringTableSize = 0;
    boolean symbolTableSegmentFound = false;
    int segmentSizePosition = 0;
    long segmentSize = 0;
    boolean linkEditSegmentFound = false;
    int segmentFileSizePosition = 0;
    int segment64FileSizePosition = 0;

    int commandsCount = header.getCommandsCount();
    for (int i = 0; i < commandsCount; i++) {
      int commandStart = map.position(); // NOPMD
      int command = ObjectFileScrubbers.getLittleEndianInt(map);
      int commandSize = ObjectFileScrubbers.getLittleEndianInt(map); // NOPMD
      switch (command) {
        case LC_SYMTAB:
          symbolTableOffset = ObjectFileScrubbers.getLittleEndianInt(map);
          symbolTableCount = ObjectFileScrubbers.getLittleEndianInt(map);
          stringTableOffset = ObjectFileScrubbers.getLittleEndianInt(map);
          stringTableSizePosition = map.position();
          stringTableSize = ObjectFileScrubbers.getLittleEndianInt(map);
          symbolTableSegmentFound = true;
          break;
        case LC_SEGMENT:
          byte[] segmentNameBytes = ObjectFileScrubbers.getBytes(map, 16);
          String segmentName = new String(segmentNameBytes, Charsets.US_ASCII);
          if (segmentName.startsWith(LINKEDIT)) {
            linkEditSegmentFound = true;
            /* vm address */ ObjectFileScrubbers.getLittleEndianInt(map);
            /* vm size */ ObjectFileScrubbers.getLittleEndianInt(map);
            /* segment file offset */ ObjectFileScrubbers.getLittleEndianInt(map);
            segmentFileSizePosition = map.position();
            segmentSize = ObjectFileScrubbers.getLittleEndianInt(map);
            /* maximum vm protection */ ObjectFileScrubbers.getLittleEndianInt(map);
            /* initial vm protection */ ObjectFileScrubbers.getLittleEndianInt(map);
            /* number of sections */ ObjectFileScrubbers.getLittleEndianInt(map);
            /* flags */ ObjectFileScrubbers.getLittleEndianInt(map);

            if (segmentSizePosition != 0) {
              throw new MachoException("multiple map segment commands map string table");
            }
            segmentSizePosition = segmentFileSizePosition;
          }
          break;
        case LC_SEGMENT_64:
          byte[] segment64NameBytes = ObjectFileScrubbers.getBytes(map, 16);
          String segment64Name = new String(segment64NameBytes, Charsets.US_ASCII);
          if (segment64Name.startsWith(LINKEDIT)) {
            linkEditSegmentFound = true;
            /* vm address */ ObjectFileScrubbers.getLittleEndianLong(map);
            /* vm size */ ObjectFileScrubbers.getLittleEndianLong(map);
            /* segment file offset */ ObjectFileScrubbers.getLittleEndianLong(map);
            segment64FileSizePosition = map.position();
            segmentSize = ObjectFileScrubbers.getLittleEndianLong(map);
            /* maximum vm protection */ ObjectFileScrubbers.getLittleEndianInt(map);
            /* initial vm protection */ ObjectFileScrubbers.getLittleEndianInt(map);
            /* number of sections */ ObjectFileScrubbers.getLittleEndianInt(map);
            /* flags */ ObjectFileScrubbers.getLittleEndianInt(map);

            if (segmentSizePosition != 0) {
              throw new MachoException("multiple map segment commands map string table");
            }
            segmentSizePosition = segment64FileSizePosition;
          }
          break;
      }
      map.position(commandStart + commandSize);
    }

    if (!linkEditSegmentFound) {
      /*The OSO entries are identified in segments named __LINKEDIT. If no segment is found with
      that name, there is nothing to scrub.*/
      return;
    }
    if (stringTableSize == 0) {
      return;
    }

    if (!isValidFilesize(header, segmentSize)) {
      throw new MachoException("32bit map segment file size too big");
    }

    if (!symbolTableSegmentFound) {
      throw new MachoException("LC_SYMTAB command not found");
    }
    if (stringTableOffset + stringTableSize != size) {
      throw new MachoException("String table does not end at end of file");
    }
    if (segmentSizePosition == 0 || segmentSize == 0) {
      throw new MachoException("LC_SEGMENT or LC_SEGMENT_64 command for string table not found");
    }

    map.position(stringTableOffset);
    if (map.get() != 0x20) {
      throw new MachoException("First character in the string table is not a space");
    }
    if (map.get() != 0x00) {
      throw new MachoException("Second character in the string table is not a NUL");
    }
    int currentStringTableOffset = map.position();

    byte[] stringTableBytes = new byte[stringTableSize];
    map.position(stringTableOffset);
    map.get(stringTableBytes);

    map.position(symbolTableOffset);

    // NB: We need to rewrite the string table as it's not deterministic and it would break
    //     caching behavior. On the other hand, the symbol table order is deterministic.

    boolean is64bit = header.getIs64Bit();
    Map<byte[], byte[]> replacementPathMap = generateReplacementMap(cellRoots);
    IntIntMap strings = new IntIntMap4a(symbolTableCount, 0.75f, NO_VALUE_MARKER);
    for (int i = 0; i < symbolTableCount; i++) {
      // Each LC_SYMTAB entry consists of the following fields:
      // - String Index: 4 bytes (offset into the string table)
      // - Type: 1 byte
      // - Section: 1 byte
      // - Description: 2 bytes
      // - Value: 8 bytes on 64bit, 4 bytes on 32bit
      int stringTableIndexPosition = map.position();
      int stringTableIndex = ObjectFileScrubbers.getLittleEndianInt(map);
      byte type = map.get();

      if (stringTableIndex >= 2) {
        int newStringTableIndex = strings.get(stringTableIndex);
        if (newStringTableIndex == NO_VALUE_MARKER) {
          ByteBuffer charByteBuffer =
              ObjectFileScrubbers.getCharByteBuffer(stringTableBytes, stringTableIndex);

          if (type == N_OSO) {
            Optional<ByteBuffer> maybeRewrittenCharByteBuffer =
                tryRewritingMatchingPath(stringTableBytes, stringTableIndex, replacementPathMap);
            if (maybeRewrittenCharByteBuffer.isPresent()) {
              charByteBuffer = maybeRewrittenCharByteBuffer.get();
            }

            int valuePosition = stringTableIndexPosition + 8;
            map.position(valuePosition);
            int lastModifiedValue = ObjectFileCommonModificationDate.COMMON_MODIFICATION_TIME_STAMP;
            if (is64bit) {
              ObjectFileScrubbers.putLittleEndianLong(map, lastModifiedValue);
            } else {
              ObjectFileScrubbers.putLittleEndianInt(map, lastModifiedValue);
            }
          }
          ObjectFileScrubbers.putCharByteBuffer(map, currentStringTableOffset, charByteBuffer);

          newStringTableIndex = currentStringTableOffset - stringTableOffset;
          strings.put(stringTableIndex, newStringTableIndex);

          currentStringTableOffset = map.position();
        }
        map.position(stringTableIndexPosition);
        ObjectFileScrubbers.putLittleEndianInt(map, newStringTableIndex);
      }

      int symtabEntrySize = 4 + 1 + 1 + 2 + (is64bit ? 8 : 4);
      int nextSymtabEntryOffset = stringTableIndexPosition + symtabEntrySize;
      map.position(nextSymtabEntryOffset);
    }

    map.position(stringTableSizePosition);
    int newStringTableSize = currentStringTableOffset - stringTableOffset;
    ObjectFileScrubbers.putLittleEndianInt(map, newStringTableSize);

    map.position(segmentSizePosition);
    long newSize = segmentSize + (newStringTableSize - stringTableSize);
    if (isValidFilesize(header, newSize)) {
      if (header.getIs64Bit()) {
        ObjectFileScrubbers.putLittleEndianLong(map, newSize);
      } else {
        ObjectFileScrubbers.putLittleEndianInt(map, (int) newSize);
      }
    } else {
      throw new MachoException("32bit scrubbed map segment file size too big");
    }

    file.truncate(currentStringTableOffset);
  }

  private static boolean isValidFilesize(MachoHeader header, long filesize) {
    return (header.getIs64Bit() || filesize <= Integer.MAX_VALUE);
  }

  /** Returns the Mach-O header provided the file is Mach-O, otherwise throws an exception. */
  protected static MachoHeader getHeader(MappedByteBuffer map) throws MachoException {
    byte[] magic = ObjectFileScrubbers.getBytes(map, MH_MAGIC.length);
    boolean is64bit;
    if (Arrays.equals(MH_MAGIC, magic) || Arrays.equals(MH_CIGAM, magic)) {
      is64bit = false;
    } else if (Arrays.equals(MH_MAGIC_64, magic) || Arrays.equals(MH_CIGAM_64, magic)) {
      is64bit = true;
    } else {
      throw new MachoException("invalid Mach-O magic");
    }

    /* CPU type */
    ObjectFileScrubbers.getLittleEndianInt(map);
    /* CPU subtype */
    ObjectFileScrubbers.getLittleEndianInt(map);
    /* File type */
    ObjectFileScrubbers.getLittleEndianInt(map);
    int commandsCount = ObjectFileScrubbers.getLittleEndianInt(map);
    /* Commands size */
    ObjectFileScrubbers.getLittleEndianInt(map);
    /* Flags */
    ObjectFileScrubbers.getLittleEndianInt(map);
    if (is64bit) {
      /* reserved */ ObjectFileScrubbers.getLittleEndianInt(map);
    }
    return ImmutableMachoHeader.of(commandsCount, is64bit);
  }

  public static class MachoException extends Exception {
    public MachoException(String msg) {
      super(msg);
    }
  }

  /**
   * Prepares a replacement map for prefixes. For example, if {@p pathMap} had two entries: 1.
   * "/path/to/repo" -> "", 2. "/path/to/repo/cell" -> "cell"
   *
   * <p>The resulting map would contain: 1. "/path/to/repo/" -> "./" 2. "/path/to/repo/cell/" ->
   * "cell/"
   *
   * <p>The above means we can do very simple (and fast!) search & replace.
   */
  public static Map<byte[], byte[]> generateReplacementMap(Map<Path, Path> pathMap) {
    // Preprocess the input map once, so we can be efficient with byte[] arrays instead of Strings
    Map<byte[], byte[]> replacementMap = new HashMap<>();
    for (Map.Entry<Path, Path> pathEntry : pathMap.entrySet()) {
      String searchPrefix = pathEntry.getKey() + "/";

      String replacementPrefix = pathEntry.getValue().toString();
      if (replacementPrefix.isEmpty()) {
        replacementPrefix = ".";
      }
      replacementPrefix = replacementPrefix + "/";

      byte[] searchPrefixBytes = searchPrefix.getBytes(Charsets.UTF_8);
      byte[] replacementBytes = replacementPrefix.getBytes(Charsets.UTF_8);

      if (replacementBytes.length > searchPrefixBytes.length) {
        throw new IllegalStateException(
            "Relativization should shorten paths, not lengthen. Prefix="
                + searchPrefix
                + "; replacement="
                + replacementPrefix);
      }

      replacementMap.put(searchPrefixBytes, replacementBytes);
    }

    return replacementMap;
  }

  /**
   * Checks whether a string matches a prefix and returns a rewritten copy if that's the case. For
   * example, given a string "/Users/fb/repo/cell" and a replacement map containing
   * "/Users/fb/repo/" -> "./", it will return a ByteBuffer wrapping "./cell".
   *
   * <p>NB: This is a perf sensitive method.
   */
  public static Optional<ByteBuffer> tryRewritingMatchingPath(
      byte[] stringBytes, int stringOffset, Map<byte[], byte[]> replacementMap) {
    int nullCharOffset = stringOffset;
    while (stringBytes[nullCharOffset] != 0x0) {
      ++nullCharOffset;
    }

    int stringLength = nullCharOffset - stringOffset;
    for (Map.Entry<byte[], byte[]> replacementEntry : replacementMap.entrySet()) {
      byte[] searchPrefix = replacementEntry.getKey();
      if (bytesStartsWith(stringBytes, stringOffset, searchPrefix)) {
        // stringBytes variable:
        //
        //  stringOffset
        //       |
        //       v
        //       +------------------+---------------+
        //       |   searchPrefix   |    suffix     |
        //       +------------------+---------------+
        //                          ^
        //                          |
        //                     suffixOffset
        //
        // replacement variable:
        //
        //       0
        //       |
        //       v
        //       +-------------+---------------+
        //       | replacement |    suffix     |
        //       +-------------+---------------+
        //                     ^
        //                     |
        //             replacement.length

        byte[] replacement = replacementEntry.getValue();

        int suffixLength = stringLength - searchPrefix.length;
        int suffixOffset = stringOffset + searchPrefix.length;

        byte[] rewrittenPath = new byte[replacement.length + suffixLength];
        System.arraycopy(replacement, 0, rewrittenPath, 0, replacement.length);
        System.arraycopy(
            stringBytes, suffixOffset, rewrittenPath, replacement.length, suffixLength);

        return Optional.of(ByteBuffer.wrap(rewrittenPath));
      }
    }

    return Optional.empty();
  }

  private static boolean bytesStartsWith(byte[] haystack, int haystackOffset, byte[] needle) {
    if (haystackOffset + needle.length > haystack.length) {
      return false;
    }

    for (int i = 0; i < needle.length; i++) {
      if (haystack[haystackOffset + i] != needle[i]) {
        return false;
      }
    }

    return true;
  }
}
