rohit/probo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add cookie pattern analysis worker for prefix auto-detection

Browse Source 
Background worker polls cookie_banners with pattern_analysis_requested_at
set, groups EXACT patterns sharing a common prefix, and merges groups of
3+ into a PREFIX pattern. Detection sets the flag when new EXACT patterns
are created. The worker relinks cookies, removes orphaned patterns, and
updates the draft version via ensureDraftVersionForBanner.

Signed-off-by: Émile Ré <emile@getprobo.com>
This commit is contained in:
Émile Ré
2026-04-29 14:44:58 +04:00
parent 9c402058ec
commit 7caa098941
6 changed files with 425 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
pkg/cookiebanner/service.go
View File

@@ -2674,6 +2674,10 @@ func (s *Service) ReportDetectedCookies(
if _, err := s.ensureDraftVersionForBanner(ctx, tx, scope, banner.ID); err != nil {
return fmt.Errorf("cannot ensure draft version: %w", err)
}
if err := banner.SetPatternAnalysisRequested(ctx, tx); err != nil {
return fmt.Errorf("cannot request pattern analysis: %w", err)
}
}
return nil

217
pkg/cookiebanner/worker.go Normal file
View File

@@ -0,0 +1,217 @@
// Copyright (c) 2026 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 cookiebanner
import (
"context"
"errors"
"fmt"
"sort"
"strings"
"time"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
)
const patternMergeThreshold = 3
type patternAnalysisHandler struct {
svc *Service
pg *pg.Client
logger *log.Logger
}
func NewPatternAnalysisWorker(
svc *Service,
pgClient *pg.Client,
logger *log.Logger,
opts ...worker.Option,
) *worker.Worker[coredata.CookieBannerPatternAnalysisTask] {
h := &patternAnalysisHandler{
svc: svc,
pg: pgClient,
logger: logger,
}
return worker.New(
"cookie-pattern-analysis-worker",
h,
logger,
opts...,
)
}
func (h *patternAnalysisHandler) Claim(ctx context.Context) (coredata.CookieBannerPatternAnalysisTask, error) {
var task coredata.CookieBannerPatternAnalysisTask
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := task.ClaimNextForUpdateSkipLocked(ctx, tx); err != nil {
return err
}
return task.ClearPatternAnalysisFlag(ctx, tx)
},
); err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
return coredata.CookieBannerPatternAnalysisTask{}, worker.ErrNoTask
}
return coredata.CookieBannerPatternAnalysisTask{}, fmt.Errorf("cannot claim pattern analysis task: %w", err)
}
return task, nil
}
func (h *patternAnalysisHandler) Process(ctx context.Context, task coredata.CookieBannerPatternAnalysisTask) error {
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
scope := coredata.NewScope(task.TenantID)
var patterns coredata.CookiePatterns
if err := patterns.LoadAllByCookieBannerID(ctx, tx, scope, task.BannerID); err != nil {
return fmt.Errorf("cannot load patterns: %w", err)
}
prefixGroups := groupByPrefix(patterns)
merged := false
for prefix, group := range prefixGroups {
if len(group) < patternMergeThreshold {
continue
}
duration := mostCommonDuration(group)
source := bestSource(group)
prefixPattern := &coredata.CookiePattern{
ID: gid.New(task.TenantID, coredata.CookiePatternEntityType),
OrganizationID: group[0].OrganizationID,
CookieBannerID: task.BannerID,
CookieCategoryID: group[0].CookieCategoryID,
Pattern: prefix,
MatchType: coredata.CookiePatternMatchTypePrefix,
DisplayName: prefix + "*",
Duration: duration,
Description: "",
Source: source,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}
inserted, err := prefixPattern.InsertIfNotExists(ctx, tx, scope)
if err != nil {
return fmt.Errorf("cannot insert prefix pattern %q: %w", prefix, err)
}
if !inserted {
continue
}
for _, exactPattern := range group {
var cookies coredata.Cookies
if err := cookies.RelinkByCookiePatternID(ctx, tx, scope, exactPattern.ID, prefixPattern.ID); err != nil {
return fmt.Errorf("cannot relink cookies from pattern %q: %w", exactPattern.Pattern, err)
}
if err := exactPattern.Delete(ctx, tx, scope); err != nil {
return fmt.Errorf("cannot delete orphaned exact pattern %q: %w", exactPattern.Pattern, err)
}
}
merged = true
h.logger.InfoCtx(
ctx,
"merged exact patterns into prefix pattern",
log.String("prefix", prefix),
log.Int("count", len(group)),
log.String("banner_id", task.BannerID.String()),
)
}
if merged {
if _, err := h.svc.ensureDraftVersionForBanner(ctx, tx, scope, task.BannerID); err != nil {
return fmt.Errorf("cannot ensure draft version: %w", err)
}
}
return nil
},
)
}
func groupByPrefix(patterns coredata.CookiePatterns) map[string][]*coredata.CookiePattern {
groups := make(map[string][]*coredata.CookiePattern)
for _, p := range patterns {
if p.MatchType != coredata.CookiePatternMatchTypeExact {
continue
}
prefix := extractPrefix(p.Pattern)
if prefix == "" {
continue
}
groups[prefix] = append(groups[prefix], p)
}
return groups
}
func extractPrefix(name string) string {
for _, sep := range []byte{'_', '-'} {
idx := strings.IndexByte(name, sep)
if idx > 0 {
return name[:idx+1]
}
}
return ""
}
func bestSource(patterns []*coredata.CookiePattern) coredata.CookieSource {
for _, p := range patterns {
if p.Source == coredata.CookieSourceScript {
return coredata.CookieSourceScript
}
}
return coredata.CookieSourcePreExisting
}
func mostCommonDuration(patterns []*coredata.CookiePattern) string {
counts := make(map[string]int)
for _, p := range patterns {
counts[p.Duration]++
}
type entry struct {
duration string
count int
}
entries := make([]entry, 0, len(counts))
for d, c := range counts {
entries = append(entries, entry{d, c})
}
sort.Slice(entries, func(i, j int) bool {
return entries[i].count > entries[j].count
})
return entries[0].duration
}

142
pkg/coredata/cookie_banner.go
View File

@@ -30,19 +30,25 @@ import (
type (
CookieBanner struct {
ID gid.GID `db:"id"`
OrganizationID gid.GID `db:"organization_id"`
Name string `db:"name"`
Origin string `db:"origin"`
State CookieBannerState `db:"state"`
PrivacyPolicyURL *string `db:"privacy_policy_url"`
CookiePolicyURL string `db:"cookie_policy_url"`
ConsentExpiryDays int `db:"consent_expiry_days"`
ConsentMode CookieConsentMode `db:"consent_mode"`
ShowBranding bool `db:"show_branding"`
DefaultLanguage string `db:"default_language"`
CreatedAt time.Time `db:"created_at"`
UpdatedAt time.Time `db:"updated_at"`
ID gid.GID `db:"id"`
OrganizationID gid.GID `db:"organization_id"`
Name string `db:"name"`
Origin string `db:"origin"`
State CookieBannerState `db:"state"`
PrivacyPolicyURL *string `db:"privacy_policy_url"`
CookiePolicyURL string `db:"cookie_policy_url"`
ConsentExpiryDays int `db:"consent_expiry_days"`
ConsentMode CookieConsentMode `db:"consent_mode"`
ShowBranding bool `db:"show_branding"`
DefaultLanguage string `db:"default_language"`
PatternAnalysisRequestedAt *time.Time `db:"pattern_analysis_requested_at"`
CreatedAt time.Time `db:"created_at"`
UpdatedAt time.Time `db:"updated_at"`
}
CookieBannerPatternAnalysisTask struct {
BannerID gid.GID
TenantID gid.TenantID
}
CookieBanners []*CookieBanner
@@ -91,6 +97,7 @@ SELECT
consent_mode,
show_branding,
default_language,
pattern_analysis_requested_at,
created_at,
updated_at
FROM
@@ -143,6 +150,7 @@ SELECT
consent_mode,
show_branding,
default_language,
pattern_analysis_requested_at,
created_at,
updated_at
FROM
@@ -196,6 +204,7 @@ SELECT
consent_mode,
show_branding,
default_language,
pattern_analysis_requested_at,
created_at,
updated_at
FROM
@@ -253,6 +262,7 @@ SELECT
consent_mode,
show_branding,
default_language,
pattern_analysis_requested_at,
created_at,
updated_at
FROM
@@ -303,6 +313,7 @@ SELECT
consent_mode,
show_branding,
default_language,
pattern_analysis_requested_at,
created_at,
updated_at
FROM
@@ -389,6 +400,7 @@ INSERT INTO cookie_banners (
consent_mode,
show_branding,
default_language,
pattern_analysis_requested_at,
created_at,
updated_at
) VALUES (
@@ -404,26 +416,28 @@ INSERT INTO cookie_banners (
@consent_mode,
@show_branding,
@default_language,
@pattern_analysis_requested_at,
@created_at,
@updated_at
)
`
args := pgx.StrictNamedArgs{
"id": b.ID,
"tenant_id": scope.GetTenantID(),
"organization_id": b.OrganizationID,
"name": b.Name,
"origin": b.Origin,
"state": b.State,
"privacy_policy_url": b.PrivacyPolicyURL,
"cookie_policy_url": b.CookiePolicyURL,
"consent_expiry_days": b.ConsentExpiryDays,
"consent_mode": b.ConsentMode,
"show_branding": b.ShowBranding,
"default_language": b.DefaultLanguage,
"created_at": b.CreatedAt,
"updated_at": b.UpdatedAt,
"id": b.ID,
"tenant_id": scope.GetTenantID(),
"organization_id": b.OrganizationID,
"name": b.Name,
"origin": b.Origin,
"state": b.State,
"privacy_policy_url": b.PrivacyPolicyURL,
"cookie_policy_url": b.CookiePolicyURL,
"consent_expiry_days": b.ConsentExpiryDays,
"consent_mode": b.ConsentMode,
"show_branding": b.ShowBranding,
"default_language": b.DefaultLanguage,
"pattern_analysis_requested_at": b.PatternAnalysisRequestedAt,
"created_at": b.CreatedAt,
"updated_at": b.UpdatedAt,
}
_, err := tx.Exec(ctx, q, args)
@@ -557,3 +571,77 @@ WHERE
return nil
}
func (t *CookieBannerPatternAnalysisTask) ClaimNextForUpdateSkipLocked(
ctx context.Context,
tx pg.Tx,
) error {
q := `
SELECT
id,
tenant_id
FROM
cookie_banners
WHERE
pattern_analysis_requested_at IS NOT NULL
ORDER BY
pattern_analysis_requested_at ASC
FOR UPDATE SKIP LOCKED
LIMIT 1;
`
var tenantIDStr string
if err := tx.QueryRow(ctx, q).Scan(&t.BannerID, &tenantIDStr); err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return ErrResourceNotFound
}
return fmt.Errorf("cannot claim banner for pattern analysis: %w", err)
}
if err := t.TenantID.UnmarshalText([]byte(tenantIDStr)); err != nil {
return fmt.Errorf("cannot parse tenant ID: %w", err)
}
return nil
}
func (t *CookieBannerPatternAnalysisTask) ClearPatternAnalysisFlag(
ctx context.Context,
tx pg.Tx,
) error {
q := `
UPDATE cookie_banners
SET pattern_analysis_requested_at = NULL
WHERE id = @id
`
args := pgx.StrictNamedArgs{"id": t.BannerID}
_, err := tx.Exec(ctx, q, args)
if err != nil {
return fmt.Errorf("cannot clear pattern analysis flag: %w", err)
}
return nil
}
func (b *CookieBanner) SetPatternAnalysisRequested(
ctx context.Context,
tx pg.Tx,
) error {
q := `
UPDATE cookie_banners
SET pattern_analysis_requested_at = NOW()
WHERE id = @id
AND pattern_analysis_requested_at IS NULL
`
args := pgx.StrictNamedArgs{"id": b.ID}
_, err := tx.Exec(ctx, q, args)
if err != nil {
return fmt.Errorf("cannot set pattern analysis requested: %w", err)
}
return nil
}

62
pkg/coredata/cookie_pattern.go
View File

@@ -381,6 +381,68 @@ INSERT INTO cookie_patterns (
return nil
}
func (cp *CookiePattern) InsertIfNotExists(
ctx context.Context,
tx pg.Tx,
scope Scoper,
) (bool, error) {
q := `
INSERT INTO cookie_patterns (
id,
tenant_id,
organization_id,
cookie_banner_id,
cookie_category_id,
pattern,
match_type,
display_name,
duration,
description,
source,
created_at,
updated_at
) VALUES (
@id,
@tenant_id,
@organization_id,
@cookie_banner_id,
@cookie_category_id,
@pattern,
@match_type,
@display_name,
@duration,
@description,
@source,
@created_at,
@updated_at
)
ON CONFLICT (cookie_banner_id, pattern) DO NOTHING
`
args := pgx.StrictNamedArgs{
"id": cp.ID,
"tenant_id": scope.GetTenantID(),
"organization_id": cp.OrganizationID,
"cookie_banner_id": cp.CookieBannerID,
"cookie_category_id": cp.CookieCategoryID,
"pattern": cp.Pattern,
"match_type": cp.MatchType,
"display_name": cp.DisplayName,
"duration": cp.Duration,
"description": cp.Description,
"source": cp.Source,
"created_at": cp.CreatedAt,
"updated_at": cp.UpdatedAt,
}
result, err := tx.Exec(ctx, q, args)
if err != nil {
return false, fmt.Errorf("cannot insert cookie pattern: %w", err)
}
return result.RowsAffected() > 0, nil
}
func (cp *CookiePattern) Update(
ctx context.Context,
tx pg.Tx,

16
pkg/coredata/migrations/20260429T102803Z.sql Normal file
View File

@@ -0,0 +1,16 @@
-- Copyright (c) 2026 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.
ALTER TABLE cookie_banners
ADD COLUMN pattern_analysis_requested_at TIMESTAMP WITH TIME ZONE;

11
pkg/probod/probod.go
View File

@@ -656,6 +656,16 @@ func (impl *Implm) Run(
},
)
cookiePatternAnalysisWorker := cookiebanner.NewPatternAnalysisWorker(cookieBannerService, pgClient, l.Named("cookie-pattern-analysis-worker"))
cookiePatternAnalysisWorkerCtx, stopCookiePatternAnalysisWorker := context.WithCancel(context.Background())
wg.Go(
func() {
if err := cookiePatternAnalysisWorker.Run(cookiePatternAnalysisWorkerCtx); err != nil {
cancel(fmt.Errorf("cookie pattern analysis worker crashed: %w", err))
}
},
)
mailingListWorker := mailman.NewMailingListWorker(mailmanService, pgClient, l.Named("mailing-list-worker"))
mailingListWorkerCtx, stopMailingListWorker := context.WithCancel(context.Background())
wg.Go(
@@ -720,6 +730,7 @@ func (impl *Implm) Run(
stopTrustCenterServer()
stopWebhookSender()
stopESignService()
stopCookiePatternAnalysisWorker()
stopMailingListWorker()
stopEvidenceDescriptionWorker()
stopDocumentPDFWorker()