[PECOFF][Driver] Add -base command line option.

llvm-svn: 186739

lld/include/lld/ReaderWriter/PECOFFTargetInfo.h

@@ -25,7 +25,7 @@ namespace lld {
 class PECOFFTargetInfo : public TargetInfo {
 public:
   PECOFFTargetInfo()
-      : _stackReserve(1024 * 1024), _stackCommit(4096),
+      : _baseAddress(0x400000), _stackReserve(1024 * 1024), _stackCommit(4096),
         _heapReserve(1024 * 1024), _heapCommit(4096),
         _subsystem(llvm::COFF::IMAGE_SUBSYSTEM_UNKNOWN), _minOSVersion(6, 0),
         _nxCompat(true), _largeAddressAware(false) {}
@@ -56,6 +56,9 @@ public:
   bool appendInputFileOrLibrary(std::string path);
   bool appendLibraryFile(StringRef path);
 
+  void setBaseAddress(uint64_t addr) { _baseAddress = addr; }
+  uint64_t getBaseAddress() const { return _baseAddress; }
+
   void setStackReserve(uint64_t size) { _stackReserve = size; }
   void setStackCommit(uint64_t size) { _stackCommit = size; }
   uint64_t getStackReserve() const { return _stackReserve; }
@@ -89,6 +92,10 @@ public:
   }
 
 private:
+  // The start address for the program. The default value for the executable is
+  // 0x400000, but can be altered using -base command line option.
+  uint64_t _baseAddress;
+
   uint64_t _stackReserve;
   uint64_t _stackCommit;
   uint64_t _heapReserve;

lld/lib/Driver/WinLinkDriver.cpp

@@ -80,23 +80,43 @@ bool checkNumber(StringRef version, const char *errorMessage,
   return true;
 }
 
-// Parse an argument for -stack or -heap. The expected string is
-// "reserveSize[,stackCommitSize]".
+// Parse an argument for -base, -stack or -heap. The expected string
+// is "<integer>[,<integer>]".
 bool parseMemoryOption(const StringRef &arg, raw_ostream &diagnostics,
                        uint64_t &reserve, uint64_t &commit) {
   StringRef reserveStr, commitStr;
   llvm::tie(reserveStr, commitStr) = arg.split(',');
-  if (!checkNumber(reserveStr, "invalid stack size: ", diagnostics))
+  if (!checkNumber(reserveStr, "invalid size: ", diagnostics))
     return false;
   reserve = atoi(reserveStr.str().c_str());
   if (!commitStr.empty()) {
-    if (!checkNumber(commitStr, "invalid stack size: ", diagnostics))
+    if (!checkNumber(commitStr, "invalid size: ", diagnostics))
       return false;
     commit = atoi(commitStr.str().c_str());
   }
   return true;
 }
 
+// Parse -base command line option. The argument for the parameter is in the
+// form of "<address>[:<size>]".
+bool parseBaseOption(PECOFFTargetInfo &info, const StringRef &arg,
+                      raw_ostream &diagnostics) {
+  // Size should be set to SizeOfImage field in the COFF header, and if it's
+  // smaller than the actual size, the linker should warn about that. Currently
+  // we just ignore the value of size parameter.
+  uint64_t addr, size;
+  if (!parseMemoryOption(arg, diagnostics, addr, size))
+    return false;
+  // It's an error if the base address is not multiple of 64K.
+  if (addr & 0xffff) {
+    diagnostics << "Base address have to be multiple of 64K, but got "
+                << addr << "\n";
+    return false;
+  }
+  info.setBaseAddress(addr);
+  return true;
+}
+
 // Parse -stack command line option
 bool parseStackOption(PECOFFTargetInfo &info, const StringRef &arg,
                       raw_ostream &diagnostics) {
@@ -280,6 +300,11 @@ bool WinLinkDriver::parse(int argc, const char *argv[],
     info.appendLLVMOption((*it)->getValue());
   }
 
+  // Handle -base
+  if (llvm::opt::Arg *arg = parsedArgs->getLastArg(OPT_base))
+    if (!parseBaseOption(info, arg->getValue(), diagnostics))
+      return true;
+
   // Handle -stack
   if (llvm::opt::Arg *arg = parsedArgs->getLastArg(OPT_stack))
     if (!parseStackOption(info, arg->getValue(), diagnostics))

lld/lib/Driver/WinLinkOptions.td

@@ -13,6 +13,10 @@ def defaultlib : Separate<["-", "/"], "defaultlib">,
     HelpText<"Add the library to the list of input files">;
 def defaultlib_c : Joined<["-", "/"], "defaultlib:">, Alias<defaultlib>;
 
+def base : Separate<["-", "/"], "base">,
+    HelpText<"Base address of the program">;
+def base_c: Joined<["-", "/"], "base:">, Alias<base>;
+
 def stack : Separate<["-", "/"], "stack">,
     HelpText<"Size of the stack">;
 def stack_c: Joined<["-", "/"], "stack:">, Alias<stack>;

lld/lib/ReaderWriter/PECOFF/WriterPECOFF.cpp

@@ -166,7 +166,7 @@ public:
 
     // The address of the executable when loaded into memory. The default for
     // DLLs is 0x10000000. The default for executables is 0x400000.
-    _peHeader.ImageBase = IMAGE_BASE;
+    _peHeader.ImageBase = targetInfo.getBaseAddress();
 
     // Sections should be page-aligned when loaded into memory, which is 4KB on
     // x86.
@@ -279,7 +279,8 @@ public:
   }
 
   void applyRelocations(uint8_t *fileBuffer,
-                        std::map<const Atom *, uint64_t> &atomRva) {
+                        std::map<const Atom *, uint64_t> &atomRva,
+                        uint64_t imageBaseAddress) {
     for (const auto *layout : _atomLayouts) {
       const DefinedAtom *atom = cast<DefinedAtom>(layout->_atom);
       for (const Reference *ref : *atom) {
@@ -297,7 +298,7 @@ public:
           break;
         case llvm::COFF::IMAGE_REL_I386_DIR32:
           // Set target's 32-bit VA.
-          *relocSite = targetAddr + IMAGE_BASE;
+          *relocSite = targetAddr + imageBaseAddress;
           break;
         case llvm::COFF::IMAGE_REL_I386_DIR32NB:
           // Set target's 32-bit RVA.
@@ -765,7 +766,8 @@ private:
   void applyAllRelocations(uint8_t *bufferStart) {
     for (auto &cp : _chunks)
       if (AtomChunk *chunk = dyn_cast<AtomChunk>(&*cp))
-        chunk->applyRelocations(bufferStart, atomRva);
+        chunk->applyRelocations(bufferStart, atomRva,
+                                _PECOFFTargetInfo.getBaseAddress());
   }
 
   void addChunk(Chunk *chunk) {

lld/test/pecoff/baseaddr.test

@@ -0,0 +1,11 @@
+# RUN: yaml2obj %p/Inputs/nop.obj.yaml > %t.obj
+#
+# RUN: lld -flavor link -out %t1 -- %t.obj \
+# RUN:   && llvm-readobj -file-headers %t1 | FileCheck -check-prefix=DEFAULT %s
+#
+# RUN: lld -flavor link -out %t1 -base 8388608 -- %t.obj \
+# RUN:   && llvm-readobj -file-headers %t1 | FileCheck -check-prefix=BASE %s
+
+DEFAULT: ImageBase: 0x400000
+
+BASE: ImageBase: 0x800000

lld/unittests/DriverTests/WinLinkDriverTest.cpp

@@ -44,6 +44,7 @@ TEST_F(WinLinkParserTest, Basic) {
   EXPECT_TRUE(_info.getInputSearchPaths().empty());
   EXPECT_EQ(6, _info.getMinOSVersion().majorVersion);
   EXPECT_EQ(0, _info.getMinOSVersion().minorVersion);
+  EXPECT_EQ(0x400000, _info.getBaseAddress());
   EXPECT_EQ(1024 * 1024ULL, _info.getStackReserve());
   EXPECT_EQ(4096ULL, _info.getStackCommit());
   EXPECT_FALSE(_info.allowRemainingUndefines());
@@ -98,6 +99,11 @@ TEST_F(WinLinkParserTest, MinMajorMinorOSVersion) {
   EXPECT_EQ(1, _info.getMinOSVersion().minorVersion);
 }
 
+TEST_F(WinLinkParserTest, Base) {
+  EXPECT_FALSE(parse("link.exe", "-base", "8388608", nullptr));
+  EXPECT_EQ(0x800000, _info.getBaseAddress());
+}
+
 TEST_F(WinLinkParserTest, StackReserve) {
   EXPECT_FALSE(parse("link.exe", "-stack", "8192", nullptr));
   EXPECT_EQ(8192ULL, _info.getStackReserve());