/*
 * 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;

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 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);
    ByteBuffer stringTable = ByteBuffer.wrap(stringTableBytes);

    map.position(symbolTableOffset);

    Map<Integer, Integer> strings = new HashMap<>();
    for (int i = 0; i < symbolTableCount; i++) {
      int stringTableIndexPosition = map.position();
      int stringTableIndex = ObjectFileScrubbers.getLittleEndianInt(map);
      byte type = map.get();
      /* section */ map.get();
      /* description */ ObjectFileScrubbers.getLittleEndianShort(map);
      int valuePosition = map.position();
      if (header.getIs64Bit()) {
        /* value */ ObjectFileScrubbers.getLittleEndianLong(map);
      } else {
        /* value */ ObjectFileScrubbers.getLittleEndianInt(map);
      }
      if (stringTableIndex < 2) {
        continue;
      }

      int position = map.position();
      try {
        int newStringTableIndex;
        if (strings.containsKey(stringTableIndex)) {
          newStringTableIndex = strings.get(stringTableIndex);
        } else {
          stringTable.position(stringTableIndex);
          String string = ObjectFileScrubbers.getAsciiString(stringTable);
          if (type == N_OSO) {
            for (Map.Entry<Path, Path> root : cellRoots.entrySet()) {
              String rootPrefix = root.getKey() + "/";
              if (string.startsWith(rootPrefix)) {
                String replacementPrefix = root.getValue().toString();
                if (replacementPrefix.equals("")) {
                  replacementPrefix = ".";
                }
                string = replacementPrefix + "/" + string.substring(rootPrefix.length());
              }
            }

            map.position(valuePosition);
            int lastModifiedValue = ObjectFileCommonModificationDate.COMMON_MODIFICATION_TIME_STAMP;
            if (header.getIs64Bit()) {
              ObjectFileScrubbers.putLittleEndianLong(map, lastModifiedValue);
            } else {
              ObjectFileScrubbers.putLittleEndianInt(map, lastModifiedValue);
            }
          }
          map.position(currentStringTableOffset);
          ObjectFileScrubbers.putAsciiString(map, string);
          newStringTableIndex = currentStringTableOffset - stringTableOffset;
          currentStringTableOffset = map.position();
          strings.put(stringTableIndex, newStringTableIndex);
        }
        map.position(stringTableIndexPosition);
        ObjectFileScrubbers.putLittleEndianInt(map, newStringTableIndex);
      } finally {
        map.position(position);
      }
    }

    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 MachoHeader.of(commandsCount, is64bit);
  }

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