// Copyright 2024 The PDFium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "fpdfsdk/fpdfxfa/cpdfxfa_page.h"

#include "core/fxcrt/observed_ptr.h"
#include "fpdfsdk/cpdfsdk_annot.h"
#include "fpdfsdk/cpdfsdk_formfillenvironment.h"
#include "fpdfsdk/cpdfsdk_helpers.h"
#include "fpdfsdk/cpdfsdk_pageview.h"
#include "fpdfsdk/fpdfxfa/cpdfxfa_context.h"
#include "fpdfsdk/fpdfxfa/cpdfxfa_widget.h"
#include "fxjs/xfa/cjx_object.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "testing/xfa_js_embedder_test.h"
#include "xfa/fxfa/cxfa_ffdocview.h"
#include "xfa/fxfa/cxfa_fffield.h"
#include "xfa/fxfa/cxfa_ffwidget.h"
#include "xfa/fxfa/parser/cxfa_node.h"

namespace {

CPDFXFA_Widget* FindXFAWidgetByName(CPDFSDK_PageView* page_view,
                                    const WideString& name) {
  for (CPDFSDK_Annot* annot : page_view->GetAnnotList()) {
    CPDFXFA_Widget* xfa = ToXFAWidget(annot);
    if (!xfa) {
      continue;
    }
    CXFA_Node* node = xfa->GetXFAFFWidget()->GetNode();
    if (node && node->JSObject()->GetCData(XFA_Attribute::Name) == name) {
      return xfa;
    }
  }
  return nullptr;
}

}  // namespace

class CPDFXFAPageEmbedderTest : public XFAJSEmbedderTest {};

// CPDFXFA_Page::DrawFocusAnnot() holds a raw CPDFSDK_Annot* across a
// CXFA_FFPageWidgetIterator loop. The iterator's MoveToNext() may lazily call
// CXFA_FFTextEdit::LoadWidget(), whose initial SetText() dispatches an XFA
// `change` event. That handler can run arbitrary script, including
// removeInstance() + xfa.host.openList(), which synchronously triggers a
// relayout that calls DeleteAnnotForFFWidget() on the focused annot's
// CPDFXFA_Widget — freeing it while DrawFocusAnnot() still holds the raw
// pointer. The trailing ToXFAWidget(pAnnot) is then a virtual call on freed
// memory.
//
// The lazy-load precondition (a widget that has XFA_WidgetStatus::kVisible set
// but for which IsLoaded() returns false at draw time) is reachable but
// requires a specific sequence of layout transitions; this test simulates it
// directly via ClearNormalWidgetForTesting() so the dangerous re-entry path is
// exercised deterministically. The unguarded raw pointer in DrawFocusAnnot()
// is the bug under test (compare the ObservedPtr<> guard already applied in
// the sibling iterator helper in the same file).
TEST_F(CPDFXFAPageEmbedderTest, DrawFocusAnnotUseAfterFree) {
  ASSERT_TRUE(OpenDocument("xfa_drawfocus_uaf.pdf"));
  ScopedPage page = LoadScopedPage(0);
  ASSERT_TRUE(page);

  CPDFSDK_FormFillEnvironment* env =
      CPDFSDKFormFillEnvironmentFromFPDFFormHandle(form_handle());
  ASSERT_TRUE(env);
  CPDFSDK_PageView* page_view =
      env->GetPageView(IPDFPageFromFPDFPage(page.get()));
  ASSERT_TRUE(page_view);

  // Focus the victim widget V — DrawFocusAnnot() will be passed its
  // CPDFXFA_Widget* as a raw pointer.
  CPDFXFA_Widget* v_annot = FindXFAWidgetByName(page_view, L"V");
  ASSERT_TRUE(v_annot);
  CXFA_FFWidget* v_widget = v_annot->GetXFAFFWidget();
  {
    ObservedPtr<CPDFSDK_Annot> obs(v_annot);
    ASSERT_TRUE(env->SetFocusAnnot(obs));
  }
  ASSERT_EQ(v_annot, page_view->GetFocusAnnot());

  // Force the trigger widget T back to "not loaded" so the iterator inside
  // DrawFocusAnnot() re-runs its LoadWidget() and fires T's `change` script.
  CPDFXFA_Widget* t_annot = FindXFAWidgetByName(page_view, L"T");
  ASSERT_TRUE(t_annot);
  CXFA_FFField* t_field = ToField(t_annot->GetXFAFFWidget());
  ASSERT_TRUE(t_field);
  t_field->ClearNormalWidgetForTesting();
  ASSERT_FALSE(t_field->IsLoaded());

  // Arm the script (it is gated on this so it does not fire during initial
  // page load / LoadFXAnnots()).
  EXPECT_TRUE(Execute("xfa.form.form1.flag = \"armed\""));

  // Reaches CPDFXFA_Page::DrawFocusAnnot() via FPDF_FFLDraw(). On the
  // unpatched tree this is a heap-use-after-free in ToXFAWidget(pAnnot).
  ScopedFPDFBitmap bitmap = RenderLoadedPage(page.get());

  // Confirm the script actually ran and the victim's annot is gone, i.e. the
  // free side of the UAF was reached during the draw above.
  EXPECT_TRUE(Execute("xfa.form.form1.flag"));
  EXPECT_FALSE(page_view->GetAnnotForFFWidget(v_widget));
}