diff --git a/core/fpdfapi/parser/cpdf_document.h b/core/fpdfapi/parser/cpdf_document.h
index adf4eca..4466df6 100644
--- a/core/fpdfapi/parser/cpdf_document.h
+++ b/core/fpdfapi/parser/cpdf_document.h
@@ -42,7 +42,10 @@
 class CPDF_Document : public CPDF_IndirectObjectHolder {
  public:
   // Type from which the XFA extension can subclass itself.
-  class Extension {};
+  class Extension {
+   public:
+    virtual ~Extension() {}
+  };
 
   explicit CPDF_Document(std::unique_ptr<CPDF_Parser> pParser);
   ~CPDF_Document() override;
diff --git a/fpdfsdk/cpdfsdk_helpers.cpp b/fpdfsdk/cpdfsdk_helpers.cpp
index 040b0f2..01d88f9 100644
--- a/fpdfsdk/cpdfsdk_helpers.cpp
+++ b/fpdfsdk/cpdfsdk_helpers.cpp
@@ -25,10 +25,6 @@
 
 constexpr char kQuadPoints[] = "QuadPoints";
 
-FPDF_DOCUMENT FPDFDocumentFromUnderlying(UnderlyingDocumentType* doc) {
-  return static_cast<FPDF_DOCUMENT>(doc);
-}
-
 bool RaiseUnSupportError(int nError) {
   CFSDK_UnsupportInfo_Adapter* pAdapter =
       CPDF_ModuleMgr::Get()->GetUnsupportInfoAdapter();
@@ -148,33 +144,20 @@
 
 }  // namespace
 
-UnderlyingDocumentType* UnderlyingFromFPDFDocument(FPDF_DOCUMENT doc) {
-  return static_cast<UnderlyingDocumentType*>(doc);
-}
-
 UnderlyingPageType* UnderlyingFromFPDFPage(FPDF_PAGE page) {
   return static_cast<UnderlyingPageType*>(page);
 }
 
 CPDF_Document* CPDFDocumentFromFPDFDocument(FPDF_DOCUMENT doc) {
-#ifdef PDF_ENABLE_XFA
-  return doc ? UnderlyingFromFPDFDocument(doc)->GetPDFDoc() : nullptr;
-#else   // PDF_ENABLE_XFA
-  return UnderlyingFromFPDFDocument(doc);
-#endif  // PDF_ENABLE_XFA
+  return reinterpret_cast<CPDF_Document*>(doc);
 }
 
 FPDF_DOCUMENT FPDFDocumentFromCPDFDocument(CPDF_Document* doc) {
 #ifdef PDF_ENABLE_XFA
-  if (!doc)
-    return nullptr;
-  if (!doc->GetExtension())
+  if (doc && !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
+  return reinterpret_cast<FPDF_DOCUMENT>(doc);
 }
 
 CPDF_Page* CPDFPageFromFPDFPage(FPDF_PAGE page) {
diff --git a/fpdfsdk/cpdfsdk_helpers.h b/fpdfsdk/cpdfsdk_helpers.h
index a9708cc..afb6dde 100644
--- a/fpdfsdk/cpdfsdk_helpers.h
+++ b/fpdfsdk/cpdfsdk_helpers.h
@@ -40,16 +40,12 @@
 // from fpdfsdk. For master, these are CPDF_ types, but for XFA, these are
 // CPDFXFA_ types.
 #ifndef PDF_ENABLE_XFA
-using UnderlyingDocumentType = CPDF_Document;
 using UnderlyingPageType = CPDF_Page;
 #else   // PDF_ENABLE_XFA
-using UnderlyingDocumentType = CPDFXFA_Context;
 using UnderlyingPageType = CPDFXFA_Page;
 #endif  // PDF_ENABLE_XFA
 
 // Conversions to/from underlying types.
-UnderlyingDocumentType* UnderlyingFromFPDFDocument(FPDF_DOCUMENT doc);
-
 UnderlyingPageType* UnderlyingFromFPDFPage(FPDF_PAGE page);
 
 // Conversions to/from FPDF_ types.
diff --git a/fpdfsdk/fpdf_catalog_unittest.cpp b/fpdfsdk/fpdf_catalog_unittest.cpp
index 648cb24..54eb4f7 100644
--- a/fpdfsdk/fpdf_catalog_unittest.cpp
+++ b/fpdfsdk/fpdf_catalog_unittest.cpp
@@ -12,6 +12,7 @@
 #include "core/fpdfapi/parser/cpdf_parser.h"
 #include "core/fpdfapi/parser/cpdf_string.h"
 #include "fpdfsdk/cpdfsdk_helpers.h"
+#include "public/cpp/fpdf_scopers.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "testing/test_support.h"
 
@@ -29,31 +30,12 @@
   }
 };
 
-#ifdef PDF_ENABLE_XFA
-class CPDF_TestXFAContext : public CPDFXFA_Context {
- public:
-  CPDF_TestXFAContext()
-      : CPDFXFA_Context(pdfium::MakeUnique<CPDF_TestDocument>()) {}
-
-  void SetRoot(CPDF_Dictionary* root) {
-    static_cast<CPDF_TestDocument*>(GetPDFDoc())->SetRoot(root);
-  }
-
-  CPDF_IndirectObjectHolder* GetHolder() { return GetPDFDoc(); }
-};
-using CPDF_TestPdfDocument = CPDF_TestXFAContext;
-#else   // PDF_ENABLE_XFA
-using CPDF_TestPdfDocument = CPDF_TestDocument;
-#endif  // PDF_ENABLE_XFA
-
 class PDFCatalogTest : public testing::Test {
  public:
   void SetUp() override {
     CPDF_ModuleMgr::Get()->Init();
-
-    m_pDoc = pdfium::MakeUnique<CPDF_TestPdfDocument>();
-
-    // Setup the root directory.
+    auto pTestDoc = pdfium::MakeUnique<CPDF_TestDocument>();
+    m_pDoc.reset(FPDFDocumentFromCPDFDocument(pTestDoc.release()));
     m_pRootObj = pdfium::MakeUnique<CPDF_Dictionary>();
   }
 
@@ -63,7 +45,7 @@
   }
 
  protected:
-  std::unique_ptr<CPDF_TestPdfDocument> m_pDoc;
+  ScopedFPDFDocument m_pDoc;
   std::unique_ptr<CPDF_Dictionary> m_pRootObj;
 };
 
@@ -71,12 +53,15 @@
   // Null doc
   EXPECT_FALSE(FPDFCatalog_IsTagged(nullptr));
 
+  CPDF_TestDocument* pTestDoc = static_cast<CPDF_TestDocument*>(
+      CPDFDocumentFromFPDFDocument(m_pDoc.get()));
+
   // No root
-  m_pDoc->SetRoot(nullptr);
+  pTestDoc->SetRoot(nullptr);
   EXPECT_FALSE(FPDFCatalog_IsTagged(m_pDoc.get()));
 
   // Empty root
-  m_pDoc->SetRoot(m_pRootObj.get());
+  pTestDoc->SetRoot(m_pRootObj.get());
   EXPECT_FALSE(FPDFCatalog_IsTagged(m_pDoc.get()));
 
   // Root with other key
diff --git a/fpdfsdk/fpdf_doc_unittest.cpp b/fpdfsdk/fpdf_doc_unittest.cpp
index bfd4262..01c7905 100644
--- a/fpdfsdk/fpdf_doc_unittest.cpp
+++ b/fpdfsdk/fpdf_doc_unittest.cpp
@@ -18,6 +18,7 @@
 #include "core/fpdfapi/parser/cpdf_string.h"
 #include "core/fpdfdoc/cpdf_dest.h"
 #include "fpdfsdk/cpdfsdk_helpers.h"
+#include "public/cpp/fpdf_scopers.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "testing/test_support.h"
 #include "third_party/base/ptr_util.h"
@@ -34,23 +35,6 @@
   CPDF_IndirectObjectHolder* GetHolder() { return this; }
 };
 
-#ifdef PDF_ENABLE_XFA
-class CPDF_TestXFAContext : public CPDFXFA_Context {
- public:
-  CPDF_TestXFAContext()
-      : CPDFXFA_Context(pdfium::MakeUnique<CPDF_TestDocument>()) {}
-
-  void SetRoot(CPDF_Dictionary* root) {
-    static_cast<CPDF_TestDocument*>(GetPDFDoc())->SetRoot(root);
-  }
-
-  CPDF_IndirectObjectHolder* GetHolder() { return GetPDFDoc(); }
-};
-using CPDF_TestPdfDocument = CPDF_TestXFAContext;
-#else   // PDF_ENABLE_XFA
-using CPDF_TestPdfDocument = CPDF_TestDocument;
-#endif  // PDF_ENABLE_XFA
-
 class PDFDocTest : public testing::Test {
  public:
   struct DictObjInfo {
@@ -60,13 +44,11 @@
 
   void SetUp() override {
     CPDF_ModuleMgr::Get()->Init();
-
-    m_pDoc = pdfium::MakeUnique<CPDF_TestPdfDocument>();
-    m_pIndirectObjs = m_pDoc->GetHolder();
-
-    // Setup the root directory.
+    auto pTestDoc = pdfium::MakeUnique<CPDF_TestDocument>();
+    m_pIndirectObjs = pTestDoc->GetHolder();
     m_pRootObj = pdfium::MakeUnique<CPDF_Dictionary>();
-    m_pDoc->SetRoot(m_pRootObj.get());
+    pTestDoc->SetRoot(m_pRootObj.get());
+    m_pDoc.reset(FPDFDocumentFromCPDFDocument(pTestDoc.release()));
   }
 
   void TearDown() override {
@@ -87,7 +69,7 @@
   }
 
  protected:
-  std::unique_ptr<CPDF_TestPdfDocument> m_pDoc;
+  ScopedFPDFDocument m_pDoc;
   UnownedPtr<CPDF_IndirectObjectHolder> m_pIndirectObjs;
   std::unique_ptr<CPDF_Dictionary> m_pRootObj;
 };
diff --git a/fpdfsdk/fpdf_editpage.cpp b/fpdfsdk/fpdf_editpage.cpp
index fc6b21f..3acbe76 100644
--- a/fpdfsdk/fpdf_editpage.cpp
+++ b/fpdfsdk/fpdf_editpage.cpp
@@ -153,8 +153,17 @@
 
 FPDF_EXPORT void FPDF_CALLCONV FPDFPage_Delete(FPDF_DOCUMENT document,
                                                int page_index) {
-  if (UnderlyingDocumentType* pDoc = UnderlyingFromFPDFDocument(document))
-    pDoc->DeletePage(page_index);
+  auto* pDoc = CPDFDocumentFromFPDFDocument(document);
+  if (!pDoc)
+    return;
+#ifdef PDF_ENABLE_XFA
+  CPDFXFA_Context* pContext =
+      static_cast<CPDFXFA_Context*>(pDoc->GetExtension());
+  if (pContext)
+    pContext->DeletePage(page_index);
+#else
+  pDoc->DeletePage(page_index);
+#endif
 }
 
 FPDF_EXPORT FPDF_PAGE FPDF_CALLCONV FPDFPage_New(FPDF_DOCUMENT document,
@@ -176,7 +185,7 @@
 
 #ifdef PDF_ENABLE_XFA
   auto pXFAPage = pdfium::MakeRetain<CPDFXFA_Page>(
-      static_cast<CPDFXFA_Context*>(document), page_index);
+      static_cast<CPDFXFA_Context*>(pDoc->GetExtension()), page_index);
   pXFAPage->LoadPDFPage(pPageDict);
   return pXFAPage.Leak();  // Caller takes ownership.
 #else  // PDF_ENABLE_XFA
diff --git a/fpdfsdk/fpdf_formfill.cpp b/fpdfsdk/fpdf_formfill.cpp
index 5d762ee..3170627 100644
--- a/fpdfsdk/fpdf_formfill.cpp
+++ b/fpdfsdk/fpdf_formfill.cpp
@@ -287,7 +287,7 @@
   if (!formInfo || formInfo->version != kRequiredVersion)
     return nullptr;
 
-  UnderlyingDocumentType* pDocument = UnderlyingFromFPDFDocument(document);
+  auto* pDocument = CPDFDocumentFromFPDFDocument(document);
   if (!pDocument)
     return nullptr;
 
@@ -295,15 +295,16 @@
   // If the CPDFXFA_Context has a FormFillEnvironment already then we've done
   // this and can just return the old Env. Otherwise, we'll end up setting a new
   // environment into the XFADocument and, that could get weird.
-  if (pDocument->GetFormFillEnv())
-    return pDocument->GetFormFillEnv();
+  auto* pContext = static_cast<CPDFXFA_Context*>(pDocument->GetExtension());
+  if (pContext->GetFormFillEnv())
+    return pContext->GetFormFillEnv();
 #endif
 
   auto pFormFillEnv = pdfium::MakeUnique<CPDFSDK_FormFillEnvironment>(
       CPDFDocumentFromFPDFDocument(document), formInfo);
 
 #ifdef PDF_ENABLE_XFA
-  pDocument->SetFormFillEnv(pFormFillEnv.get());
+  pContext->SetFormFillEnv(pFormFillEnv.get());
 #endif  // PDF_ENABLE_XFA
 
   return pFormFillEnv.release();  // Caller takes ownership.
diff --git a/fpdfsdk/fpdf_save.cpp b/fpdfsdk/fpdf_save.cpp
index 7da48f3..2844f5a 100644
--- a/fpdfsdk/fpdf_save.cpp
+++ b/fpdfsdk/fpdf_save.cpp
@@ -227,7 +227,8 @@
     return 0;
 
 #ifdef PDF_ENABLE_XFA
-  CPDFXFA_Context* pContext = static_cast<CPDFXFA_Context*>(document);
+  CPDFXFA_Context* pContext =
+      static_cast<CPDFXFA_Context*>(pPDFDoc->GetExtension());
   std::vector<RetainPtr<IFX_SeekableStream>> fileList;
   SendPreSaveToXFADoc(pContext, &fileList);
 #endif  // PDF_ENABLE_XFA
diff --git a/fpdfsdk/fpdf_view.cpp b/fpdfsdk/fpdf_view.cpp
index 0015716..00c59bb 100644
--- a/fpdfsdk/fpdf_view.cpp
+++ b/fpdfsdk/fpdf_view.cpp
@@ -256,7 +256,9 @@
 
 #ifdef PDF_ENABLE_XFA
 FPDF_EXPORT FPDF_BOOL FPDF_CALLCONV FPDF_LoadXFA(FPDF_DOCUMENT document) {
-  return document && static_cast<CPDFXFA_Context*>(document)->LoadXFADoc();
+  auto* pDoc = CPDFDocumentFromFPDFDocument(document);
+  return pDoc &&
+         static_cast<CPDFXFA_Context*>(pDoc->GetExtension())->LoadXFADoc();
 }
 #endif  // PDF_ENABLE_XFA
 
@@ -320,21 +322,30 @@
 }
 
 FPDF_EXPORT int FPDF_CALLCONV FPDF_GetPageCount(FPDF_DOCUMENT document) {
-  UnderlyingDocumentType* pDoc = UnderlyingFromFPDFDocument(document);
-  return pDoc ? pDoc->GetPageCount() : 0;
+  auto* pDoc = CPDFDocumentFromFPDFDocument(document);
+  if (!pDoc)
+    return 0;
+#ifdef PDF_ENABLE_XFA
+  auto* pContext = static_cast<CPDFXFA_Context*>(pDoc->GetExtension());
+  return pContext ? pContext->GetPageCount() : 0;
+#else
+  return pDoc->GetPageCount();
+#endif
 }
 
 FPDF_EXPORT FPDF_PAGE FPDF_CALLCONV FPDF_LoadPage(FPDF_DOCUMENT document,
                                                   int page_index) {
-  UnderlyingDocumentType* pDoc = UnderlyingFromFPDFDocument(document);
+  auto* pDoc = CPDFDocumentFromFPDFDocument(document);
   if (!pDoc)
     return nullptr;
 
-  if (page_index < 0 || page_index >= pDoc->GetPageCount())
+  if (page_index < 0 || page_index >= FPDF_GetPageCount(document))
     return nullptr;
 
 #ifdef PDF_ENABLE_XFA
-  return pDoc->GetXFAPage(page_index).Leak();
+  return static_cast<CPDFXFA_Context*>(pDoc->GetExtension())
+      ->GetXFAPage(page_index)
+      .Leak();
 #else   // PDF_ENABLE_XFA
   CPDF_Dictionary* pDict = pDoc->GetPage(page_index);
   if (!pDict)
@@ -740,7 +751,14 @@
 }
 
 FPDF_EXPORT void FPDF_CALLCONV FPDF_CloseDocument(FPDF_DOCUMENT document) {
-  delete UnderlyingFromFPDFDocument(document);
+  auto* pDoc = CPDFDocumentFromFPDFDocument(document);
+#if PDF_ENABLE_XFA
+  // Deleting the extension will delete the document
+  if (pDoc)
+    delete pDoc->GetExtension();
+#else
+  delete pDoc;
+#endif
 }
 
 FPDF_EXPORT unsigned long FPDF_CALLCONV FPDF_GetLastError() {
@@ -918,15 +936,16 @@
                                                       int page_index,
                                                       double* width,
                                                       double* height) {
-  UnderlyingDocumentType* pDoc = UnderlyingFromFPDFDocument(document);
+  auto* pDoc = CPDFDocumentFromFPDFDocument(document);
   if (!pDoc)
     return false;
 
 #ifdef PDF_ENABLE_XFA
-  int count = pDoc->GetPageCount();
-  if (page_index < 0 || page_index >= count)
+  if (page_index < 0 || page_index >= FPDF_GetPageCount(document))
     return false;
-  RetainPtr<CPDFXFA_Page> pPage = pDoc->GetXFAPage(page_index);
+  RetainPtr<CPDFXFA_Page> pPage =
+      static_cast<CPDFXFA_Context*>(pDoc->GetExtension())
+          ->GetXFAPage(page_index);
   if (!pPage)
     return false;
   *width = pPage->GetPageWidth();
diff --git a/public/fpdfview.h b/public/fpdfview.h
index 20f710f..9ebfd8d 100644
--- a/public/fpdfview.h
+++ b/public/fpdfview.h
@@ -40,7 +40,7 @@
 typedef void* FPDF_BOOKMARK;
 typedef void* FPDF_CLIPPATH;
 typedef void* FPDF_DEST;
-typedef void* FPDF_DOCUMENT;
+typedef struct fpdf_document_t__* FPDF_DOCUMENT;
 typedef void* FPDF_FONT;
 typedef void* FPDF_LINK;
 typedef void* FPDF_PAGE;
diff --git a/testing/xfa_js_embedder_test.cpp b/testing/xfa_js_embedder_test.cpp
index 995c0ed..375497f 100644
--- a/testing/xfa_js_embedder_test.cpp
+++ b/testing/xfa_js_embedder_test.cpp
@@ -44,7 +44,9 @@
 }
 
 CXFA_Document* XFAJSEmbedderTest::GetXFADocument() {
-  return UnderlyingFromFPDFDocument(document())->GetXFADoc()->GetXFADoc();
+  auto* pDoc = CPDFDocumentFromFPDFDocument(document());
+  auto* pContext = static_cast<CPDFXFA_Context*>(pDoc->GetExtension());
+  return pContext->GetXFADoc()->GetXFADoc();
 }
 
 bool XFAJSEmbedderTest::OpenDocumentWithOptions(