mirror of
https://github.com/phishingclub/phishingclub.git
synced 2026-09-18 07:02:19 +02:00
perf, fix random recipient table scan
Signed-off-by: RonniSkansing <rskansing@gmail.com>
This commit is contained in:
@@ -2,6 +2,7 @@ package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strconv"
|
||||
@@ -744,19 +745,31 @@ func (r *Recipient) GetRandomByCompanyID(
|
||||
) (*model.Recipient, error) {
|
||||
var dbRecipient database.Recipient
|
||||
|
||||
db := r.DB.Table(database.RECIPIENT_TABLE)
|
||||
|
||||
// apply company filter
|
||||
db = whereCompany(db, database.RECIPIENT_TABLE, companyID)
|
||||
|
||||
// exclude specific recipient if provided
|
||||
if excludeRecipientID != nil {
|
||||
db = db.Where(fmt.Sprintf("%s != ?", TableColumnID(database.RECIPIENT_TABLE)), excludeRecipientID)
|
||||
// build the company (and optional exclude) filter into a fresh query each call
|
||||
filtered := func() *gorm.DB {
|
||||
q := r.DB.Table(database.RECIPIENT_TABLE)
|
||||
q = whereCompany(q, database.RECIPIENT_TABLE, companyID)
|
||||
if excludeRecipientID != nil {
|
||||
q = q.Where(fmt.Sprintf("%s != ?", TableColumnID(database.RECIPIENT_TABLE)), excludeRecipientID)
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
// order randomly and get one
|
||||
res := db.Order("RANDOM()").Limit(1).First(&dbRecipient)
|
||||
|
||||
// pick a random rowid threshold and take the first matching row at or after it.
|
||||
// this rides the rowid index; ORDER BY RANDOM instead sorts the whole table on
|
||||
// every call, which gets slower as the recipients table grows. if the threshold
|
||||
// lands past the last match, wrap to the first matching row.
|
||||
res := filtered().
|
||||
Where(fmt.Sprintf(
|
||||
"%[1]s.rowid >= (abs(random()) %% (SELECT max(rowid) + 1 FROM %[1]s))",
|
||||
database.RECIPIENT_TABLE,
|
||||
)).
|
||||
Order(database.RECIPIENT_TABLE + ".rowid").
|
||||
Limit(1).
|
||||
First(&dbRecipient)
|
||||
if errors.Is(res.Error, gorm.ErrRecordNotFound) {
|
||||
res = filtered().Order(database.RECIPIENT_TABLE + ".rowid").Limit(1).First(&dbRecipient)
|
||||
}
|
||||
if res.Error != nil {
|
||||
return nil, res.Error
|
||||
}
|
||||
|
||||
@@ -78,8 +78,21 @@ func (t *Template) CreateMail(
|
||||
apiSender,
|
||||
)
|
||||
|
||||
// add random recipient data to template context (excluding current recipient)
|
||||
(*data)["RandomRecipient"] = t.getRandomRecipientData(ctx, companyID, &rid)
|
||||
// add random recipient data only when the email actually uses it (the lookup is
|
||||
// costly), excluding the current recipient
|
||||
if email != nil {
|
||||
subject := ""
|
||||
if v, err := email.MailHeaderSubject.Get(); err == nil {
|
||||
subject = v.String()
|
||||
}
|
||||
content := ""
|
||||
if v, err := email.Content.Get(); err == nil {
|
||||
content = v.String()
|
||||
}
|
||||
if usesRandomRecipient(subject, content) {
|
||||
(*data)["RandomRecipient"] = t.getRandomRecipientData(ctx, companyID, &rid)
|
||||
}
|
||||
}
|
||||
|
||||
// full report endpoint URL for this recipient on the campaign domain, so a report
|
||||
// header can carry a ready-to-call link instead of just the token
|
||||
@@ -524,12 +537,15 @@ func (t *Template) CreatePhishingPageWithCampaignAndRecipient(
|
||||
nil, // apiSender
|
||||
)
|
||||
|
||||
// add random recipient data to template context (excluding current recipient)
|
||||
var excludeRecipientID *uuid.UUID
|
||||
if rid, err := recipient.ID.Get(); err == nil {
|
||||
excludeRecipientID = &rid
|
||||
// add random recipient data only when the page template uses it (the lookup is
|
||||
// costly), excluding the current recipient
|
||||
if usesRandomRecipient(contentToRender) {
|
||||
var excludeRecipientID *uuid.UUID
|
||||
if rid, err := recipient.ID.Get(); err == nil {
|
||||
excludeRecipientID = &rid
|
||||
}
|
||||
(*data)["RandomRecipient"] = t.getRandomRecipientData(ctx, companyID, excludeRecipientID)
|
||||
}
|
||||
(*data)["RandomRecipient"] = t.getRandomRecipientData(ctx, companyID, excludeRecipientID)
|
||||
|
||||
// direct URLs to each stage of the flow, available when this page is served as part
|
||||
// of a real campaign. before and after stay empty when the flow has no such stage.
|
||||
@@ -953,6 +969,18 @@ func (t *Template) TemplateFuncsWithDeviceCode(
|
||||
return funcs
|
||||
}
|
||||
|
||||
// usesRandomRecipient reports whether any of the template texts reference the
|
||||
// RandomRecipient variable. The random recipient lookup is costly, so it is only
|
||||
// done when a template actually uses it.
|
||||
func usesRandomRecipient(texts ...string) bool {
|
||||
for _, s := range texts {
|
||||
if strings.Contains(s, "RandomRecipient") {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// getRandomRecipientData gets a random recipient from a company and returns a map of their data
|
||||
func (t *Template) getRandomRecipientData(ctx context.Context, companyID *uuid.UUID, excludeRecipientID *uuid.UUID) map[string]string {
|
||||
data := map[string]string{
|
||||
|
||||
Reference in New Issue
Block a user