1 //===-- WindowsResource.cpp -------------------------------------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements the .res file class.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "llvm/Object/WindowsResource.h"
14 #include "llvm/Object/COFF.h"
15 #include "llvm/Support/FileOutputBuffer.h"
16 #include "llvm/Support/FormatVariadic.h"
17 #include "llvm/Support/MathExtras.h"
18 #include "llvm/Support/ScopedPrinter.h"
19 #include <ctime>
20 #include <queue>
21 #include <system_error>
22 
23 using namespace llvm;
24 using namespace object;
25 
26 namespace llvm {
27 namespace object {
28 
29 #define RETURN_IF_ERROR(X)                                                     \
30   if (auto EC = X)                                                             \
31     return EC;
32 
33 #define UNWRAP_REF_OR_RETURN(Name, Expr)                                       \
34   auto Name##OrErr = Expr;                                                     \
35   if (!Name##OrErr)                                                            \
36     return Name##OrErr.takeError();                                            \
37   const auto &Name = *Name##OrErr;
38 
39 #define UNWRAP_OR_RETURN(Name, Expr)                                           \
40   auto Name##OrErr = Expr;                                                     \
41   if (!Name##OrErr)                                                            \
42     return Name##OrErr.takeError();                                            \
43   auto Name = *Name##OrErr;
44 
45 const uint32_t MIN_HEADER_SIZE = 7 * sizeof(uint32_t) + 2 * sizeof(uint16_t);
46 
47 // COFF files seem to be inconsistent with alignment between sections, just use
48 // 8-byte because it makes everyone happy.
49 const uint32_t SECTION_ALIGNMENT = sizeof(uint64_t);
50 
51 WindowsResource::WindowsResource(MemoryBufferRef Source)
52     : Binary(Binary::ID_WinRes, Source) {
53   size_t LeadingSize = WIN_RES_MAGIC_SIZE + WIN_RES_NULL_ENTRY_SIZE;
54   BBS = BinaryByteStream(Data.getBuffer().drop_front(LeadingSize),
55                          support::little);
56 }
57 
58 // static
59 Expected<std::unique_ptr<WindowsResource>>
60 WindowsResource::createWindowsResource(MemoryBufferRef Source) {
61   if (Source.getBufferSize() < WIN_RES_MAGIC_SIZE + WIN_RES_NULL_ENTRY_SIZE)
62     return make_error<GenericBinaryError>(
63         Source.getBufferIdentifier() + ": too small to be a resource file",
64         object_error::invalid_file_type);
65   std::unique_ptr<WindowsResource> Ret(new WindowsResource(Source));
66   return std::move(Ret);
67 }
68 
69 Expected<ResourceEntryRef> WindowsResource::getHeadEntry() {
70   if (BBS.getLength() < sizeof(WinResHeaderPrefix) + sizeof(WinResHeaderSuffix))
71     return make_error<EmptyResError>(getFileName() + " contains no entries",
72                                      object_error::unexpected_eof);
73   return ResourceEntryRef::create(BinaryStreamRef(BBS), this);
74 }
75 
76 ResourceEntryRef::ResourceEntryRef(BinaryStreamRef Ref,
77                                    const WindowsResource *Owner)
78     : Reader(Ref), Owner(Owner) {}
79 
80 Expected<ResourceEntryRef>
81 ResourceEntryRef::create(BinaryStreamRef BSR, const WindowsResource *Owner) {
82   auto Ref = ResourceEntryRef(BSR, Owner);
83   if (auto E = Ref.loadNext())
84     return std::move(E);
85   return Ref;
86 }
87 
88 Error ResourceEntryRef::moveNext(bool &End) {
89   // Reached end of all the entries.
90   if (Reader.bytesRemaining() == 0) {
91     End = true;
92     return Error::success();
93   }
94   RETURN_IF_ERROR(loadNext());
95 
96   return Error::success();
97 }
98 
99 static Error readStringOrId(BinaryStreamReader &Reader, uint16_t &ID,
100                             ArrayRef<UTF16> &Str, bool &IsString) {
101   uint16_t IDFlag;
102   RETURN_IF_ERROR(Reader.readInteger(IDFlag));
103   IsString = IDFlag != 0xffff;
104 
105   if (IsString) {
106     Reader.setOffset(
107         Reader.getOffset() -
108         sizeof(uint16_t)); // Re-read the bytes which we used to check the flag.
109     RETURN_IF_ERROR(Reader.readWideString(Str));
110   } else
111     RETURN_IF_ERROR(Reader.readInteger(ID));
112 
113   return Error::success();
114 }
115 
116 Error ResourceEntryRef::loadNext() {
117   const WinResHeaderPrefix *Prefix;
118   RETURN_IF_ERROR(Reader.readObject(Prefix));
119 
120   if (Prefix->HeaderSize < MIN_HEADER_SIZE)
121     return make_error<GenericBinaryError>(Owner->getFileName() +
122                                               ": header size too small",
123                                           object_error::parse_failed);
124 
125   RETURN_IF_ERROR(readStringOrId(Reader, TypeID, Type, IsStringType));
126 
127   RETURN_IF_ERROR(readStringOrId(Reader, NameID, Name, IsStringName));
128 
129   RETURN_IF_ERROR(Reader.padToAlignment(WIN_RES_HEADER_ALIGNMENT));
130 
131   RETURN_IF_ERROR(Reader.readObject(Suffix));
132 
133   RETURN_IF_ERROR(Reader.readArray(Data, Prefix->DataSize));
134 
135   RETURN_IF_ERROR(Reader.padToAlignment(WIN_RES_DATA_ALIGNMENT));
136 
137   return Error::success();
138 }
139 
140 WindowsResourceParser::WindowsResourceParser(bool MinGW)
141     : Root(false), MinGW(MinGW) {}
142 
143 void printResourceTypeName(uint16_t TypeID, raw_ostream &OS) {
144   switch (TypeID) {
145   case  1: OS << "CURSOR (ID 1)"; break;
146   case  2: OS << "BITMAP (ID 2)"; break;
147   case  3: OS << "ICON (ID 3)"; break;
148   case  4: OS << "MENU (ID 4)"; break;
149   case  5: OS << "DIALOG (ID 5)"; break;
150   case  6: OS << "STRINGTABLE (ID 6)"; break;
151   case  7: OS << "FONTDIR (ID 7)"; break;
152   case  8: OS << "FONT (ID 8)"; break;
153   case  9: OS << "ACCELERATOR (ID 9)"; break;
154   case 10: OS << "RCDATA (ID 10)"; break;
155   case 11: OS << "MESSAGETABLE (ID 11)"; break;
156   case 12: OS << "GROUP_CURSOR (ID 12)"; break;
157   case 14: OS << "GROUP_ICON (ID 14)"; break;
158   case 16: OS << "VERSIONINFO (ID 16)"; break;
159   case 17: OS << "DLGINCLUDE (ID 17)"; break;
160   case 19: OS << "PLUGPLAY (ID 19)"; break;
161   case 20: OS << "VXD (ID 20)"; break;
162   case 21: OS << "ANICURSOR (ID 21)"; break;
163   case 22: OS << "ANIICON (ID 22)"; break;
164   case 23: OS << "HTML (ID 23)"; break;
165   case 24: OS << "MANIFEST (ID 24)"; break;
166   default: OS << "ID " << TypeID; break;
167   }
168 }
169 
170 static bool convertUTF16LEToUTF8String(ArrayRef<UTF16> Src, std::string &Out) {
171   if (!sys::IsBigEndianHost)
172     return convertUTF16ToUTF8String(Src, Out);
173 
174   std::vector<UTF16> EndianCorrectedSrc;
175   EndianCorrectedSrc.resize(Src.size() + 1);
176   llvm::copy(Src, EndianCorrectedSrc.begin() + 1);
177   EndianCorrectedSrc[0] = UNI_UTF16_BYTE_ORDER_MARK_SWAPPED;
178   return convertUTF16ToUTF8String(makeArrayRef(EndianCorrectedSrc), Out);
179 }
180 
181 static std::string makeDuplicateResourceError(
182     const ResourceEntryRef &Entry, StringRef File1, StringRef File2) {
183   std::string Ret;
184   raw_string_ostream OS(Ret);
185 
186   OS << "duplicate resource:";
187 
188   OS << " type ";
189   if (Entry.checkTypeString()) {
190     std::string UTF8;
191     if (!convertUTF16LEToUTF8String(Entry.getTypeString(), UTF8))
192       UTF8 = "(failed conversion from UTF16)";
193     OS << '\"' << UTF8 << '\"';
194   } else
195     printResourceTypeName(Entry.getTypeID(), OS);
196 
197   OS << "/name ";
198   if (Entry.checkNameString()) {
199     std::string UTF8;
200     if (!convertUTF16LEToUTF8String(Entry.getNameString(), UTF8))
201       UTF8 = "(failed conversion from UTF16)";
202     OS << '\"' << UTF8 << '\"';
203   } else {
204     OS << "ID " << Entry.getNameID();
205   }
206 
207   OS << "/language " << Entry.getLanguage() << ", in " << File1 << " and in "
208      << File2;
209 
210   return OS.str();
211 }
212 
213 static void printStringOrID(const WindowsResourceParser::StringOrID &S,
214                             raw_string_ostream &OS, bool IsType, bool IsID) {
215   if (S.IsString) {
216     std::string UTF8;
217     if (!convertUTF16LEToUTF8String(S.String, UTF8))
218       UTF8 = "(failed conversion from UTF16)";
219     OS << '\"' << UTF8 << '\"';
220   } else if (IsType)
221     printResourceTypeName(S.ID, OS);
222   else if (IsID)
223     OS << "ID " << S.ID;
224   else
225     OS << S.ID;
226 }
227 
228 static std::string makeDuplicateResourceError(
229     const std::vector<WindowsResourceParser::StringOrID> &Context,
230     StringRef File1, StringRef File2) {
231   std::string Ret;
232   raw_string_ostream OS(Ret);
233 
234   OS << "duplicate resource:";
235 
236   if (Context.size() >= 1) {
237     OS << " type ";
238     printStringOrID(Context[0], OS, /* IsType */ true, /* IsID */ true);
239   }
240 
241   if (Context.size() >= 2) {
242     OS << "/name ";
243     printStringOrID(Context[1], OS, /* IsType */ false, /* IsID */ true);
244   }
245 
246   if (Context.size() >= 3) {
247     OS << "/language ";
248     printStringOrID(Context[2], OS, /* IsType */ false, /* IsID */ false);
249   }
250   OS << ", in " << File1 << " and in " << File2;
251 
252   return OS.str();
253 }
254 
255 // MinGW specific. Remove default manifests (with language zero) if there are
256 // other manifests present, and report an error if there are more than one
257 // manifest with a non-zero language code.
258 // GCC has the concept of a default manifest resource object, which gets
259 // linked in implicitly if present. This default manifest has got language
260 // id zero, and should be dropped silently if there's another manifest present.
261 // If the user resources surprisignly had a manifest with language id zero,
262 // we should also ignore the duplicate default manifest.
263 void WindowsResourceParser::cleanUpManifests(
264     std::vector<std::string> &Duplicates) {
265   auto TypeIt = Root.IDChildren.find(/* RT_MANIFEST */ 24);
266   if (TypeIt == Root.IDChildren.end())
267     return;
268 
269   TreeNode *TypeNode = TypeIt->second.get();
270   auto NameIt =
271       TypeNode->IDChildren.find(/* CREATEPROCESS_MANIFEST_RESOURCE_ID */ 1);
272   if (NameIt == TypeNode->IDChildren.end())
273     return;
274 
275   TreeNode *NameNode = NameIt->second.get();
276   if (NameNode->IDChildren.size() <= 1)
277     return; // None or one manifest present, all good.
278 
279   // If we have more than one manifest, drop the language zero one if present,
280   // and check again.
281   auto LangZeroIt = NameNode->IDChildren.find(0);
282   if (LangZeroIt != NameNode->IDChildren.end() &&
283       LangZeroIt->second->IsDataNode) {
284     uint32_t RemovedIndex = LangZeroIt->second->DataIndex;
285     NameNode->IDChildren.erase(LangZeroIt);
286     Data.erase(Data.begin() + RemovedIndex);
287     Root.shiftDataIndexDown(RemovedIndex);
288 
289     // If we're now down to one manifest, all is good.
290     if (NameNode->IDChildren.size() <= 1)
291       return;
292   }
293 
294   // More than one non-language-zero manifest
295   auto FirstIt = NameNode->IDChildren.begin();
296   uint32_t FirstLang = FirstIt->first;
297   TreeNode *FirstNode = FirstIt->second.get();
298   auto LastIt = NameNode->IDChildren.rbegin();
299   uint32_t LastLang = LastIt->first;
300   TreeNode *LastNode = LastIt->second.get();
301   Duplicates.push_back(
302       ("duplicate non-default manifests with languages " + Twine(FirstLang) +
303        " in " + InputFilenames[FirstNode->Origin] + " and " + Twine(LastLang) +
304        " in " + InputFilenames[LastNode->Origin])
305           .str());
306 }
307 
308 // Ignore duplicates of manifests with language zero (the default manifest),
309 // in case the user has provided a manifest with that language id. See
310 // the function comment above for context. Only returns true if MinGW is set
311 // to true.
312 bool WindowsResourceParser::shouldIgnoreDuplicate(
313     const ResourceEntryRef &Entry) const {
314   return MinGW && !Entry.checkTypeString() &&
315          Entry.getTypeID() == /* RT_MANIFEST */ 24 &&
316          !Entry.checkNameString() &&
317          Entry.getNameID() == /* CREATEPROCESS_MANIFEST_RESOURCE_ID */ 1 &&
318          Entry.getLanguage() == 0;
319 }
320 
321 bool WindowsResourceParser::shouldIgnoreDuplicate(
322     const std::vector<StringOrID> &Context) const {
323   return MinGW && Context.size() == 3 && !Context[0].IsString &&
324          Context[0].ID == /* RT_MANIFEST */ 24 && !Context[1].IsString &&
325          Context[1].ID == /* CREATEPROCESS_MANIFEST_RESOURCE_ID */ 1 &&
326          !Context[2].IsString && Context[2].ID == 0;
327 }
328 
329 Error WindowsResourceParser::parse(WindowsResource *WR,
330                                    std::vector<std::string> &Duplicates) {
331   auto EntryOrErr = WR->getHeadEntry();
332   if (!EntryOrErr) {
333     auto E = EntryOrErr.takeError();
334     if (E.isA<EmptyResError>()) {
335       // Check if the .res file contains no entries.  In this case we don't have
336       // to throw an error but can rather just return without parsing anything.
337       // This applies for files which have a valid PE header magic and the
338       // mandatory empty null resource entry.  Files which do not fit this
339       // criteria would have already been filtered out by
340       // WindowsResource::createWindowsResource().
341       consumeError(std::move(E));
342       return Error::success();
343     }
344     return E;
345   }
346 
347   ResourceEntryRef Entry = EntryOrErr.get();
348   uint32_t Origin = InputFilenames.size();
349   InputFilenames.push_back(WR->getFileName());
350   bool End = false;
351   while (!End) {
352 
353     TreeNode *Node;
354     bool IsNewNode = Root.addEntry(Entry, Origin, Data, StringTable, Node);
355     if (!IsNewNode) {
356       if (!shouldIgnoreDuplicate(Entry))
357         Duplicates.push_back(makeDuplicateResourceError(
358             Entry, InputFilenames[Node->Origin], WR->getFileName()));
359     }
360 
361     RETURN_IF_ERROR(Entry.moveNext(End));
362   }
363 
364   return Error::success();
365 }
366 
367 Error WindowsResourceParser::parse(ResourceSectionRef &RSR, StringRef Filename,
368                                    std::vector<std::string> &Duplicates) {
369   UNWRAP_REF_OR_RETURN(BaseTable, RSR.getBaseTable());
370   uint32_t Origin = InputFilenames.size();
371   InputFilenames.push_back(Filename);
372   std::vector<StringOrID> Context;
373   return addChildren(Root, RSR, BaseTable, Origin, Context, Duplicates);
374 }
375 
376 void WindowsResourceParser::printTree(raw_ostream &OS) const {
377   ScopedPrinter Writer(OS);
378   Root.print(Writer, "Resource Tree");
379 }
380 
381 bool WindowsResourceParser::TreeNode::addEntry(
382     const ResourceEntryRef &Entry, uint32_t Origin,
383     std::vector<std::vector<uint8_t>> &Data,
384     std::vector<std::vector<UTF16>> &StringTable, TreeNode *&Result) {
385   TreeNode &TypeNode = addTypeNode(Entry, StringTable);
386   TreeNode &NameNode = TypeNode.addNameNode(Entry, StringTable);
387   return NameNode.addLanguageNode(Entry, Origin, Data, Result);
388 }
389 
390 Error WindowsResourceParser::addChildren(TreeNode &Node,
391                                          ResourceSectionRef &RSR,
392                                          const coff_resource_dir_table &Table,
393                                          uint32_t Origin,
394                                          std::vector<StringOrID> &Context,
395                                          std::vector<std::string> &Duplicates) {
396 
397   for (int i = 0; i < Table.NumberOfNameEntries + Table.NumberOfIDEntries;
398        i++) {
399     UNWRAP_REF_OR_RETURN(Entry, RSR.getTableEntry(Table, i));
400     TreeNode *Child;
401 
402     if (Entry.Offset.isSubDir()) {
403 
404       // Create a new subdirectory and recurse
405       if (i < Table.NumberOfNameEntries) {
406         UNWRAP_OR_RETURN(NameString, RSR.getEntryNameString(Entry));
407         Child = &Node.addNameChild(NameString, StringTable);
408         Context.push_back(StringOrID(NameString));
409       } else {
410         Child = &Node.addIDChild(Entry.Identifier.ID);
411         Context.push_back(StringOrID(Entry.Identifier.ID));
412       }
413 
414       UNWRAP_REF_OR_RETURN(NextTable, RSR.getEntrySubDir(Entry));
415       Error E =
416           addChildren(*Child, RSR, NextTable, Origin, Context, Duplicates);
417       if (E)
418         return E;
419       Context.pop_back();
420 
421     } else {
422 
423       // Data leaves are supposed to have a numeric ID as identifier (language).
424       if (Table.NumberOfNameEntries > 0)
425         return createStringError(object_error::parse_failed,
426                                  "unexpected string key for data object");
427 
428       // Try adding a data leaf
429       UNWRAP_REF_OR_RETURN(DataEntry, RSR.getEntryData(Entry));
430       TreeNode *Child;
431       Context.push_back(StringOrID(Entry.Identifier.ID));
432       bool Added = Node.addDataChild(Entry.Identifier.ID, Table.MajorVersion,
433                                      Table.MinorVersion, Table.Characteristics,
434                                      Origin, Data.size(), Child);
435       if (Added) {
436         UNWRAP_OR_RETURN(Contents, RSR.getContents(DataEntry));
437         Data.push_back(ArrayRef<uint8_t>(
438             reinterpret_cast<const uint8_t *>(Contents.data()),
439             Contents.size()));
440       } else {
441         if (!shouldIgnoreDuplicate(Context))
442           Duplicates.push_back(makeDuplicateResourceError(
443               Context, InputFilenames[Child->Origin], InputFilenames.back()));
444       }
445       Context.pop_back();
446 
447     }
448   }
449   return Error::success();
450 }
451 
452 WindowsResourceParser::TreeNode::TreeNode(uint32_t StringIndex)
453     : StringIndex(StringIndex) {}
454 
455 WindowsResourceParser::TreeNode::TreeNode(uint16_t MajorVersion,
456                                           uint16_t MinorVersion,
457                                           uint32_t Characteristics,
458                                           uint32_t Origin, uint32_t DataIndex)
459     : IsDataNode(true), DataIndex(DataIndex), MajorVersion(MajorVersion),
460       MinorVersion(MinorVersion), Characteristics(Characteristics),
461       Origin(Origin) {}
462 
463 std::unique_ptr<WindowsResourceParser::TreeNode>
464 WindowsResourceParser::TreeNode::createStringNode(uint32_t Index) {
465   return std::unique_ptr<TreeNode>(new TreeNode(Index));
466 }
467 
468 std::unique_ptr<WindowsResourceParser::TreeNode>
469 WindowsResourceParser::TreeNode::createIDNode() {
470   return std::unique_ptr<TreeNode>(new TreeNode(0));
471 }
472 
473 std::unique_ptr<WindowsResourceParser::TreeNode>
474 WindowsResourceParser::TreeNode::createDataNode(uint16_t MajorVersion,
475                                                 uint16_t MinorVersion,
476                                                 uint32_t Characteristics,
477                                                 uint32_t Origin,
478                                                 uint32_t DataIndex) {
479   return std::unique_ptr<TreeNode>(new TreeNode(
480       MajorVersion, MinorVersion, Characteristics, Origin, DataIndex));
481 }
482 
483 WindowsResourceParser::TreeNode &WindowsResourceParser::TreeNode::addTypeNode(
484     const ResourceEntryRef &Entry,
485     std::vector<std::vector<UTF16>> &StringTable) {
486   if (Entry.checkTypeString())
487     return addNameChild(Entry.getTypeString(), StringTable);
488   else
489     return addIDChild(Entry.getTypeID());
490 }
491 
492 WindowsResourceParser::TreeNode &WindowsResourceParser::TreeNode::addNameNode(
493     const ResourceEntryRef &Entry,
494     std::vector<std::vector<UTF16>> &StringTable) {
495   if (Entry.checkNameString())
496     return addNameChild(Entry.getNameString(), StringTable);
497   else
498     return addIDChild(Entry.getNameID());
499 }
500 
501 bool WindowsResourceParser::TreeNode::addLanguageNode(
502     const ResourceEntryRef &Entry, uint32_t Origin,
503     std::vector<std::vector<uint8_t>> &Data, TreeNode *&Result) {
504   bool Added = addDataChild(Entry.getLanguage(), Entry.getMajorVersion(),
505                             Entry.getMinorVersion(), Entry.getCharacteristics(),
506                             Origin, Data.size(), Result);
507   if (Added)
508     Data.push_back(Entry.getData());
509   return Added;
510 }
511 
512 bool WindowsResourceParser::TreeNode::addDataChild(
513     uint32_t ID, uint16_t MajorVersion, uint16_t MinorVersion,
514     uint32_t Characteristics, uint32_t Origin, uint32_t DataIndex,
515     TreeNode *&Result) {
516   auto NewChild = createDataNode(MajorVersion, MinorVersion, Characteristics,
517                                  Origin, DataIndex);
518   auto ElementInserted = IDChildren.emplace(ID, std::move(NewChild));
519   Result = ElementInserted.first->second.get();
520   return ElementInserted.second;
521 }
522 
523 WindowsResourceParser::TreeNode &WindowsResourceParser::TreeNode::addIDChild(
524     uint32_t ID) {
525   auto Child = IDChildren.find(ID);
526   if (Child == IDChildren.end()) {
527     auto NewChild = createIDNode();
528     WindowsResourceParser::TreeNode &Node = *NewChild;
529     IDChildren.emplace(ID, std::move(NewChild));
530     return Node;
531   } else
532     return *(Child->second);
533 }
534 
535 WindowsResourceParser::TreeNode &WindowsResourceParser::TreeNode::addNameChild(
536     ArrayRef<UTF16> NameRef, std::vector<std::vector<UTF16>> &StringTable) {
537   std::string NameString;
538   convertUTF16LEToUTF8String(NameRef, NameString);
539 
540   auto Child = StringChildren.find(NameString);
541   if (Child == StringChildren.end()) {
542     auto NewChild = createStringNode(StringTable.size());
543     StringTable.push_back(NameRef);
544     WindowsResourceParser::TreeNode &Node = *NewChild;
545     StringChildren.emplace(NameString, std::move(NewChild));
546     return Node;
547   } else
548     return *(Child->second);
549 }
550 
551 void WindowsResourceParser::TreeNode::print(ScopedPrinter &Writer,
552                                             StringRef Name) const {
553   ListScope NodeScope(Writer, Name);
554   for (auto const &Child : StringChildren) {
555     Child.second->print(Writer, Child.first);
556   }
557   for (auto const &Child : IDChildren) {
558     Child.second->print(Writer, to_string(Child.first));
559   }
560 }
561 
562 // This function returns the size of the entire resource tree, including
563 // directory tables, directory entries, and data entries.  It does not include
564 // the directory strings or the relocations of the .rsrc section.
565 uint32_t WindowsResourceParser::TreeNode::getTreeSize() const {
566   uint32_t Size = (IDChildren.size() + StringChildren.size()) *
567                   sizeof(coff_resource_dir_entry);
568 
569   // Reached a node pointing to a data entry.
570   if (IsDataNode) {
571     Size += sizeof(coff_resource_data_entry);
572     return Size;
573   }
574 
575   // If the node does not point to data, it must have a directory table pointing
576   // to other nodes.
577   Size += sizeof(coff_resource_dir_table);
578 
579   for (auto const &Child : StringChildren) {
580     Size += Child.second->getTreeSize();
581   }
582   for (auto const &Child : IDChildren) {
583     Size += Child.second->getTreeSize();
584   }
585   return Size;
586 }
587 
588 // Shift DataIndex of all data children with an Index greater or equal to the
589 // given one, to fill a gap from removing an entry from the Data vector.
590 void WindowsResourceParser::TreeNode::shiftDataIndexDown(uint32_t Index) {
591   if (IsDataNode && DataIndex >= Index) {
592     DataIndex--;
593   } else {
594     for (auto &Child : IDChildren)
595       Child.second->shiftDataIndexDown(Index);
596     for (auto &Child : StringChildren)
597       Child.second->shiftDataIndexDown(Index);
598   }
599 }
600 
601 class WindowsResourceCOFFWriter {
602 public:
603   WindowsResourceCOFFWriter(COFF::MachineTypes MachineType,
604                             const WindowsResourceParser &Parser, Error &E);
605   std::unique_ptr<MemoryBuffer> write(uint32_t TimeDateStamp);
606 
607 private:
608   void performFileLayout();
609   void performSectionOneLayout();
610   void performSectionTwoLayout();
611   void writeCOFFHeader(uint32_t TimeDateStamp);
612   void writeFirstSectionHeader();
613   void writeSecondSectionHeader();
614   void writeFirstSection();
615   void writeSecondSection();
616   void writeSymbolTable();
617   void writeStringTable();
618   void writeDirectoryTree();
619   void writeDirectoryStringTable();
620   void writeFirstSectionRelocations();
621   std::unique_ptr<WritableMemoryBuffer> OutputBuffer;
622   char *BufferStart;
623   uint64_t CurrentOffset = 0;
624   COFF::MachineTypes MachineType;
625   const WindowsResourceParser::TreeNode &Resources;
626   const ArrayRef<std::vector<uint8_t>> Data;
627   uint64_t FileSize;
628   uint32_t SymbolTableOffset;
629   uint32_t SectionOneSize;
630   uint32_t SectionOneOffset;
631   uint32_t SectionOneRelocations;
632   uint32_t SectionTwoSize;
633   uint32_t SectionTwoOffset;
634   const ArrayRef<std::vector<UTF16>> StringTable;
635   std::vector<uint32_t> StringTableOffsets;
636   std::vector<uint32_t> DataOffsets;
637   std::vector<uint32_t> RelocationAddresses;
638 };
639 
640 WindowsResourceCOFFWriter::WindowsResourceCOFFWriter(
641     COFF::MachineTypes MachineType, const WindowsResourceParser &Parser,
642     Error &E)
643     : MachineType(MachineType), Resources(Parser.getTree()),
644       Data(Parser.getData()), StringTable(Parser.getStringTable()) {
645   performFileLayout();
646 
647   OutputBuffer = WritableMemoryBuffer::getNewMemBuffer(
648       FileSize, "internal .obj file created from .res files");
649 }
650 
651 void WindowsResourceCOFFWriter::performFileLayout() {
652   // Add size of COFF header.
653   FileSize = COFF::Header16Size;
654 
655   // one .rsrc section header for directory tree, another for resource data.
656   FileSize += 2 * COFF::SectionSize;
657 
658   performSectionOneLayout();
659   performSectionTwoLayout();
660 
661   // We have reached the address of the symbol table.
662   SymbolTableOffset = FileSize;
663 
664   FileSize += COFF::Symbol16Size;     // size of the @feat.00 symbol.
665   FileSize += 4 * COFF::Symbol16Size; // symbol + aux for each section.
666   FileSize += Data.size() * COFF::Symbol16Size; // 1 symbol per resource.
667   FileSize += 4; // four null bytes for the string table.
668 }
669 
670 void WindowsResourceCOFFWriter::performSectionOneLayout() {
671   SectionOneOffset = FileSize;
672 
673   SectionOneSize = Resources.getTreeSize();
674   uint32_t CurrentStringOffset = SectionOneSize;
675   uint32_t TotalStringTableSize = 0;
676   for (auto const &String : StringTable) {
677     StringTableOffsets.push_back(CurrentStringOffset);
678     uint32_t StringSize = String.size() * sizeof(UTF16) + sizeof(uint16_t);
679     CurrentStringOffset += StringSize;
680     TotalStringTableSize += StringSize;
681   }
682   SectionOneSize += alignTo(TotalStringTableSize, sizeof(uint32_t));
683 
684   // account for the relocations of section one.
685   SectionOneRelocations = FileSize + SectionOneSize;
686   FileSize += SectionOneSize;
687   FileSize +=
688       Data.size() * COFF::RelocationSize; // one relocation for each resource.
689   FileSize = alignTo(FileSize, SECTION_ALIGNMENT);
690 }
691 
692 void WindowsResourceCOFFWriter::performSectionTwoLayout() {
693   // add size of .rsrc$2 section, which contains all resource data on 8-byte
694   // alignment.
695   SectionTwoOffset = FileSize;
696   SectionTwoSize = 0;
697   for (auto const &Entry : Data) {
698     DataOffsets.push_back(SectionTwoSize);
699     SectionTwoSize += alignTo(Entry.size(), sizeof(uint64_t));
700   }
701   FileSize += SectionTwoSize;
702   FileSize = alignTo(FileSize, SECTION_ALIGNMENT);
703 }
704 
705 std::unique_ptr<MemoryBuffer>
706 WindowsResourceCOFFWriter::write(uint32_t TimeDateStamp) {
707   BufferStart = OutputBuffer->getBufferStart();
708 
709   writeCOFFHeader(TimeDateStamp);
710   writeFirstSectionHeader();
711   writeSecondSectionHeader();
712   writeFirstSection();
713   writeSecondSection();
714   writeSymbolTable();
715   writeStringTable();
716 
717   return std::move(OutputBuffer);
718 }
719 
720 // According to COFF specification, if the Src has a size equal to Dest,
721 // it's okay to *not* copy the trailing zero.
722 static void coffnamecpy(char (&Dest)[COFF::NameSize], StringRef Src) {
723   assert(Src.size() <= COFF::NameSize &&
724          "Src is not larger than COFF::NameSize");
725   strncpy(Dest, Src.data(), (size_t)COFF::NameSize);
726 }
727 
728 void WindowsResourceCOFFWriter::writeCOFFHeader(uint32_t TimeDateStamp) {
729   // Write the COFF header.
730   auto *Header = reinterpret_cast<coff_file_header *>(BufferStart);
731   Header->Machine = MachineType;
732   Header->NumberOfSections = 2;
733   Header->TimeDateStamp = TimeDateStamp;
734   Header->PointerToSymbolTable = SymbolTableOffset;
735   // One symbol for every resource plus 2 for each section and 1 for @feat.00
736   Header->NumberOfSymbols = Data.size() + 5;
737   Header->SizeOfOptionalHeader = 0;
738   // cvtres.exe sets 32BIT_MACHINE even for 64-bit machine types. Match it.
739   Header->Characteristics = COFF::IMAGE_FILE_32BIT_MACHINE;
740 }
741 
742 void WindowsResourceCOFFWriter::writeFirstSectionHeader() {
743   // Write the first section header.
744   CurrentOffset += sizeof(coff_file_header);
745   auto *SectionOneHeader =
746       reinterpret_cast<coff_section *>(BufferStart + CurrentOffset);
747   coffnamecpy(SectionOneHeader->Name, ".rsrc$01");
748   SectionOneHeader->VirtualSize = 0;
749   SectionOneHeader->VirtualAddress = 0;
750   SectionOneHeader->SizeOfRawData = SectionOneSize;
751   SectionOneHeader->PointerToRawData = SectionOneOffset;
752   SectionOneHeader->PointerToRelocations = SectionOneRelocations;
753   SectionOneHeader->PointerToLinenumbers = 0;
754   SectionOneHeader->NumberOfRelocations = Data.size();
755   SectionOneHeader->NumberOfLinenumbers = 0;
756   SectionOneHeader->Characteristics += COFF::IMAGE_SCN_CNT_INITIALIZED_DATA;
757   SectionOneHeader->Characteristics += COFF::IMAGE_SCN_MEM_READ;
758 }
759 
760 void WindowsResourceCOFFWriter::writeSecondSectionHeader() {
761   // Write the second section header.
762   CurrentOffset += sizeof(coff_section);
763   auto *SectionTwoHeader =
764       reinterpret_cast<coff_section *>(BufferStart + CurrentOffset);
765   coffnamecpy(SectionTwoHeader->Name, ".rsrc$02");
766   SectionTwoHeader->VirtualSize = 0;
767   SectionTwoHeader->VirtualAddress = 0;
768   SectionTwoHeader->SizeOfRawData = SectionTwoSize;
769   SectionTwoHeader->PointerToRawData = SectionTwoOffset;
770   SectionTwoHeader->PointerToRelocations = 0;
771   SectionTwoHeader->PointerToLinenumbers = 0;
772   SectionTwoHeader->NumberOfRelocations = 0;
773   SectionTwoHeader->NumberOfLinenumbers = 0;
774   SectionTwoHeader->Characteristics = COFF::IMAGE_SCN_CNT_INITIALIZED_DATA;
775   SectionTwoHeader->Characteristics += COFF::IMAGE_SCN_MEM_READ;
776 }
777 
778 void WindowsResourceCOFFWriter::writeFirstSection() {
779   // Write section one.
780   CurrentOffset += sizeof(coff_section);
781 
782   writeDirectoryTree();
783   writeDirectoryStringTable();
784   writeFirstSectionRelocations();
785 
786   CurrentOffset = alignTo(CurrentOffset, SECTION_ALIGNMENT);
787 }
788 
789 void WindowsResourceCOFFWriter::writeSecondSection() {
790   // Now write the .rsrc$02 section.
791   for (auto const &RawDataEntry : Data) {
792     llvm::copy(RawDataEntry, BufferStart + CurrentOffset);
793     CurrentOffset += alignTo(RawDataEntry.size(), sizeof(uint64_t));
794   }
795 
796   CurrentOffset = alignTo(CurrentOffset, SECTION_ALIGNMENT);
797 }
798 
799 void WindowsResourceCOFFWriter::writeSymbolTable() {
800   // Now write the symbol table.
801   // First, the feat symbol.
802   auto *Symbol = reinterpret_cast<coff_symbol16 *>(BufferStart + CurrentOffset);
803   coffnamecpy(Symbol->Name.ShortName, "@feat.00");
804   Symbol->Value = 0x11;
805   Symbol->SectionNumber = 0xffff;
806   Symbol->Type = COFF::IMAGE_SYM_DTYPE_NULL;
807   Symbol->StorageClass = COFF::IMAGE_SYM_CLASS_STATIC;
808   Symbol->NumberOfAuxSymbols = 0;
809   CurrentOffset += sizeof(coff_symbol16);
810 
811   // Now write the .rsrc1 symbol + aux.
812   Symbol = reinterpret_cast<coff_symbol16 *>(BufferStart + CurrentOffset);
813   coffnamecpy(Symbol->Name.ShortName, ".rsrc$01");
814   Symbol->Value = 0;
815   Symbol->SectionNumber = 1;
816   Symbol->Type = COFF::IMAGE_SYM_DTYPE_NULL;
817   Symbol->StorageClass = COFF::IMAGE_SYM_CLASS_STATIC;
818   Symbol->NumberOfAuxSymbols = 1;
819   CurrentOffset += sizeof(coff_symbol16);
820   auto *Aux = reinterpret_cast<coff_aux_section_definition *>(BufferStart +
821                                                               CurrentOffset);
822   Aux->Length = SectionOneSize;
823   Aux->NumberOfRelocations = Data.size();
824   Aux->NumberOfLinenumbers = 0;
825   Aux->CheckSum = 0;
826   Aux->NumberLowPart = 0;
827   Aux->Selection = 0;
828   CurrentOffset += sizeof(coff_aux_section_definition);
829 
830   // Now write the .rsrc2 symbol + aux.
831   Symbol = reinterpret_cast<coff_symbol16 *>(BufferStart + CurrentOffset);
832   coffnamecpy(Symbol->Name.ShortName, ".rsrc$02");
833   Symbol->Value = 0;
834   Symbol->SectionNumber = 2;
835   Symbol->Type = COFF::IMAGE_SYM_DTYPE_NULL;
836   Symbol->StorageClass = COFF::IMAGE_SYM_CLASS_STATIC;
837   Symbol->NumberOfAuxSymbols = 1;
838   CurrentOffset += sizeof(coff_symbol16);
839   Aux = reinterpret_cast<coff_aux_section_definition *>(BufferStart +
840                                                         CurrentOffset);
841   Aux->Length = SectionTwoSize;
842   Aux->NumberOfRelocations = 0;
843   Aux->NumberOfLinenumbers = 0;
844   Aux->CheckSum = 0;
845   Aux->NumberLowPart = 0;
846   Aux->Selection = 0;
847   CurrentOffset += sizeof(coff_aux_section_definition);
848 
849   // Now write a symbol for each relocation.
850   for (unsigned i = 0; i < Data.size(); i++) {
851     auto RelocationName = formatv("$R{0:X-6}", i & 0xffffff).sstr<COFF::NameSize>();
852     Symbol = reinterpret_cast<coff_symbol16 *>(BufferStart + CurrentOffset);
853     coffnamecpy(Symbol->Name.ShortName, RelocationName);
854     Symbol->Value = DataOffsets[i];
855     Symbol->SectionNumber = 2;
856     Symbol->Type = COFF::IMAGE_SYM_DTYPE_NULL;
857     Symbol->StorageClass = COFF::IMAGE_SYM_CLASS_STATIC;
858     Symbol->NumberOfAuxSymbols = 0;
859     CurrentOffset += sizeof(coff_symbol16);
860   }
861 }
862 
863 void WindowsResourceCOFFWriter::writeStringTable() {
864   // Just 4 null bytes for the string table.
865   auto COFFStringTable = reinterpret_cast<void *>(BufferStart + CurrentOffset);
866   memset(COFFStringTable, 0, 4);
867 }
868 
869 void WindowsResourceCOFFWriter::writeDirectoryTree() {
870   // Traverse parsed resource tree breadth-first and write the corresponding
871   // COFF objects.
872   std::queue<const WindowsResourceParser::TreeNode *> Queue;
873   Queue.push(&Resources);
874   uint32_t NextLevelOffset =
875       sizeof(coff_resource_dir_table) + (Resources.getStringChildren().size() +
876                                          Resources.getIDChildren().size()) *
877                                             sizeof(coff_resource_dir_entry);
878   std::vector<const WindowsResourceParser::TreeNode *> DataEntriesTreeOrder;
879   uint32_t CurrentRelativeOffset = 0;
880 
881   while (!Queue.empty()) {
882     auto CurrentNode = Queue.front();
883     Queue.pop();
884     auto *Table = reinterpret_cast<coff_resource_dir_table *>(BufferStart +
885                                                               CurrentOffset);
886     Table->Characteristics = CurrentNode->getCharacteristics();
887     Table->TimeDateStamp = 0;
888     Table->MajorVersion = CurrentNode->getMajorVersion();
889     Table->MinorVersion = CurrentNode->getMinorVersion();
890     auto &IDChildren = CurrentNode->getIDChildren();
891     auto &StringChildren = CurrentNode->getStringChildren();
892     Table->NumberOfNameEntries = StringChildren.size();
893     Table->NumberOfIDEntries = IDChildren.size();
894     CurrentOffset += sizeof(coff_resource_dir_table);
895     CurrentRelativeOffset += sizeof(coff_resource_dir_table);
896 
897     // Write the directory entries immediately following each directory table.
898     for (auto const &Child : StringChildren) {
899       auto *Entry = reinterpret_cast<coff_resource_dir_entry *>(BufferStart +
900                                                                 CurrentOffset);
901       Entry->Identifier.setNameOffset(
902           StringTableOffsets[Child.second->getStringIndex()]);
903       if (Child.second->checkIsDataNode()) {
904         Entry->Offset.DataEntryOffset = NextLevelOffset;
905         NextLevelOffset += sizeof(coff_resource_data_entry);
906         DataEntriesTreeOrder.push_back(Child.second.get());
907       } else {
908         Entry->Offset.SubdirOffset = NextLevelOffset + (1 << 31);
909         NextLevelOffset += sizeof(coff_resource_dir_table) +
910                            (Child.second->getStringChildren().size() +
911                             Child.second->getIDChildren().size()) *
912                                sizeof(coff_resource_dir_entry);
913         Queue.push(Child.second.get());
914       }
915       CurrentOffset += sizeof(coff_resource_dir_entry);
916       CurrentRelativeOffset += sizeof(coff_resource_dir_entry);
917     }
918     for (auto const &Child : IDChildren) {
919       auto *Entry = reinterpret_cast<coff_resource_dir_entry *>(BufferStart +
920                                                                 CurrentOffset);
921       Entry->Identifier.ID = Child.first;
922       if (Child.second->checkIsDataNode()) {
923         Entry->Offset.DataEntryOffset = NextLevelOffset;
924         NextLevelOffset += sizeof(coff_resource_data_entry);
925         DataEntriesTreeOrder.push_back(Child.second.get());
926       } else {
927         Entry->Offset.SubdirOffset = NextLevelOffset + (1 << 31);
928         NextLevelOffset += sizeof(coff_resource_dir_table) +
929                            (Child.second->getStringChildren().size() +
930                             Child.second->getIDChildren().size()) *
931                                sizeof(coff_resource_dir_entry);
932         Queue.push(Child.second.get());
933       }
934       CurrentOffset += sizeof(coff_resource_dir_entry);
935       CurrentRelativeOffset += sizeof(coff_resource_dir_entry);
936     }
937   }
938 
939   RelocationAddresses.resize(Data.size());
940   // Now write all the resource data entries.
941   for (auto DataNodes : DataEntriesTreeOrder) {
942     auto *Entry = reinterpret_cast<coff_resource_data_entry *>(BufferStart +
943                                                                CurrentOffset);
944     RelocationAddresses[DataNodes->getDataIndex()] = CurrentRelativeOffset;
945     Entry->DataRVA = 0; // Set to zero because it is a relocation.
946     Entry->DataSize = Data[DataNodes->getDataIndex()].size();
947     Entry->Codepage = 0;
948     Entry->Reserved = 0;
949     CurrentOffset += sizeof(coff_resource_data_entry);
950     CurrentRelativeOffset += sizeof(coff_resource_data_entry);
951   }
952 }
953 
954 void WindowsResourceCOFFWriter::writeDirectoryStringTable() {
955   // Now write the directory string table for .rsrc$01
956   uint32_t TotalStringTableSize = 0;
957   for (auto &String : StringTable) {
958     uint16_t Length = String.size();
959     support::endian::write16le(BufferStart + CurrentOffset, Length);
960     CurrentOffset += sizeof(uint16_t);
961     auto *Start = reinterpret_cast<UTF16 *>(BufferStart + CurrentOffset);
962     llvm::copy(String, Start);
963     CurrentOffset += Length * sizeof(UTF16);
964     TotalStringTableSize += Length * sizeof(UTF16) + sizeof(uint16_t);
965   }
966   CurrentOffset +=
967       alignTo(TotalStringTableSize, sizeof(uint32_t)) - TotalStringTableSize;
968 }
969 
970 void WindowsResourceCOFFWriter::writeFirstSectionRelocations() {
971 
972   // Now write the relocations for .rsrc$01
973   // Five symbols already in table before we start, @feat.00 and 2 for each
974   // .rsrc section.
975   uint32_t NextSymbolIndex = 5;
976   for (unsigned i = 0; i < Data.size(); i++) {
977     auto *Reloc =
978         reinterpret_cast<coff_relocation *>(BufferStart + CurrentOffset);
979     Reloc->VirtualAddress = RelocationAddresses[i];
980     Reloc->SymbolTableIndex = NextSymbolIndex++;
981     switch (MachineType) {
982     case COFF::IMAGE_FILE_MACHINE_ARMNT:
983       Reloc->Type = COFF::IMAGE_REL_ARM_ADDR32NB;
984       break;
985     case COFF::IMAGE_FILE_MACHINE_AMD64:
986       Reloc->Type = COFF::IMAGE_REL_AMD64_ADDR32NB;
987       break;
988     case COFF::IMAGE_FILE_MACHINE_I386:
989       Reloc->Type = COFF::IMAGE_REL_I386_DIR32NB;
990       break;
991     case COFF::IMAGE_FILE_MACHINE_ARM64:
992       Reloc->Type = COFF::IMAGE_REL_ARM64_ADDR32NB;
993       break;
994     default:
995       llvm_unreachable("unknown machine type");
996     }
997     CurrentOffset += sizeof(coff_relocation);
998   }
999 }
1000 
1001 Expected<std::unique_ptr<MemoryBuffer>>
1002 writeWindowsResourceCOFF(COFF::MachineTypes MachineType,
1003                          const WindowsResourceParser &Parser,
1004                          uint32_t TimeDateStamp) {
1005   Error E = Error::success();
1006   WindowsResourceCOFFWriter Writer(MachineType, Parser, E);
1007   if (E)
1008     return std::move(E);
1009   return Writer.write(TimeDateStamp);
1010 }
1011 
1012 } // namespace object
1013 } // namespace llvm
1014