4 files changed, 14 insertions, 3 deletions
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; }