7 files changed, 30 insertions, 25 deletions
diff --git a/fpdfsdk/cpdfsdk_document.cpp b/fpdfsdk/cpdfsdk_document.cpp
index 544c1e6e19..a0febafd48 100644
--- a/fpdfsdk/cpdfsdk_document.cpp
+++ b/fpdfsdk/cpdfsdk_document.cpp
@@ -39,14 +39,17 @@ CPDFSDK_Document::CPDFSDK_Document(UnderlyingDocumentType* pDoc,
 CPDFSDK_Document::~CPDFSDK_Document() {
   m_bBeingDestroyed = TRUE;
 
+  ClearAllFocusedAnnots();
+  for (auto& it : m_pageMap)
+    delete it.second;
+  m_pageMap.clear();
+}
+
+void CPDFSDK_Document::ClearAllFocusedAnnots() {
   for (auto& it : m_pageMap) {
     if (it.second->IsValidSDKAnnot(GetFocusAnnot()))
       KillFocusAnnot(0);
   }
-
-  for (auto& it : m_pageMap)
-    delete it.second;
-  m_pageMap.clear();
 }
 
 CPDFSDK_PageView* CPDFSDK_Document::GetPageView(
diff --git a/fpdfsdk/cpdfsdk_document.h b/fpdfsdk/cpdfsdk_document.h
index 697ee48f9d..561bfc198d 100644
--- a/fpdfsdk/cpdfsdk_document.h
+++ b/fpdfsdk/cpdfsdk_document.h
@@ -69,6 +69,7 @@ class CPDFSDK_Document : public CFX_Observable<CPDFSDK_Document> {
   CPDFSDK_Annot* GetFocusAnnot() { return m_pFocusAnnot.Get(); }
   FX_BOOL SetFocusAnnot(CPDFSDK_Annot::ObservedPtr* pAnnot);
   FX_BOOL KillFocusAnnot(uint32_t nFlag);
+  void ClearAllFocusedAnnots();
 
   FX_BOOL ExtractPages(const std::vector<uint16_t>& arrExtraPages,
                        CPDF_Document* pDstDoc);
diff --git a/fpdfsdk/cpdfsdk_environment.h b/fpdfsdk/cpdfsdk_environment.h
index 9943d38c67..bed757c2ba 100644
--- a/fpdfsdk/cpdfsdk_environment.h
+++ b/fpdfsdk/cpdfsdk_environment.h
@@ -146,8 +146,7 @@ class CPDFSDK_Environment final {
   void JS_docgotoPage(int nPageNum);
 
   FX_BOOL IsJSInitiated() const { return m_pInfo && m_pInfo->m_pJsPlatform; }
-  void SetSDKDocument(CPDFSDK_Document* pFXDoc) { m_pSDKDoc = pFXDoc; }
-  CPDFSDK_Document* GetSDKDocument() const { return m_pSDKDoc; }
+  CPDFSDK_Document* GetSDKDocument() const { return m_pSDKDoc.get(); }
   UnderlyingDocumentType* GetUnderlyingDocument() const {
     return m_pUnderlyingDoc;
   }
@@ -166,10 +165,7 @@ class CPDFSDK_Environment final {
   std::unique_ptr<CPDFSDK_ActionHandler> m_pActionHandler;
   std::unique_ptr<IJS_Runtime> m_pJSRuntime;
   FPDF_FORMFILLINFO* const m_pInfo;
-  // Ownership of |m_pSDKDoc| depends on if this is XFA. If we're in XFA then
-  // the object is owned by the CPDFXFA_Document. In non-xfa then we own
-  // the pointer.
-  CPDFSDK_Document* m_pSDKDoc;
+  std::unique_ptr<CPDFSDK_Document> m_pSDKDoc;
   UnderlyingDocumentType* const m_pUnderlyingDoc;
   std::unique_ptr<CFFL_InteractiveFormFiller> m_pFormFiller;
   std::unique_ptr<CFX_SystemHandler> m_pSysHandler;
diff --git a/fpdfsdk/fpdfformfill.cpp b/fpdfsdk/fpdfformfill.cpp
index 6fea1af42b..904a3c867e 100644
--- a/fpdfsdk/fpdfformfill.cpp
+++ b/fpdfsdk/fpdfformfill.cpp
@@ -254,8 +254,7 @@ FPDFDOC_InitFormFillEnvironment(FPDF_DOCUMENT document,
   CPDFSDK_Environment* pEnv = new CPDFSDK_Environment(pDocument, formInfo);
 
 #ifdef PDF_ENABLE_XFA
-  // Ownership of the SDKDocument is passed to the CPDFXFA_Document.
-  pDocument->SetSDKDoc(pdfium::WrapUnique(pEnv->GetSDKDocument()));
+  pDocument->SetSDKDoc(pEnv->GetSDKDocument());
   CPDFXFA_App::GetInstance()->AddFormFillEnv(pEnv);
 #endif  // PDF_ENABLE_XFA
 
@@ -271,10 +270,12 @@ FPDFDOC_ExitFormFillEnvironment(FPDF_FORMHANDLE hHandle) {
 
 #ifdef PDF_ENABLE_XFA
   CPDFXFA_App::GetInstance()->RemoveFormFillEnv(pEnv);
-#else   // PDF_ENABLE_XFA
+
+  // Reset the focused annotations and remove the SDK document from the
+  // XFA document.
   if (CPDFSDK_Document* pSDKDoc = pEnv->GetSDKDocument()) {
-    pEnv->SetSDKDocument(nullptr);
-    delete pSDKDoc;
+    pSDKDoc->ClearAllFocusedAnnots();
+    pSDKDoc->GetXFADocument()->SetSDKDoc(nullptr);
   }
 #endif  // PDF_ENABLE_XFA
 
diff --git a/fpdfsdk/fpdfxfa/cpdfxfa_docenvironment.cpp b/fpdfsdk/fpdfxfa/cpdfxfa_docenvironment.cpp
index 1b27e16551..fed5409624 100644
--- a/fpdfsdk/fpdfxfa/cpdfxfa_docenvironment.cpp
+++ b/fpdfsdk/fpdfxfa/cpdfxfa_docenvironment.cpp
@@ -258,7 +258,11 @@ FX_BOOL CPDFXFA_DocEnvironment::PopupMenu(CXFA_FFWidget* hWidget,
 
 void CPDFXFA_DocEnvironment::PageViewEvent(CXFA_FFPageView* pPageView,
                                            uint32_t dwFlags) {
-  CPDFSDK_Environment* pEnv = m_pDocument->GetSDKDoc()->GetEnv();
+  CPDFSDK_Document* pSDKDoc = m_pDocument->GetSDKDoc();
+  if (!pSDKDoc)
+    return;
+
+  CPDFSDK_Environment* pEnv = pSDKDoc->GetEnv();
   if (!pEnv)
     return;
 
diff --git a/fpdfsdk/fpdfxfa/cpdfxfa_document.cpp b/fpdfsdk/fpdfxfa/cpdfxfa_document.cpp
index e49cfde8c4..121d060453 100644
--- a/fpdfsdk/fpdfxfa/cpdfxfa_document.cpp
+++ b/fpdfsdk/fpdfxfa/cpdfxfa_document.cpp
@@ -33,6 +33,7 @@ CPDFXFA_Document::CPDFXFA_Document(std::unique_ptr<CPDF_Document> pPDFDoc,
                                    CPDFXFA_App* pProvider)
     : m_iDocType(DOCTYPE_PDF),
       m_pPDFDoc(std::move(pPDFDoc)),
+      m_pSDKDoc(nullptr),
       m_pXFADocView(nullptr),
       m_pApp(pProvider),
       m_nLoadStatus(FXFA_LOADSTATUS_PRELOAD),
@@ -42,6 +43,11 @@ CPDFXFA_Document::CPDFXFA_Document(std::unique_ptr<CPDF_Document> pPDFDoc,
 CPDFXFA_Document::~CPDFXFA_Document() {
   m_nLoadStatus = FXFA_LOADSTATUS_CLOSING;
 
+  if (m_pSDKDoc) {
+    m_pSDKDoc->ClearAllFocusedAnnots();
+    m_pSDKDoc = nullptr;
+  }
+
   if (m_pXFADoc) {
     CXFA_FFApp* pApp = m_pApp->GetXFAApp();
     if (pApp) {
@@ -191,10 +197,6 @@ void CPDFXFA_Document::RemovePage(CPDFXFA_Page* page) {
   m_XFAPageList.SetAt(page->GetPageIndex(), nullptr);
 }
 
-void CPDFXFA_Document::SetSDKDoc(std::unique_ptr<CPDFSDK_Document> pSDKDoc) {
-  m_pSDKDoc.reset(pSDKDoc.release());
-}
-
 void CPDFXFA_Document::ClearChangeMark() {
   if (m_pSDKDoc)
     m_pSDKDoc->ClearChangeMark();
diff --git a/fpdfsdk/fpdfxfa/cpdfxfa_document.h b/fpdfsdk/fpdfxfa/cpdfxfa_document.h
index f67a880a36..fa415f917c 100644
--- a/fpdfsdk/fpdfxfa/cpdfxfa_document.h
+++ b/fpdfsdk/fpdfxfa/cpdfxfa_document.h
@@ -40,8 +40,8 @@ class CPDFXFA_Document {
   CXFA_FFDocView* GetXFADocView() { return m_pXFADocView; }
   int GetDocType() const { return m_iDocType; }
 
-  CPDFSDK_Document* GetSDKDoc() const { return m_pSDKDoc.get(); }
-  void SetSDKDoc(std::unique_ptr<CPDFSDK_Document> pSDKDoc);
+  CPDFSDK_Document* GetSDKDoc() const { return m_pSDKDoc; }
+  void SetSDKDoc(CPDFSDK_Document* pSDKDoc) { m_pSDKDoc = pSDKDoc; }
 
   void DeletePage(int page_index);
   int GetPageCount() const;
@@ -78,10 +78,8 @@ class CPDFXFA_Document {
   int m_iDocType;
 
   std::unique_ptr<CPDF_Document> m_pPDFDoc;
-  // |m_pSDKDoc| must be destroyed before |m_pPDFDoc| since it needs to access
-  // it to kill focused annotations.
-  std::unique_ptr<CPDFSDK_Document> m_pSDKDoc;
   std::unique_ptr<CXFA_FFDoc> m_pXFADoc;
+  CPDFSDK_Document* m_pSDKDoc;    // not owned.
   CXFA_FFDocView* m_pXFADocView;  // not owned.
   CPDFXFA_App* const m_pApp;
   CFX_ArrayTemplate<CPDFXFA_Page*> m_XFAPageList;