7 files changed, 29 insertions, 11 deletions
diff --git a/fpdfsdk/cpdfsdk_pageview.cpp b/fpdfsdk/cpdfsdk_pageview.cpp
index e648e596c4..87856fb0f0 100644
--- a/fpdfsdk/cpdfsdk_pageview.cpp
+++ b/fpdfsdk/cpdfsdk_pageview.cpp
@@ -89,7 +89,8 @@ void CPDFSDK_PageView::PageView_OnDraw(CFX_RenderDevice* pDevice,
   if (!pPage)
     return;
 
-  if (pPage->GetContext()->GetFormType() == FormType::kXFAFull) {
+  auto* pContext = static_cast<CPDFXFA_Context*>(pPage->GetDocumentExtension());
+  if (pContext->GetFormType() == FormType::kXFAFull) {
     CFX_RectF rectClip(
         static_cast<float>(pClip.left), static_cast<float>(pClip.top),
         static_cast<float>(pClip.Width()), static_cast<float>(pClip.Height()));
@@ -177,7 +178,11 @@ bool CPDFSDK_PageView::DeleteAnnot(CPDFSDK_Annot* pAnnot) {
     return false;
 
   CPDFXFA_Page* pPage = pAnnot->GetPDFXFAPage();
-  if (!pPage || !pPage->GetContext()->ContainsXFAForm())
+  if (!pPage)
+    return false;
+
+  auto* pContext = static_cast<CPDFXFA_Context*>(pPage->GetDocumentExtension());
+  if (!pContext->ContainsXFAForm())
     return false;
 
   CPDFSDK_Annot::ObservedPtr pObserved(pAnnot);
@@ -204,7 +209,7 @@ bool CPDFSDK_PageView::DeleteAnnot(CPDFSDK_Annot* pAnnot) {
 CPDF_Document* CPDFSDK_PageView::GetPDFDocument() {
   if (m_page) {
 #ifdef PDF_ENABLE_XFA
-    return m_page->GetContext()->GetPDFDoc();
+    return m_page->GetDocumentExtension()->GetPDFDoc();
 #else   // PDF_ENABLE_XFA
     return m_page->m_pDocument.Get();
 #endif  // PDF_ENABLE_XFA
@@ -530,7 +535,9 @@ int CPDFSDK_PageView::GetPageIndex() const {
     return -1;
 
 #ifdef PDF_ENABLE_XFA
-  switch (m_page->GetContext()->GetFormType()) {
+  auto* pContext =
+      static_cast<CPDFXFA_Context*>(m_page->GetDocumentExtension());
+  switch (pContext->GetFormType()) {
     case FormType::kXFAFull: {
       CXFA_FFPageView* pPageView = m_page->GetXFAPageView();
       return pPageView ? pPageView->GetPageIndex() : -1;
diff --git a/fpdfsdk/fpdf_formfill.cpp b/fpdfsdk/fpdf_formfill.cpp
index 202835cb57..b99bd5f40d 100644
--- a/fpdfsdk/fpdf_formfill.cpp
+++ b/fpdfsdk/fpdf_formfill.cpp
@@ -138,10 +138,10 @@ void FFLCommon(FPDF_FORMHANDLE hHandle,
     return;
 
 #ifdef PDF_ENABLE_XFA
-  CPDFXFA_Context* pContext = pPage->GetContext();
-  if (!pContext)
+  CPDF_Document::Extension* pExtension = pPage->GetDocumentExtension();
+  if (!pExtension)
     return;
-  CPDF_Document* pPDFDoc = pContext->GetPDFDoc();
+  CPDF_Document* pPDFDoc = pExtension->GetPDFDoc();
   if (!pPDFDoc)
     return;
   CPDFSDK_FormFillEnvironment* pFormFillEnv =
diff --git a/fpdfsdk/fpdf_text.cpp b/fpdfsdk/fpdf_text.cpp
index dade84a005..c06885a1a7 100644
--- a/fpdfsdk/fpdf_text.cpp
+++ b/fpdfsdk/fpdf_text.cpp
@@ -52,7 +52,7 @@ FPDF_EXPORT FPDF_TEXTPAGE FPDF_CALLCONV FPDFText_LoadPage(FPDF_PAGE page) {
 
 #ifdef PDF_ENABLE_XFA
   CPDFXFA_Page* pPage = UnderlyingFromFPDFPage(page);
-  CPDFXFA_Context* pContext = pPage->GetContext();
+  auto* pContext = static_cast<CPDFXFA_Context*>(pPage->GetDocumentExtension());
   CPDF_ViewerPreferences viewRef(pContext->GetPDFDoc());
 #else  // PDF_ENABLE_XFA
   CPDF_ViewerPreferences viewRef(pPDFPage->m_pDocument.Get());
diff --git a/fpdfsdk/fpdfxfa/cpdfxfa_context.cpp b/fpdfsdk/fpdfxfa/cpdfxfa_context.cpp
index 97288c67f0..05bb56ca99 100644
--- a/fpdfsdk/fpdfxfa/cpdfxfa_context.cpp
+++ b/fpdfsdk/fpdfxfa/cpdfxfa_context.cpp
@@ -180,6 +180,10 @@ RetainPtr<CPDFXFA_Page> CPDFXFA_Context::GetXFAPage(
   return nullptr;
 }
 
+CPDF_Document* CPDFXFA_Context::GetPDFDoc() const {
+  return m_pPDFDoc.get();
+}
+
 void CPDFXFA_Context::DeletePage(int page_index) {
   // Delete from the document first because, if GetPage was never called for
   // this |page_index| then |m_XFAPageList| may have size < |page_index| even
diff --git a/fpdfsdk/fpdfxfa/cpdfxfa_context.h b/fpdfsdk/fpdfxfa/cpdfxfa_context.h
index 23f1ac2220..c7fe02d8fb 100644
--- a/fpdfsdk/fpdfxfa/cpdfxfa_context.h
+++ b/fpdfsdk/fpdfxfa/cpdfxfa_context.h
@@ -40,7 +40,6 @@ class CPDFXFA_Context : public CPDF_Document::Extension,
   ~CPDFXFA_Context() override;
 
   bool LoadXFADoc();
-  CPDF_Document* GetPDFDoc() { return m_pPDFDoc.get(); }
   CXFA_FFDoc* GetXFADoc() { return m_pXFADoc.get(); }
   CXFA_FFDocView* GetXFADocView() { return m_pXFADocView.Get(); }
   FormType GetFormType() const { return m_FormType; }
@@ -48,7 +47,6 @@ class CPDFXFA_Context : public CPDF_Document::Extension,
     return m_FormType == FormType::kXFAFull ||
            m_FormType == FormType::kXFAForeground;
   }
-  CJS_Runtime* GetCJSRuntime() const;
   CXFA_FFApp* GetXFAApp() { return m_pXFAApp.get(); }
 
   CPDFSDK_FormFillEnvironment* GetFormFillEnv() const {
@@ -61,6 +59,7 @@ class CPDFXFA_Context : public CPDF_Document::Extension,
   void ClearChangeMark();
 
   // CPDF_Document::Extension:
+  CPDF_Document* GetPDFDoc() const override;
   int GetPageCount() const override;
   void DeletePage(int page_index) override;
 
@@ -108,6 +107,7 @@ class CPDFXFA_Context : public CPDF_Document::Extension,
   }
 
  private:
+  CJS_Runtime* GetCJSRuntime() const;
   void CloseXFADoc();
 
   FormType m_FormType = FormType::kNone;
diff --git a/fpdfsdk/fpdfxfa/cpdfxfa_page.cpp b/fpdfsdk/fpdfxfa/cpdfxfa_page.cpp
index 8268ed8759..809214dec6 100644
--- a/fpdfsdk/fpdfxfa/cpdfxfa_page.cpp
+++ b/fpdfsdk/fpdfxfa/cpdfxfa_page.cpp
@@ -86,6 +86,10 @@ bool CPDFXFA_Page::LoadPDFPage(CPDF_Dictionary* pageDict) {
   return true;
 }
 
+CPDF_Document::Extension* CPDFXFA_Page::GetDocumentExtension() const {
+  return m_pContext.Get();
+}
+
 float CPDFXFA_Page::GetPageWidth() const {
   if (!m_pPDFPage && !m_pXFAPageView)
     return 0.0f;
diff --git a/fpdfsdk/fpdfxfa/cpdfxfa_page.h b/fpdfsdk/fpdfxfa/cpdfxfa_page.h
index 030b0684b6..4f4d6b0a22 100644
--- a/fpdfsdk/fpdfxfa/cpdfxfa_page.h
+++ b/fpdfsdk/fpdfxfa/cpdfxfa_page.h
@@ -10,6 +10,7 @@
 #include <memory>
 
 #include "core/fpdfapi/page/cpdf_page.h"
+#include "core/fpdfapi/parser/cpdf_document.h"
 #include "core/fxcrt/fx_coordinates.h"
 #include "core/fxcrt/fx_system.h"
 #include "core/fxcrt/retain_ptr.h"
@@ -28,7 +29,9 @@ class CPDFXFA_Page : public CPDF_Page::Extension {
   bool LoadPage();
   bool LoadPDFPage(CPDF_Dictionary* pageDict);
 
-  CPDFXFA_Context* GetContext() const { return m_pContext.Get(); }
+  // CPDF_Page::Extension:
+  CPDF_Document::Extension* GetDocumentExtension() const override;
+
   int GetPageIndex() const { return m_iPageIndex; }
   CPDF_Page* GetPDFPage() const { return m_pPDFPage.get(); }
   CXFA_FFPageView* GetXFAPageView() const { return m_pXFAPageView; }