Add tracker mapping worker and initiator domain extraction

Poll-based worker that maps org-scoped tracker patterns to the common knowledge base via pattern matching and domain-based attribution. Populates initiator_domain on detected trackers at report time. Resolves org-scoped vendors through the common third party link. Signed-off-by: Émile Ré <emile@getprobo.com> Signed-off-by: Émile Ré <emile@probo.com>
2026-05-13 17:26:15 +04:00
parent f0483d5280
commit d6f268cb06
7 changed files with 498 additions and 15 deletions
--- a/pkg/cookiebanner/service.go
+++ b/pkg/cookiebanner/service.go
@@ -2122,20 +2122,21 @@ func (s *Service) reportDetectedTracker(
 		*matchedPatternIDs = append(*matchedPatternIDs, matchedPattern.ID)
 	} else {
 		newPattern := &coredata.TrackerPattern{
-			ID:               gid.New(scope.GetTenantID(), coredata.TrackerPatternEntityType),
-			OrganizationID:   banner.OrganizationID,
-			CookieBannerID:   banner.ID,
-			CookieCategoryID: uncategorisedID,
-			TrackerType:      info.TrackerType,
-			Pattern:          info.Identifier,
-			MatchType:        coredata.TrackerPatternMatchTypeExact,
-			DisplayName:      info.Identifier,
-			Description:      "",
-			MaxAgeSeconds:    info.MaxAgeSeconds,
-			Source:           info.Source,
-			LastMatchedAt:    &now,
-			CreatedAt:        now,
-			UpdatedAt:        now,
+			ID:                 gid.New(scope.GetTenantID(), coredata.TrackerPatternEntityType),
+			OrganizationID:     banner.OrganizationID,
+			CookieBannerID:     banner.ID,
+			CookieCategoryID:   uncategorisedID,
+			TrackerType:        info.TrackerType,
+			Pattern:            info.Identifier,
+			MatchType:          coredata.TrackerPatternMatchTypeExact,
+			DisplayName:        info.Identifier,
+			Description:        "",
+			MaxAgeSeconds:      info.MaxAgeSeconds,
+			Source:             info.Source,
+			LastMatchedAt:      &now,
+			MappingRequestedAt: &now,
+			CreatedAt:          now,
+			UpdatedAt:          now,
 		}
 		wasInserted, err := newPattern.InsertIfNotExists(ctx, tx, scope)
 		if err != nil {
@@ -2153,6 +2154,13 @@ func (s *Service) reportDetectedTracker(
 		}
 	}

+	var initiatorDomain *string
+	if info.InitiatorURL != nil {
+		if domain := uri.ExtractDomain(*info.InitiatorURL); domain != "" {
+			initiatorDomain = &domain
+		}
+	}
+
 	tracker := &coredata.DetectedTracker{
 		ID:               gid.New(scope.GetTenantID(), coredata.DetectedTrackerEntityType),
 		CookieBannerID:   banner.ID,
@@ -2163,6 +2171,7 @@ func (s *Service) reportDetectedTracker(
 		Source:           info.Source,
 		ValueSize:        info.ValueSize,
 		InitiatorURL:     info.InitiatorURL,
+		InitiatorDomain:  initiatorDomain,
 		LastDetectedAt:   now,
 		CreatedAt:        now,
 		UpdatedAt:        now,
--- a/pkg/coredata/detected_tracker.go
+++ b/pkg/coredata/detected_tracker.go
@@ -16,6 +16,7 @@ package coredata

 import (
 	"context"
+	"errors"
 	"fmt"
 	"maps"
 	"time"
@@ -151,6 +152,33 @@ WHERE
 	return count, nil
 }

+func (dts *DetectedTrackers) LoadCommonThirdPartyIDByTrackerPatternID(
+	ctx context.Context,
+	conn pg.Querier,
+	trackerPatternID gid.GID,
+) (*gid.GID, error) {
+	q := `
+SELECT DISTINCT ctpd.common_third_party_id
+FROM detected_trackers dt
+JOIN common_third_party_domains ctpd ON ctpd.domain = dt.initiator_domain
+WHERE dt.tracker_pattern_id = @tracker_pattern_id
+  AND dt.initiator_domain IS NOT NULL
+LIMIT 1;
+`
+
+	args := pgx.StrictNamedArgs{"tracker_pattern_id": trackerPatternID}
+
+	var commonThirdPartyID gid.GID
+	if err := conn.QueryRow(ctx, q, args).Scan(&commonThirdPartyID); err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return nil, nil
+		}
+		return nil, fmt.Errorf("cannot load common third party ID by tracker pattern: %w", err)
+	}
+
+	return &commonThirdPartyID, nil
+}
+
 func (dts *DetectedTrackers) RelinkByTrackerPatternID(
 	ctx context.Context,
 	tx pg.Tx,
--- a/pkg/coredata/migrations/20260513T130929Z.sql
+++ b/pkg/coredata/migrations/20260513T130929Z.sql
@@ -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 tracker_patterns
+    ADD COLUMN mapping_requested_at TIMESTAMP WITH TIME ZONE;
--- a/pkg/coredata/third_party.go
+++ b/pkg/coredata/third_party.go
@@ -1315,3 +1315,75 @@ ORDER BY name ASC

 	return nil
 }
+
+func (v *Vendor) LoadByOrganizationIDAndCommonThirdPartyID(
+	ctx context.Context,
+	conn pg.Querier,
+	scope Scoper,
+	organizationID gid.GID,
+	commonThirdPartyID gid.GID,
+) error {
+	q := `
+SELECT
+	id,
+	tenant_id,
+	organization_id,
+	common_third_party_id,
+	name,
+	description,
+	category,
+	headquarter_address,
+	legal_name,
+	website_url,
+	privacy_policy_url,
+	service_level_agreement_url,
+	data_processing_agreement_url,
+	business_associate_agreement_url,
+	subprocessors_list_url,
+	certifications,
+	countries,
+	business_owner_profile_id,
+	security_owner_profile_id,
+	status_page_url,
+	terms_of_service_url,
+	security_page_url,
+	trust_page_url,
+	show_on_trust_center,
+	created_at,
+	updated_at
+FROM
+	vendors
+WHERE
+	%s
+	AND organization_id = @organization_id
+	AND common_third_party_id = @common_third_party_id
+	AND snapshot_id IS NULL
+LIMIT 1;
+`
+
+	q = fmt.Sprintf(q, scope.SQLFragment())
+
+	args := pgx.StrictNamedArgs{
+		"organization_id":       organizationID,
+		"common_third_party_id": commonThirdPartyID,
+	}
+	maps.Copy(args, scope.SQLArguments())
+
+	rows, err := conn.Query(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot query vendor by common third party: %w", err)
+	}
+	defer rows.Close()
+
+	vendor, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Vendor])
+	if err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return ErrResourceNotFound
+		}
+		return fmt.Errorf("cannot collect vendor by common third party: %w", err)
+	}
+
+	*v = vendor
+
+	return nil
+}
--- a/pkg/coredata/tracker_pattern.go
+++ b/pkg/coredata/tracker_pattern.go
@@ -46,6 +46,7 @@ type (
 		MaxAgeSeconds          *int                    `db:"max_age_seconds"`
 		Source                 *CookieSource           `db:"source"`
 		LastMatchedAt          *time.Time              `db:"last_matched_at"`
+		MappingRequestedAt     *time.Time              `db:"mapping_requested_at"`
 		CreatedAt              time.Time               `db:"created_at"`
 		UpdatedAt              time.Time               `db:"updated_at"`
 	}
@@ -114,6 +115,7 @@ SELECT
 	max_age_seconds,
 	source,
 	last_matched_at,
+	mapping_requested_at,
 	created_at,
 	updated_at
 FROM
@@ -173,6 +175,7 @@ SELECT
 	max_age_seconds,
 	source,
 	last_matched_at,
+	mapping_requested_at,
 	created_at,
 	updated_at
 FROM
@@ -239,6 +242,7 @@ SELECT
 	max_age_seconds,
 	source,
 	last_matched_at,
+	mapping_requested_at,
 	created_at,
 	updated_at
 FROM
@@ -315,6 +319,7 @@ INSERT INTO tracker_patterns (
 	max_age_seconds,
 	source,
 	last_matched_at,
+	mapping_requested_at,
 	created_at,
 	updated_at
 ) VALUES (
@@ -334,6 +339,7 @@ INSERT INTO tracker_patterns (
 	@max_age_seconds,
 	@source,
 	@last_matched_at,
+	@mapping_requested_at,
 	@created_at,
 	@updated_at
 )
@@ -356,6 +362,7 @@ INSERT INTO tracker_patterns (
 		"max_age_seconds":           tp.MaxAgeSeconds,
 		"source":                    tp.Source,
 		"last_matched_at":           tp.LastMatchedAt,
+		"mapping_requested_at":      tp.MappingRequestedAt,
 		"created_at":                tp.CreatedAt,
 		"updated_at":                tp.UpdatedAt,
 	}
@@ -396,6 +403,7 @@ INSERT INTO tracker_patterns (
 	max_age_seconds,
 	source,
 	last_matched_at,
+	mapping_requested_at,
 	created_at,
 	updated_at
 ) VALUES (
@@ -415,6 +423,7 @@ INSERT INTO tracker_patterns (
 	@max_age_seconds,
 	@source,
 	@last_matched_at,
+	@mapping_requested_at,
 	@created_at,
 	@updated_at
 )
@@ -438,6 +447,7 @@ ON CONFLICT (cookie_banner_id, tracker_type, pattern, COALESCE(max_age_seconds,
 		"max_age_seconds":           tp.MaxAgeSeconds,
 		"source":                    tp.Source,
 		"last_matched_at":           tp.LastMatchedAt,
+		"mapping_requested_at":      tp.MappingRequestedAt,
 		"created_at":                tp.CreatedAt,
 		"updated_at":                tp.UpdatedAt,
 	}
@@ -556,6 +566,7 @@ SELECT
 	max_age_seconds,
 	source,
 	last_matched_at,
+	mapping_requested_at,
 	created_at,
 	updated_at
 FROM
@@ -654,6 +665,7 @@ SELECT
 	max_age_seconds,
 	source,
 	last_matched_at,
+	mapping_requested_at,
 	created_at,
 	updated_at
 FROM
@@ -765,6 +777,7 @@ SELECT
 	max_age_seconds,
 	source,
 	last_matched_at,
+	mapping_requested_at,
 	created_at,
 	updated_at
 FROM
@@ -897,3 +910,135 @@ WHERE

 	return nil
 }
+
+func (tp *TrackerPattern) LoadNextForMappingForUpdateSkipLocked(
+	ctx context.Context,
+	tx pg.Tx,
+) error {
+	q := `
+SELECT
+	id,
+	organization_id,
+	cookie_banner_id,
+	cookie_category_id,
+	common_tracker_pattern_id,
+	third_party_id,
+	tracker_type,
+	pattern,
+	match_type,
+	display_name,
+	description,
+	excluded,
+	max_age_seconds,
+	source,
+	last_matched_at,
+	mapping_requested_at,
+	created_at,
+	updated_at
+FROM
+	tracker_patterns
+WHERE
+	mapping_requested_at IS NOT NULL
+ORDER BY
+	mapping_requested_at ASC
+FOR UPDATE SKIP LOCKED
+LIMIT 1;
+`
+
+	rows, err := tx.Query(ctx, q)
+	if err != nil {
+		return fmt.Errorf("cannot query tracker patterns for mapping: %w", err)
+	}
+
+	pattern, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[TrackerPattern])
+	if err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return ErrResourceNotFound
+		}
+		return fmt.Errorf("cannot collect tracker pattern for mapping: %w", err)
+	}
+
+	*tp = pattern
+
+	return nil
+}
+
+func (tp *TrackerPattern) ClearMappingRequestedAt(
+	ctx context.Context,
+	tx pg.Tx,
+) error {
+	q := `
+UPDATE tracker_patterns
+SET mapping_requested_at = NULL
+WHERE id = @id
+`
+
+	args := pgx.StrictNamedArgs{"id": tp.ID}
+
+	_, err := tx.Exec(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot clear mapping requested at: %w", err)
+	}
+
+	tp.MappingRequestedAt = nil
+
+	return nil
+}
+
+func (tp *TrackerPattern) SetMappingRequested(
+	ctx context.Context,
+	tx pg.Tx,
+) error {
+	q := `
+UPDATE tracker_patterns
+SET mapping_requested_at = NOW()
+WHERE id = @id
+  AND mapping_requested_at IS NULL
+`
+
+	args := pgx.StrictNamedArgs{"id": tp.ID}
+
+	_, err := tx.Exec(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot set mapping requested: %w", err)
+	}
+
+	return nil
+}
+
+func (tp *TrackerPattern) UpdateMapping(
+	ctx context.Context,
+	tx pg.Tx,
+	commonTrackerPatternID *gid.GID,
+	thirdPartyID *gid.GID,
+) error {
+	q := `
+UPDATE tracker_patterns
+SET
+	common_tracker_pattern_id = @common_tracker_pattern_id,
+	third_party_id = @third_party_id,
+	mapping_requested_at = NULL,
+	updated_at = @updated_at
+WHERE id = @id
+`
+
+	now := time.Now()
+	args := pgx.StrictNamedArgs{
+		"id":                        tp.ID,
+		"common_tracker_pattern_id": commonTrackerPatternID,
+		"third_party_id":            thirdPartyID,
+		"updated_at":                now,
+	}
+
+	_, err := tx.Exec(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot update tracker pattern mapping: %w", err)
+	}
+
+	tp.CommonTrackerPatternID = commonTrackerPatternID
+	tp.ThirdPartyID = thirdPartyID
+	tp.MappingRequestedAt = nil
+	tp.UpdatedAt = now
+
+	return nil
+}
--- a/pkg/probo/tracker_mapping_worker.go
+++ b/pkg/probo/tracker_mapping_worker.go
@@ -0,0 +1,202 @@
+// 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 probo
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"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"
+)
+
+type trackerMappingHandler struct {
+	pg     *pg.Client
+	logger *log.Logger
+}
+
+func NewTrackerMappingWorker(
+	pgClient *pg.Client,
+	logger *log.Logger,
+	opts ...worker.Option,
+) *worker.Worker[coredata.TrackerPattern] {
+	h := &trackerMappingHandler{
+		pg:     pgClient,
+		logger: logger,
+	}
+
+	return worker.New(
+		"tracker-mapping-worker",
+		h,
+		logger,
+		opts...,
+	)
+}
+
+func (h *trackerMappingHandler) Claim(ctx context.Context) (coredata.TrackerPattern, error) {
+	var tp coredata.TrackerPattern
+
+	if err := h.pg.WithTx(
+		ctx,
+		func(ctx context.Context, tx pg.Tx) error {
+			if err := tp.LoadNextForMappingForUpdateSkipLocked(ctx, tx); err != nil {
+				return err
+			}
+
+			return tp.ClearMappingRequestedAt(ctx, tx)
+		},
+	); err != nil {
+		if errors.Is(err, coredata.ErrResourceNotFound) {
+			return coredata.TrackerPattern{}, worker.ErrNoTask
+		}
+		return coredata.TrackerPattern{}, fmt.Errorf("cannot claim tracker mapping task: %w", err)
+	}
+
+	return tp, nil
+}
+
+func (h *trackerMappingHandler) Process(ctx context.Context, tp coredata.TrackerPattern) error {
+	return h.pg.WithTx(
+		ctx,
+		func(ctx context.Context, tx pg.Tx) error {
+			var commonPatternID *gid.GID
+			var thirdPartyID *gid.GID
+
+			commonPatternID, thirdPartyID = h.matchByPattern(ctx, tx, tp)
+
+			if commonPatternID == nil {
+				commonPatternID, thirdPartyID = h.matchByDomain(ctx, tx, tp)
+			}
+
+			if commonPatternID != nil && thirdPartyID == nil {
+				var cp coredata.CommonTrackerPattern
+				if err := cp.LoadByID(ctx, tx, *commonPatternID); err == nil {
+					thirdPartyID = h.resolveThirdParty(ctx, tx, tp, &cp)
+				}
+			}
+
+			if commonPatternID != nil || thirdPartyID != nil {
+				if err := tp.UpdateMapping(ctx, tx, commonPatternID, thirdPartyID); err != nil {
+					return fmt.Errorf("cannot update tracker pattern mapping: %w", err)
+				}
+
+				h.logger.InfoCtx(
+					ctx,
+					"mapped tracker pattern",
+					log.String("pattern", tp.Pattern),
+					log.String("tracker_pattern_id", tp.ID.String()),
+				)
+			}
+
+			return nil
+		},
+	)
+}
+
+func (h *trackerMappingHandler) matchByPattern(
+	ctx context.Context,
+	conn pg.Querier,
+	tp coredata.TrackerPattern,
+) (*gid.GID, *gid.GID) {
+	var commonPattern coredata.CommonTrackerPattern
+	if err := commonPattern.LoadByPattern(ctx, conn, tp.TrackerType, tp.Pattern, tp.MaxAgeSeconds); err != nil {
+		if !errors.Is(err, coredata.ErrResourceNotFound) {
+			h.logger.ErrorCtx(ctx, "cannot load common tracker pattern", log.Error(err))
+		}
+		return nil, nil
+	}
+
+	var thirdPartyID *gid.GID
+	if commonPattern.CommonThirdPartyID != nil {
+		thirdPartyID = h.resolveThirdParty(ctx, conn, tp, &commonPattern)
+	}
+
+	return &commonPattern.ID, thirdPartyID
+}
+
+func (h *trackerMappingHandler) matchByDomain(
+	ctx context.Context,
+	tx pg.Tx,
+	tp coredata.TrackerPattern,
+) (*gid.GID, *gid.GID) {
+	var trackers coredata.DetectedTrackers
+	commonThirdPartyID, err := trackers.LoadCommonThirdPartyIDByTrackerPatternID(ctx, tx, tp.ID)
+	if err != nil {
+		h.logger.ErrorCtx(ctx, "cannot load common third party ID from domain", log.Error(err))
+		return nil, nil
+	}
+
+	if commonThirdPartyID == nil {
+		return nil, nil
+	}
+
+	now := time.Now()
+	commonPattern := coredata.CommonTrackerPattern{
+		ID:                 gid.New(gid.NilTenant, coredata.CommonTrackerPatternEntityType),
+		CommonThirdPartyID: commonThirdPartyID,
+		TrackerType:        tp.TrackerType,
+		Pattern:            tp.Pattern,
+		MatchType:          tp.MatchType,
+		Description:        tp.Description,
+		MaxAgeSeconds:      tp.MaxAgeSeconds,
+		Confidence:         0.7,
+		CreatedAt:          now,
+		UpdatedAt:          now,
+	}
+
+	if _, err := commonPattern.Upsert(ctx, tx); err != nil {
+		h.logger.ErrorCtx(
+			ctx,
+			"cannot upsert common tracker pattern from domain match",
+			log.Error(err),
+		)
+		return nil, nil
+	}
+
+	thirdPartyID := h.resolveThirdParty(ctx, tx, tp, &commonPattern)
+
+	return &commonPattern.ID, thirdPartyID
+}
+
+func (h *trackerMappingHandler) resolveThirdParty(
+	ctx context.Context,
+	conn pg.Querier,
+	tp coredata.TrackerPattern,
+	commonPattern *coredata.CommonTrackerPattern,
+) *gid.GID {
+	if commonPattern.CommonThirdPartyID == nil {
+		return nil
+	}
+
+	scope := coredata.NewScopeFromObjectID(tp.ID)
+
+	var vendor coredata.Vendor
+	if err := vendor.LoadByOrganizationIDAndCommonThirdPartyID(
+		ctx,
+		conn,
+		scope,
+		tp.OrganizationID,
+		*commonPattern.CommonThirdPartyID,
+	); err != nil {
+		return nil
+	}
+
+	return &vendor.ID
+}
--- a/pkg/probod/probod.go
+++ b/pkg/probod/probod.go
@@ -673,7 +673,7 @@ func (impl *Implm) Run(
 		},
 	)

-	trackerPatternAnalysisWorker := cookiebanner.NewPatternAnalysisWorker(cookieBannerService, pgClient, l.Named("tracker-pattern-analysis-worker"))
+	trackerPatternAnalysisWorker := cookiebanner.NewPatternAnalysisWorker(cookieBannerService, pgClient, l)
 	trackerPatternAnalysisWorkerCtx, stopTrackerPatternAnalysisWorker := context.WithCancel(context.Background())
 	wg.Go(
 		func() {
@@ -683,6 +683,16 @@ func (impl *Implm) Run(
 		},
 	)

+	trackerMappingWorker := probo.NewTrackerMappingWorker(pgClient, l)
+	trackerMappingWorkerCtx, stopTrackerMappingWorker := context.WithCancel(context.Background())
+	wg.Go(
+		func() {
+			if err := trackerMappingWorker.Run(trackerMappingWorkerCtx); err != nil {
+				cancel(fmt.Errorf("tracker mapping worker crashed: %w", err))
+			}
+		},
+	)
+
 	mailingListWorker := mailman.NewMailingListWorker(mailmanService, pgClient, l.Named("mailing-list-worker"))
 	mailingListWorkerCtx, stopMailingListWorker := context.WithCancel(context.Background())
 	wg.Go(
@@ -748,6 +758,7 @@ func (impl *Implm) Run(
 	stopWebhookSender()
 	stopESignService()
 	stopTrackerPatternAnalysisWorker()
+	stopTrackerMappingWorker()
 	stopMailingListWorker()
 	stopEvidenceDescriptionWorker()
 	stopDocumentPDFWorker()