4 files changed, 18 insertions, 4 deletions
diff --git a/fpdfsdk/cpdfsdk_helpers.cpp b/fpdfsdk/cpdfsdk_helpers.cpp
index 210fcf8b84..f7a875a3e9 100644
--- a/fpdfsdk/cpdfsdk_helpers.cpp
+++ b/fpdfsdk/cpdfsdk_helpers.cpp
@@ -156,10 +156,6 @@ CPDF_Document* CPDFDocumentFromFPDFDocument(FPDF_DOCUMENT doc) {
 }
 
 FPDF_DOCUMENT FPDFDocumentFromCPDFDocument(CPDF_Document* doc) {
-#ifdef PDF_ENABLE_XFA
-  if (doc && !doc->GetExtension())
-    doc->SetExtension(pdfium::MakeUnique<CPDFXFA_Context>(doc));
-#endif  // PDF_ENABLE_XFA
   return reinterpret_cast<FPDF_DOCUMENT>(doc);
 }
 
diff --git a/fpdfsdk/fpdf_dataavail.cpp b/fpdfsdk/fpdf_dataavail.cpp
index 6f201d69f1..d7bcf8f5c7 100644
--- a/fpdfsdk/fpdf_dataavail.cpp
+++ b/fpdfsdk/fpdf_dataavail.cpp
@@ -17,6 +17,10 @@
 #include "public/fpdf_formfill.h"
 #include "third_party/base/ptr_util.h"
 
+#ifdef PDF_ENABLE_XFA
+#include "fpdfsdk/fpdfxfa/cpdfxfa_context.h"
+#endif  // PDF_ENABLE_XFA
+
 // These checks are here because core/ and public/ cannot depend on each other.
 static_assert(CPDF_DataAvail::DataError == PDF_DATA_ERROR,
               "CPDF_DataAvail::DataError value mismatch");
@@ -153,6 +157,11 @@ FPDFAvail_GetDocument(FPDF_AVAIL avail, FPDF_BYTESTRING password) {
     ProcessParseError(error);
     return nullptr;
   }
+
+#ifdef PDF_ENABLE_XFA
+  document->SetExtension(pdfium::MakeUnique<CPDFXFA_Context>(document.get()));
+#endif  // PDF_ENABLE_XFA
+
   CheckUnSupportError(document.get(), FPDF_ERR_SUCCESS);
   return FPDFDocumentFromCPDFDocument(document.release());
 }
diff --git a/fpdfsdk/fpdf_editpage.cpp b/fpdfsdk/fpdf_editpage.cpp
index 9a4d5c72a1..9198e1d2e4 100644
--- a/fpdfsdk/fpdf_editpage.cpp
+++ b/fpdfsdk/fpdf_editpage.cpp
@@ -150,6 +150,10 @@ FPDF_EXPORT FPDF_DOCUMENT FPDF_CALLCONV FPDF_CreateNewDocument() {
     pInfoDict->SetNewFor<CPDF_String>("Creator", L"PDFium");
   }
 
+#ifdef PDF_ENABLE_XFA
+  pDoc->SetExtension(pdfium::MakeUnique<CPDFXFA_Context>(pDoc.get()));
+#endif  // PDF_ENABLE_XFA
+
   // Caller takes ownership of pDoc.
   return FPDFDocumentFromCPDFDocument(pDoc.release());
 }
diff --git a/fpdfsdk/fpdf_view.cpp b/fpdfsdk/fpdf_view.cpp
index ef242295b7..f0ceb1f378 100644
--- a/fpdfsdk/fpdf_view.cpp
+++ b/fpdfsdk/fpdf_view.cpp
@@ -152,6 +152,11 @@ FPDF_DOCUMENT LoadDocumentImpl(
     return nullptr;
   }
   CheckUnSupportError(pDocument.get(), error);
+
+#ifdef PDF_ENABLE_XFA
+  pDocument->SetExtension(pdfium::MakeUnique<CPDFXFA_Context>(pDocument.get()));
+#endif  // PDF_ENABLE_XFA
+
   return FPDFDocumentFromCPDFDocument(pDocument.release());
 }