astra
/
plc-mirror
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add whitelist for PDS hosts and update repo PDS pointer on appropriate occasions

main
Max Ignatenko 2024-02-22 12:05:44 +00:00
parent 600dac7694
commit 8f32c494f7
8 changed files with 108 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
cmd/consumer/consumer.go
View File

@ -28,6 +28,7 @@ import (
"github.com/uabluerail/indexer/models"
"github.com/uabluerail/indexer/pds"
"github.com/uabluerail/indexer/repo"
"github.com/uabluerail/indexer/util/resolver"
)
type BadRecord struct {
@ -245,10 +246,32 @@ func (c *Consumer) processMessage(ctx context.Context, typ string, r io.Reader,
if err != nil {
return fmt.Errorf("repo.EnsureExists(%q): %w", payload.Repo, err)
}
if repoInfo.PDS != models.ID(c.remote.ID) {
log.Error().Str("did", payload.Repo).Str("rev", payload.Rev).
Msgf("Commit from an incorrect PDS, skipping")
return nil
if repoInfo.PDS != c.remote.ID {
u, err := resolver.GetPDSEndpoint(ctx, payload.Repo)
if err == nil {
cur, err := pds.EnsureExists(ctx, c.db, u.String())
if err == nil {
if repoInfo.PDS != cur.ID {
// Repo was migrated, lets update our record.
err := c.db.Model(repoInfo).Where(&repo.Repo{ID: repoInfo.ID}).Updates(&repo.Repo{PDS: cur.ID}).Error
if err != nil {
log.Error().Err(err).Msgf("Repo %q was migrated to %q, but updating the repo has failed: %s", payload.Repo, cur.Host, err)
}
}
repoInfo.PDS = cur.ID
} else {
log.Error().Err(err).Msgf("Failed to get PDS record for %q: %s", u, err)
}
} else {
log.Error().Err(err).Msgf("Failed to get PDS endpoint for repo %q: %s", payload.Repo, err)
}
if repoInfo.PDS != c.remote.ID {
// We checked a recent version of DID doc and this is still not a correct PDS.
log.Error().Str("did", payload.Repo).Str("rev", payload.Rev).
Msgf("Commit from an incorrect PDS, skipping")
return nil
}
}
if created {
reposDiscovered.WithLabelValues(c.remote.Host).Inc()
@ -436,7 +459,11 @@ func (c *Consumer) processMessage(ctx context.Context, typ string, r io.Reader,
log.Error().Msgf("Unknown #info message %q: %+v", payload.Name, payload)
}
default:
log.Warn().Msgf("Unknown message type received: %s", typ)
b, err := io.ReadAll(r)
if err != nil {
log.Error().Err(err).Msgf("Failed to read message payload: %s", err)
}
log.Warn().Msgf("Unknown message type received: %s payload=%q", typ, string(b))
}
return nil
}

3
cmd/consumer/main.go
View File

@ -60,6 +60,9 @@ func runMain(ctx context.Context) error {
}
// TODO: check for changes and start/stop consumers as needed
for _, remote := range remotes {
if remote.Disabled {
continue
}
c, err := NewConsumer(ctx, &remote, db)
if err != nil {
return fmt.Errorf("failed to create a consumer for %q: %w", remote.Host, err)

11
cmd/lister/lister.go
View File

@ -57,13 +57,22 @@ func (l *Lister) run(ctx context.Context) {
remote := pds.PDS{}
if err := db.Model(&remote).
Where("last_list is null or last_list < ?", time.Now().Add(-l.listRefreshInterval)).
Where("disabled=false and (last_list is null or last_list < ?)", time.Now().Add(-l.listRefreshInterval)).
Take(&remote).Error; err != nil {
if !errors.Is(err, gorm.ErrRecordNotFound) {
log.Error().Err(err).Msgf("Failed to query DB for a PDS to list repos from: %s", err)
}
break
}
if !pds.IsWhitelisted(remote.Host) {
log.Info().Msgf("PDS %q is not whitelisted, disabling it", remote.Host)
if err := db.Model(&remote).Where(&pds.PDS{ID: remote.ID}).Updates(&pds.PDS{Disabled: true}).Error; err != nil {
log.Error().Err(err).Msgf("Failed to disable PDS %q: %s", remote.Host, err)
}
break
}
client := xrpcauth.NewAnonymousClient(ctx)
client.Host = remote.Host

7
cmd/record-indexer/scheduler.go
View File

@ -3,6 +3,7 @@ package main
import (
"context"
"fmt"
"slices"
"time"
"github.com/rs/zerolog"
@ -101,7 +102,11 @@ func (s *Scheduler) fillQueue(ctx context.Context) error {
if err := s.db.Find(&remotes).Error; err != nil {
return fmt.Errorf("failed to get the list of PDSs: %w", err)
}
perPDSLimit := 0
remotes = slices.DeleteFunc(remotes, func(pds pds.PDS) bool {
return pds.Disabled
})
perPDSLimit := maxQueueLen
if len(remotes) > 0 {
perPDSLimit = maxQueueLen * 2 / len(remotes)
}

26
cmd/record-indexer/workerpool.go
View File

@ -4,7 +4,6 @@ import (
"bytes"
"context"
"fmt"
"net/url"
"regexp"
"strings"
"time"
@ -143,27 +142,20 @@ func (p *WorkerPool) doWork(ctx context.Context, work WorkItem) error {
log := zerolog.Ctx(ctx)
defer close(work.signal)
doc, err := resolver.GetDocument(ctx, work.Repo.DID)
u, err := resolver.GetPDSEndpoint(ctx, work.Repo.DID)
if err != nil {
return fmt.Errorf("resolving did %q: %w", work.Repo.DID, err)
return err
}
pdsHost := ""
for _, srv := range doc.Service {
if srv.Type != "AtprotoPersonalDataServer" {
continue
}
pdsHost = srv.ServiceEndpoint
}
if pdsHost == "" {
return fmt.Errorf("did not find any PDS in DID Document")
}
u, err := url.Parse(pdsHost)
remote, err := pds.EnsureExists(ctx, p.db, u.String())
if err != nil {
return fmt.Errorf("PDS endpoint (%q) is an invalid URL: %w", pdsHost, err)
return fmt.Errorf("failed to get PDS records for %q: %w", u, err)
}
if u.Host == "" {
return fmt.Errorf("PDS endpoint (%q) doesn't have a host part", pdsHost)
if work.Repo.PDS != remote.ID {
if err := p.db.Model(&work.Repo).Where(&repo.Repo{ID: work.Repo.ID}).Updates(&repo.Repo{PDS: remote.ID}).Error; err != nil {
return fmt.Errorf("failed to update repo's PDS to %q: %w", u, err)
}
work.Repo.PDS = remote.ID
}
client := xrpcauth.NewAnonymousClient(ctx)

19
pds/pds.go
View File

@ -3,6 +3,7 @@ package pds
import (
"context"
"fmt"
"path/filepath"
"time"
"gorm.io/gorm"
@ -12,6 +13,11 @@ import (
const Unknown models.ID = 0
var whitelist []string = []string{
"https://bsky.social",
"https://*.bsky.network",
}
type PDS struct {
ID models.ID `gorm:"primarykey"`
CreatedAt time.Time
@ -21,6 +27,7 @@ type PDS struct {
FirstCursorSinceReset int64
LastList time.Time
CrawlLimit int
Disabled bool
}
func AutoMigrate(db *gorm.DB) error {
@ -28,9 +35,21 @@ func AutoMigrate(db *gorm.DB) error {
}
func EnsureExists(ctx context.Context, db *gorm.DB, host string) (*PDS, error) {
if !IsWhitelisted(host) {
return nil, fmt.Errorf("host %q is not whitelisted", host)
}
remote := PDS{Host: host}
if err := db.Model(&remote).Where(&PDS{Host: host}).FirstOrCreate(&remote).Error; err != nil {
return nil, fmt.Errorf("failed to get PDS record from DB for %q: %w", remote.Host, err)
}
return &remote, nil
}
func IsWhitelisted(host string) bool {
for _, p := range whitelist {
if match, _ := filepath.Match(p, host); match {
return true
}
}
return false
}

20
repo/repo.go
View File

@ -5,7 +5,6 @@ import (
"encoding/json"
"errors"
"fmt"
"net/url"
"time"
"gorm.io/gorm"
@ -66,28 +65,11 @@ func EnsureExists(ctx context.Context, db *gorm.DB, did string) (*Repo, bool, er
// if we do - compare PDS IDs
// if they don't match - also reset FirstRevSinceReset
doc, err := resolver.GetDocument(ctx, did)
u, err := resolver.GetPDSEndpoint(ctx, did)
if err != nil {
return nil, false, fmt.Errorf("fetching DID Document: %w", err)
}
pdsHost := ""
for _, srv := range doc.Service {
if srv.Type != "AtprotoPersonalDataServer" {
continue
}
pdsHost = srv.ServiceEndpoint
}
if pdsHost == "" {
return nil, false, fmt.Errorf("did not find any PDS in DID Document")
}
u, err := url.Parse(pdsHost)
if err != nil {
return nil, false, fmt.Errorf("PDS endpoint (%q) is an invalid URL: %w", pdsHost, err)
}
if u.Host == "" {
return nil, false, fmt.Errorf("PDS endpoint (%q) doesn't have a host part", pdsHost)
}
remote, err := pds.EnsureExists(ctx, db, u.String())
if err != nil {
return nil, false, fmt.Errorf("failed to get PDS record from DB for %q: %w", remote.Host, err)

28
util/resolver/resolver.go
View File

@ -3,6 +3,8 @@ package resolver
import (
"context"
"errors"
"fmt"
"net/url"
"os"
"github.com/bluesky-social/indigo/api"
@ -56,3 +58,29 @@ func (r *fallbackResolver) FlushCacheFor(did string) {
res.FlushCacheFor(did)
}
}
func GetPDSEndpoint(ctx context.Context, did string) (*url.URL, error) {
doc, err := GetDocument(ctx, did)
if err != nil {
return nil, fmt.Errorf("resolving did %q: %w", did, err)
}
pdsHost := ""
for _, srv := range doc.Service {
if srv.Type != "AtprotoPersonalDataServer" {
continue
}
pdsHost = srv.ServiceEndpoint
}
if pdsHost == "" {
return nil, fmt.Errorf("did not find any PDS in DID Document")
}
u, err := url.Parse(pdsHost)
if err != nil {
return nil, fmt.Errorf("PDS endpoint (%q) is an invalid URL: %w", pdsHost, err)
}
if u.Host == "" {
return nil, fmt.Errorf("PDS endpoint (%q) doesn't have a host part", pdsHost)
}
return u, nil
}