From 7cc1144192a6a7ce4c82ee310430e03ae970f6a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= <emile@getprobo.com>
Date: Fri, 16 Jan 2026 12:45:44 +0400
Subject: [PATCH] Put backn membership middleware
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Émile Ré <emile@getprobo.com>
---
 .../compliancepage/membership_middleware.go   | 74 +++++++++++++++++++
 pkg/server/api/trust/v1/resolver.go           |  1 +
 pkg/trust/service.go                          | 27 +++++++
 3 files changed, 102 insertions(+)
 create mode 100644 pkg/server/api/compliancepage/membership_middleware.go

diff --git a/pkg/server/api/compliancepage/membership_middleware.go b/pkg/server/api/compliancepage/membership_middleware.go
new file mode 100644
index 000000000..b3632d804
--- /dev/null
+++ b/pkg/server/api/compliancepage/membership_middleware.go
@@ -0,0 +1,74 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package compliancepage
+
+import (
+	"context"
+	"errors"
+	"net/http"
+
+	"github.com/99designs/gqlgen/graphql"
+	"github.com/vektah/gqlparser/v2/gqlerror"
+	"go.gearno.de/kit/httpserver"
+	"go.gearno.de/kit/log"
+	"go.probo.inc/probo/pkg/server/api/authn"
+	"go.probo.inc/probo/pkg/server/gqlutils"
+	"go.probo.inc/probo/pkg/trust"
+)
+
+func NewMembershipMiddleware(trustSvc *trust.Service, logger *log.Logger) func(next http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(
+			func(w http.ResponseWriter, r *http.Request) {
+				ctx := r.Context()
+
+				identity := authn.IdentityFromContext(r.Context())
+				if identity == nil {
+					next.ServeHTTP(w, r)
+					return
+				}
+				compliancePage := CompliancePageFromContext(ctx)
+
+				membership, err := trustSvc.GetMembershipByCompliancePageIDAndEmail(ctx, compliancePage.ID, identity.EmailAddress)
+				if err != nil {
+					if errors.Is(err, trust.ErrMembershipNotFound) {
+						next.ServeHTTP(w, r)
+						return
+					}
+
+					logger.ErrorCtx(ctx, "cannot get membership by page id and email", log.Error(err))
+					httpserver.RenderJSON(
+						w,
+						http.StatusInternalServerError,
+						&graphql.Response{
+							Errors: gqlerror.List{
+								gqlutils.Internal(ctx),
+							},
+						},
+					)
+					return
+				}
+
+				if membership.Active {
+					ctx = context.WithValue(ctx, complianceMembershipKey, membership)
+					next.ServeHTTP(w, r.WithContext(ctx))
+					return
+				}
+
+				next.ServeHTTP(w, r)
+			},
+		)
+	}
+}
diff --git a/pkg/server/api/trust/v1/resolver.go b/pkg/server/api/trust/v1/resolver.go
index 81f4acebf..f6214eec2 100644
--- a/pkg/server/api/trust/v1/resolver.go
+++ b/pkg/server/api/trust/v1/resolver.go
@@ -82,6 +82,7 @@ func NewMux(
 
 	r.Use(compliancepage.NewCompliancePagePresenceMiddleware())
 	r.Use(authn.NewSessionMiddleware(iamSvc, cookieConfig))
+	r.Use(compliancepage.NewMembershipMiddleware(trustSvc, logger))
 
 	graphqlHandler := NewGraphQLHandler(iamSvc, trustSvc, logger, baseURL, cookieConfig)
 
diff --git a/pkg/trust/service.go b/pkg/trust/service.go
index d389df7e5..e7679ec99 100644
--- a/pkg/trust/service.go
+++ b/pkg/trust/service.go
@@ -29,6 +29,7 @@ import (
 	"go.probo.inc/probo/pkg/gid"
 	"go.probo.inc/probo/pkg/html2pdf"
 	"go.probo.inc/probo/pkg/iam"
+	"go.probo.inc/probo/pkg/mail"
 	"go.probo.inc/probo/pkg/probo"
 	"go.probo.inc/probo/pkg/slack"
 )
@@ -268,3 +269,29 @@ func (s *Service) GetCustomDomainByOrganizationID(ctx context.Context, organizat
 
 	return customDomain, err
 }
+
+func (s *Service) GetMembershipByCompliancePageIDAndEmail(ctx context.Context, compliancePageID gid.GID, email mail.Addr) (*coredata.TrustCenterAccess, error) {
+	membership := &coredata.TrustCenterAccess{}
+
+	err := s.pg.WithConn(
+		ctx,
+		func(conn pg.Conn) error {
+			return membership.LoadByTrustCenterIDAndEmail(
+				ctx,
+				conn,
+				coredata.NewScopeFromObjectID(compliancePageID),
+				compliancePageID,
+				email,
+			)
+		},
+	)
+	if err != nil {
+		if errors.Is(err, coredata.ErrResourceNotFound) {
+			return nil, ErrMembershipNotFound
+		}
+
+		return nil, err
+	}
+
+	return membership, nil
+}