3 files changed, 16 insertions, 4 deletions
diff --git a/core/fpdfapi/parser/cpdf_document.h b/core/fpdfapi/parser/cpdf_document.h
index d73dbc1df9..adf4eca75d 100644
--- a/core/fpdfapi/parser/cpdf_document.h
+++ b/core/fpdfapi/parser/cpdf_document.h
@@ -41,9 +41,15 @@ class JBig2_DocumentContext;
 
 class CPDF_Document : public CPDF_IndirectObjectHolder {
  public:
+  // Type from which the XFA extension can subclass itself.
+  class Extension {};
+
   explicit CPDF_Document(std::unique_ptr<CPDF_Parser> pParser);
   ~CPDF_Document() override;
 
+  Extension* GetExtension() const { return m_pExtension.Get(); }
+  void SetExtension(Extension* pExt) { m_pExtension = pExt; }
+
   CPDF_Parser* GetParser() const { return m_pParser.get(); }
   const CPDF_Dictionary* GetRoot() const { return m_pRootDict; }
   CPDF_Dictionary* GetRoot() { return m_pRootDict; }
@@ -153,6 +159,7 @@ class CPDF_Document : public CPDF_IndirectObjectHolder {
   std::unique_ptr<JBig2_DocumentContext> m_pCodecContext;
   std::unique_ptr<CPDF_LinkList> m_pLinksContext;
   std::vector<uint32_t> m_PageList;
+  UnownedPtr<Extension> m_pExtension;
 };
 
 #endif  // CORE_FPDFAPI_PARSER_CPDF_DOCUMENT_H_
diff --git a/fpdfsdk/cpdfsdk_helpers.cpp b/fpdfsdk/cpdfsdk_helpers.cpp
index 4c417e15e5..040b0f2e5a 100644
--- a/fpdfsdk/cpdfsdk_helpers.cpp
+++ b/fpdfsdk/cpdfsdk_helpers.cpp
@@ -166,9 +166,12 @@ CPDF_Document* CPDFDocumentFromFPDFDocument(FPDF_DOCUMENT doc) {
 
 FPDF_DOCUMENT FPDFDocumentFromCPDFDocument(CPDF_Document* doc) {
 #ifdef PDF_ENABLE_XFA
-  return doc ? FPDFDocumentFromUnderlying(
-                   new CPDFXFA_Context(pdfium::WrapUnique(doc)))
-             : nullptr;
+  if (!doc)
+    return nullptr;
+  if (!doc->GetExtension())
+    doc->SetExtension(new CPDFXFA_Context(pdfium::WrapUnique(doc)));
+  return FPDFDocumentFromUnderlying(
+      static_cast<CPDFXFA_Context*>(doc->GetExtension()));
 #else   // PDF_ENABLE_XFA
   return FPDFDocumentFromUnderlying(doc);
 #endif  // PDF_ENABLE_XFA
diff --git a/fpdfsdk/fpdfxfa/cpdfxfa_context.h b/fpdfsdk/fpdfxfa/cpdfxfa_context.h
index 188f1fbe2f..cb42de0576 100644
--- a/fpdfsdk/fpdfxfa/cpdfxfa_context.h
+++ b/fpdfsdk/fpdfxfa/cpdfxfa_context.h
@@ -10,6 +10,7 @@
 #include <memory>
 #include <vector>
 
+#include "core/fpdfapi/parser/cpdf_document.h"
 #include "core/fxcrt/fx_system.h"
 #include "core/fxcrt/observable.h"
 #include "core/fxcrt/unowned_ptr.h"
@@ -32,7 +33,8 @@ enum LoadStatus {
   FXFA_LOADSTATUS_CLOSED
 };
 
-class CPDFXFA_Context : public IXFA_AppProvider {
+class CPDFXFA_Context : public CPDF_Document::Extension,
+                        public IXFA_AppProvider {
  public:
   explicit CPDFXFA_Context(std::unique_ptr<CPDF_Document> pPDFDoc);
   ~CPDFXFA_Context() override;