update remote account get/deref logic

1 files changed, 180 insertions, 106 deletions

diff --git a/internal/federation/dereferencing/account.go b/internal/federation/dereferencing/account.go
index 581c95de2..27591d857 100644
--- a/internal/federation/dereferencing/account.go
+++ b/internal/federation/dereferencing/account.go
@@ -26,6 +26,7 @@ import (
 	"io"
 	"net/url"
 	"strings"
+	"sync"
 
 	"github.com/sirupsen/logrus"
 	"github.com/superseriousbusiness/activity/streams"
@@ -44,30 +45,6 @@ func instanceAccount(account *gtsmodel.Account) bool {
 		(account.Username == "internal.fetch" && strings.Contains(account.Note, "internal service actor"))
 }
 
-// EnrichRemoteAccount takes an account that's already been inserted into the database in a minimal form,
-// and populates it with additional fields, media, etc.
-//
-// EnrichRemoteAccount is mostly useful for calling after an account has been initially created by
-// the federatingDB's Create function, or during the federated authorization flow.
-func (d *deref) EnrichRemoteAccount(ctx context.Context, username string, account *gtsmodel.Account) (*gtsmodel.Account, error) {
-	// if we're dealing with an instance account, we don't need to update anything
-	if instanceAccount(account) {
-		return account, nil
-	}
-
-	if err := d.PopulateAccountFields(ctx, account, username, false); err != nil {
-		return nil, err
-	}
-
-	updated, err := d.db.UpdateAccount(ctx, account)
-	if err != nil {
-		logrus.Errorf("EnrichRemoteAccount: error updating account: %s", err)
-		return account, nil
-	}
-
-	return updated, nil
-}
-
 // GetRemoteAccount completely dereferences a remote account, converts it to a GtS model account,
 // puts it in the database, and returns it to a caller. The boolean indicates whether the account is new
 // to us or not. If we haven't seen the account before, bool will be true. If we have seen the account before,
@@ -77,60 +54,73 @@ func (d *deref) EnrichRemoteAccount(ctx context.Context, username string, accoun
 // the remote instance again.
 //
 // SIDE EFFECTS: remote account will be stored in the database, or updated if it already exists (and refresh is true).
-func (d *deref) GetRemoteAccount(ctx context.Context, username string, remoteAccountID *url.URL, refresh bool) (*gtsmodel.Account, bool, error) {
+func (d *deref) GetRemoteAccount(ctx context.Context, username string, remoteAccountID *url.URL, refresh bool, blocking bool) (*gtsmodel.Account, error) {
 	new := true
 
-	// check if we already have the account in our db
-	maybeAccount, err := d.db.GetAccountByURI(ctx, remoteAccountID.String())
+	// check if we already have the account in our db, and just return it unless we'd doing a refresh
+	remoteAccount, err := d.db.GetAccountByURI(ctx, remoteAccountID.String())
 	if err == nil {
-		// we've seen this account before so it's not new
 		new = false
 		if !refresh {
-			// we're not being asked to refresh, but just in case we don't have the avatar/header cached yet....
-			maybeAccount, err = d.EnrichRemoteAccount(ctx, username, maybeAccount)
-			return maybeAccount, new, err
+			// make sure the account fields are populated before returning:
+			// even if we're not doing a refresh, the caller might want to block
+			// until everything is loaded
+			err = d.populateAccountFields(ctx, remoteAccount, username, refresh, blocking)
+			return remoteAccount, err
 		}
 	}
 
-	accountable, err := d.dereferenceAccountable(ctx, username, remoteAccountID)
-	if err != nil {
-		return nil, new, fmt.Errorf("FullyDereferenceAccount: error dereferencing accountable: %s", err)
-	}
-
-	gtsAccount, err := d.typeConverter.ASRepresentationToAccount(ctx, accountable, refresh)
-	if err != nil {
-		return nil, new, fmt.Errorf("FullyDereferenceAccount: error converting accountable to account: %s", err)
-	}
-
 	if new {
-		// generate a new id since we haven't seen this account before, and do a put
-		ulid, err := id.NewRandomULID()
+		// we haven't seen this account before: dereference it from remote
+		accountable, err := d.dereferenceAccountable(ctx, username, remoteAccountID)
 		if err != nil {
-			return nil, new, fmt.Errorf("FullyDereferenceAccount: error generating new id for account: %s", err)
+			return nil, fmt.Errorf("GetRemoteAccount: error dereferencing accountable: %s", err)
 		}
-		gtsAccount.ID = ulid
 
-		if err := d.PopulateAccountFields(ctx, gtsAccount, username, refresh); err != nil {
-			return nil, new, fmt.Errorf("FullyDereferenceAccount: error populating further account fields: %s", err)
+		newAccount, err := d.typeConverter.ASRepresentationToAccount(ctx, accountable, refresh)
+		if err != nil {
+			return nil, fmt.Errorf("GetRemoteAccount: error converting accountable to account: %s", err)
 		}
 
-		if err := d.db.Put(ctx, gtsAccount); err != nil {
-			return nil, new, fmt.Errorf("FullyDereferenceAccount: error putting new account: %s", err)
+		ulid, err := id.NewRandomULID()
+		if err != nil {
+			return nil, fmt.Errorf("GetRemoteAccount: error generating new id for account: %s", err)
 		}
-	} else {
-		// take the id we already have and do an update
-		gtsAccount.ID = maybeAccount.ID
-		if err := d.PopulateAccountFields(ctx, gtsAccount, username, refresh); err != nil {
-			return nil, new, fmt.Errorf("FullyDereferenceAccount: error populating further account fields: %s", err)
+		newAccount.ID = ulid
+
+		if err := d.populateAccountFields(ctx, newAccount, username, refresh, blocking); err != nil {
+			return nil, fmt.Errorf("GetRemoteAccount: error populating further account fields: %s", err)
 		}
 
-		gtsAccount, err = d.db.UpdateAccount(ctx, gtsAccount)
-		if err != nil {
-			return nil, false, fmt.Errorf("EnrichRemoteAccount: error updating account: %s", err)
+		if err := d.db.Put(ctx, newAccount); err != nil {
+			return nil, fmt.Errorf("GetRemoteAccount: error putting new account: %s", err)
 		}
+
+		return newAccount, nil
+	}
+
+	// we have seen this account before, but we have to refresh it
+	refreshedAccountable, err := d.dereferenceAccountable(ctx, username, remoteAccountID)
+	if err != nil {
+		return nil, fmt.Errorf("GetRemoteAccount: error dereferencing refreshedAccountable: %s", err)
 	}
 
-	return gtsAccount, new, nil
+	refreshedAccount, err := d.typeConverter.ASRepresentationToAccount(ctx, refreshedAccountable, refresh)
+	if err != nil {
+		return nil, fmt.Errorf("GetRemoteAccount: error converting refreshedAccountable to refreshedAccount: %s", err)
+	}
+	refreshedAccount.ID = remoteAccount.ID
+
+	if err := d.populateAccountFields(ctx, refreshedAccount, username, refresh, blocking); err != nil {
+		return nil, fmt.Errorf("GetRemoteAccount: error populating further refreshedAccount fields: %s", err)
+	}
+
+	updatedAccount, err := d.db.UpdateAccount(ctx, refreshedAccount)
+	if err != nil {
+		return nil, fmt.Errorf("GetRemoteAccount: error updating refreshedAccount: %s", err)
+	}
+
+	return updatedAccount, nil
 }
 
 // dereferenceAccountable calls remoteAccountID with a GET request, and tries to parse whatever
@@ -201,93 +191,177 @@ func (d *deref) dereferenceAccountable(ctx context.Context, username string, rem
 	return nil, fmt.Errorf("DereferenceAccountable: type name %s not supported", t.GetTypeName())
 }
 
-// PopulateAccountFields populates any fields on the given account that weren't populated by the initial
+// populateAccountFields populates any fields on the given account that weren't populated by the initial
 // dereferencing. This includes things like header and avatar etc.
-func (d *deref) PopulateAccountFields(ctx context.Context, account *gtsmodel.Account, requestingUsername string, refresh bool) error {
-	l := logrus.WithFields(logrus.Fields{
-		"func":               "PopulateAccountFields",
-		"requestingUsername": requestingUsername,
-	})
+func (d *deref) populateAccountFields(ctx context.Context, account *gtsmodel.Account, requestingUsername string, refresh bool, blocking bool) error {
+	// if we're dealing with an instance account, just bail, we don't need to do anything
+	if instanceAccount(account) {
+		return nil
+	}
 
 	accountURI, err := url.Parse(account.URI)
 	if err != nil {
-		return fmt.Errorf("PopulateAccountFields: couldn't parse account URI %s: %s", account.URI, err)
+		return fmt.Errorf("populateAccountFields: couldn't parse account URI %s: %s", account.URI, err)
 	}
+
 	if blocked, err := d.db.IsDomainBlocked(ctx, accountURI.Host); blocked || err != nil {
-		return fmt.Errorf("PopulateAccountFields: domain %s is blocked", accountURI.Host)
+		return fmt.Errorf("populateAccountFields: domain %s is blocked", accountURI.Host)
 	}
 
 	t, err := d.transportController.NewTransportForUsername(ctx, requestingUsername)
 	if err != nil {
-		return fmt.Errorf("PopulateAccountFields: error getting transport for user: %s", err)
+		return fmt.Errorf("populateAccountFields: error getting transport for user: %s", err)
 	}
 
 	// fetch the header and avatar
-	if err := d.fetchHeaderAndAviForAccount(ctx, account, t, refresh); err != nil {
-		// if this doesn't work, just skip it -- we can do it later
-		l.Debugf("error fetching header/avi for account: %s", err)
+	if err := d.fetchRemoteAccountMedia(ctx, account, t, refresh, blocking); err != nil {
+		return fmt.Errorf("populateAccountFields: error fetching header/avi for account: %s", err)
 	}
 
 	return nil
 }
 
-// fetchHeaderAndAviForAccount fetches the header and avatar for a remote account, using a transport
-// on behalf of requestingUsername.
+// fetchRemoteAccountMedia fetches and stores the header and avatar for a remote account,
+// using a transport on behalf of requestingUsername.
 //
 // targetAccount's AvatarMediaAttachmentID and HeaderMediaAttachmentID will be updated as necessary.
 //
-// SIDE EFFECTS: remote header and avatar will be stored in local storage.
-func (d *deref) fetchHeaderAndAviForAccount(ctx context.Context, targetAccount *gtsmodel.Account, t transport.Transport, refresh bool) error {
+// If refresh is true, then the media will be fetched again even if it's already been fetched before.
+//
+// If blocking is true, then the calls to the media manager made by this function will be blocking:
+// in other words, the function won't return until the header and the avatar have been fully processed.
+func (d *deref) fetchRemoteAccountMedia(ctx context.Context, targetAccount *gtsmodel.Account, t transport.Transport, refresh bool, blocking bool) error {
 	accountURI, err := url.Parse(targetAccount.URI)
 	if err != nil {
-		return fmt.Errorf("fetchHeaderAndAviForAccount: couldn't parse account URI %s: %s", targetAccount.URI, err)
+		return fmt.Errorf("fetchRemoteAccountMedia: couldn't parse account URI %s: %s", targetAccount.URI, err)
 	}
+
 	if blocked, err := d.db.IsDomainBlocked(ctx, accountURI.Host); blocked || err != nil {
-		return fmt.Errorf("fetchHeaderAndAviForAccount: domain %s is blocked", accountURI.Host)
+		return fmt.Errorf("fetchRemoteAccountMedia: domain %s is blocked", accountURI.Host)
 	}
 
 	if targetAccount.AvatarRemoteURL != "" && (targetAccount.AvatarMediaAttachmentID == "" || refresh) {
-		avatarIRI, err := url.Parse(targetAccount.AvatarRemoteURL)
-		if err != nil {
-			return err
-		}
+		var processingMedia *media.ProcessingMedia
 
-		data := func(innerCtx context.Context) (io.Reader, int, error) {
-			return t.DereferenceMedia(innerCtx, avatarIRI)
+		// first check if we're already processing this media
+		d.dereferencingAvatarsLock.Lock()
+		if alreadyProcessing, ok := d.dereferencingAvatars[targetAccount.ID]; ok {
+			// we're already on it, no worries
+			processingMedia = alreadyProcessing
 		}
-
-		avatar := true
-		processingMedia, err := d.mediaManager.ProcessMedia(ctx, data, targetAccount.ID, &media.AdditionalMediaInfo{
-			RemoteURL: &targetAccount.AvatarRemoteURL,
-			Avatar:    &avatar,
-		})
-		if err != nil {
-			return err
+		d.dereferencingAvatarsLock.Unlock()
+
+		if processingMedia == nil {
+			// we're not already processing it so start now
+			avatarIRI, err := url.Parse(targetAccount.AvatarRemoteURL)
+			if err != nil {
+				return err
+			}
+
+			data := func(innerCtx context.Context) (io.Reader, int, error) {
+				return t.DereferenceMedia(innerCtx, avatarIRI)
+			}
+
+			avatar := true
+			newProcessing, err := d.mediaManager.ProcessMedia(ctx, data, targetAccount.ID, &media.AdditionalMediaInfo{
+				RemoteURL: &targetAccount.AvatarRemoteURL,
+				Avatar:    &avatar,
+			})
+			if err != nil {
+				return err
+			}
+			targetAccount.AvatarMediaAttachmentID = newProcessing.AttachmentID()
+
+			// store it in our map to indicate it's in process
+			d.dereferencingAvatarsLock.Lock()
+			d.dereferencingAvatars[targetAccount.ID] = newProcessing
+			d.dereferencingAvatarsLock.Unlock()
+
+			processingMedia = newProcessing
 		}
 
-		targetAccount.AvatarMediaAttachmentID = processingMedia.AttachmentID()
+		// block until loaded if required...
+		if blocking {
+			if err := lockAndLoad(ctx, d.dereferencingAvatarsLock, processingMedia, d.dereferencingAvatars, targetAccount.ID); err != nil {
+				return err
+			}
+		} else {
+			// ...otherwise do it async
+			go func() {
+				if err := lockAndLoad(ctx, d.dereferencingAvatarsLock, processingMedia, d.dereferencingAvatars, targetAccount.ID); err != nil {
+					logrus.Errorf("fetchRemoteAccountMedia: error during async lock and load of avatar: %s", err)
+				}
+			}()
+		}
 	}
 
 	if targetAccount.HeaderRemoteURL != "" && (targetAccount.HeaderMediaAttachmentID == "" || refresh) {
-		headerIRI, err := url.Parse(targetAccount.HeaderRemoteURL)
-		if err != nil {
-			return err
-		}
+		var processingMedia *media.ProcessingMedia
 
-		data := func(innerCtx context.Context) (io.Reader, int, error) {
-			return t.DereferenceMedia(innerCtx, headerIRI)
+		// first check if we're already processing this media
+		d.dereferencingHeadersLock.Lock()
+		if alreadyProcessing, ok := d.dereferencingHeaders[targetAccount.ID]; ok {
+			// we're already on it, no worries
+			processingMedia = alreadyProcessing
 		}
-
-		header := true
-		processingMedia, err := d.mediaManager.ProcessMedia(ctx, data, targetAccount.ID, &media.AdditionalMediaInfo{
-			RemoteURL: &targetAccount.HeaderRemoteURL,
-			Header:    &header,
-		})
-		if err != nil {
-			return err
+		d.dereferencingHeadersLock.Unlock()
+
+		if processingMedia == nil {
+			// we're not already processing it so start now
+			headerIRI, err := url.Parse(targetAccount.HeaderRemoteURL)
+			if err != nil {
+				return err
+			}
+
+			data := func(innerCtx context.Context) (io.Reader, int, error) {
+				return t.DereferenceMedia(innerCtx, headerIRI)
+			}
+
+			header := true
+			newProcessing, err := d.mediaManager.ProcessMedia(ctx, data, targetAccount.ID, &media.AdditionalMediaInfo{
+				RemoteURL: &targetAccount.HeaderRemoteURL,
+				Header:    &header,
+			})
+			if err != nil {
+				return err
+			}
+			targetAccount.HeaderMediaAttachmentID = newProcessing.AttachmentID()
+
+			// store it in our map to indicate it's in process
+			d.dereferencingHeadersLock.Lock()
+			d.dereferencingHeaders[targetAccount.ID] = newProcessing
+			d.dereferencingHeadersLock.Unlock()
+
+			processingMedia = newProcessing
 		}
 
-		targetAccount.HeaderMediaAttachmentID = processingMedia.AttachmentID()
+		// block until loaded if required...
+		if blocking {
+			if err := lockAndLoad(ctx, d.dereferencingHeadersLock, processingMedia, d.dereferencingHeaders, targetAccount.ID); err != nil {
+				return err
+			}
+		} else {
+			// ...otherwise do it async
+			go func() {
+				if err := lockAndLoad(ctx, d.dereferencingHeadersLock, processingMedia, d.dereferencingHeaders, targetAccount.ID); err != nil {
+					logrus.Errorf("fetchRemoteAccountMedia: error during async lock and load of header: %s", err)
+				}
+			}()
+		}
 	}
+
 	return nil
 }
+
+func lockAndLoad(ctx context.Context, lock *sync.Mutex, processing *media.ProcessingMedia, processingMap map[string]*media.ProcessingMedia, accountID string) error {
+	// whatever happens, remove the in-process media from the map
+	defer func() {
+		lock.Lock()
+		delete(processingMap, accountID)
+		lock.Unlock()
+	}()
+
+	// try and load it
+	_, err := processing.LoadAttachment(ctx)
+	return err
+}