1 files changed, 68 insertions, 27 deletions

diff --git a/internal/federation/federatingdb/create.go b/internal/federation/federatingdb/create.go
index d9834b144..17ddf7c92 100644
--- a/internal/federation/federatingdb/create.go
+++ b/internal/federation/federatingdb/create.go
@@ -20,6 +20,7 @@ package federatingdb
 import (
 	"context"
 	"errors"
+	"slices"
 
 	"github.com/superseriousbusiness/activity/streams/vocab"
 	"github.com/superseriousbusiness/gotosocial/internal/ap"
@@ -138,6 +139,10 @@ func (f *federatingDB) createPollOptionables(
 		// then re-used in each further iteration.
 		inReplyTo *gtsmodel.Status
 
+		// existing poll vote by requesting user
+		// in receiving user's poll, if it exists.
+		vote *gtsmodel.PollVote
+
 		// the resulting slices of Poll.Option
 		// choice indices passed into the new
 		// created PollVote object.
@@ -169,15 +174,17 @@ func (f *federatingDB) createPollOptionables(
 			case inReplyTo.PollID == "":
 				return gtserror.Newf("poll vote in status %s without poll", statusURI)
 
-			// We don't own the poll ...
-			case !*inReplyTo.Local:
-				return gtserror.Newf("poll vote in remote status %s", statusURI)
+			// Ensure poll isn't closed.
+			case inReplyTo.Poll.Closed():
+				return gtserror.Newf("poll vote in closed poll %s", statusURI)
+
+			// Ensure receiver actually owns poll.
+			case inReplyTo.AccountID != receiver.ID:
+				return gtserror.Newf("receiving account %s does not own poll %s", receiver.URI, statusURI)
 			}
 
-			// Check whether user has already vote in this poll.
-			// (we only check this for the first object, as multiple
-			// may be sent in response to a multiple-choice poll).
-			vote, err := f.state.DB.GetPollVoteBy(
+			// Try load any existing vote(s) by user.
+			vote, err = f.state.DB.GetPollVoteBy(
 				gtscontext.SetBarebones(ctx),
 				inReplyTo.PollID,
 				requester.ID,
@@ -185,11 +192,6 @@ func (f *federatingDB) createPollOptionables(
 			if err != nil && !errors.Is(err, db.ErrNoEntries) {
 				return gtserror.Newf("error getting status %s poll votes from database: %w", statusURI, err)
 			}
-
-			if vote != nil {
-				log.Warnf(ctx, "%s has already voted in poll %s", requester.URI, statusURI)
-				return nil // this is a useful warning for admins to report to us from logs
-			}
 		}
 
 		if statusURI != inReplyTo.URI {
@@ -210,21 +212,60 @@ func (f *federatingDB) createPollOptionables(
 		choices = append(choices, choice)
 	}
 
-	// Enqueue message to the fedi API worker with poll vote(s).
-	f.state.Workers.Federator.Queue.Push(&messages.FromFediAPI{
-		APActivityType: ap.ActivityCreate,
-		APObjectType:   ap.ActivityQuestion,
-		GTSModel: &gtsmodel.PollVote{
-			ID:        id.NewULID(),
-			Choices:   choices,
-			AccountID: requester.ID,
-			Account:   requester,
-			PollID:    inReplyTo.PollID,
-			Poll:      inReplyTo.Poll,
-		},
-		Receiving:  receiver,
-		Requesting: requester,
-	})
+	if vote != nil {
+		// Ensure this isn't a multiple vote
+		// by same account in the same poll.
+		if !*inReplyTo.Poll.Multiple {
+			log.Warnf(ctx, "%s has already voted in single-choice poll %s", requester.URI, inReplyTo.URI)
+			return nil // this is a useful warning for admins to report to us from logs
+		}
+
+		// Drop all existing votes from the new slice of choices.
+		choices = slices.DeleteFunc(choices, func(choice int) bool {
+			return slices.Contains(vote.Choices, choice)
+		})
+
+		// Update poll with new choices but *not* voters.
+		inReplyTo.Poll.IncrementVotes(choices, false)
+
+		// Append new vote choices to existing
+		// vote, then sort them by index size.
+		vote.Choices = append(vote.Choices, choices...)
+		slices.Sort(vote.Choices)
+
+		// Populate the poll vote,
+		// later used in fedi worker.
+		vote.Poll = inReplyTo.Poll
+		vote.Account = requester
+
+		// TODO: would it be useful to store an UpdatedAt field
+		// on poll votes? i'm not sure we'd actually use it...
+
+		// Enqueue an update event for poll vote to fedi API worker.
+		f.state.Workers.Federator.Queue.Push(&messages.FromFediAPI{
+			APActivityType: ap.ActivityUpdate,
+			APObjectType:   ap.ActivityQuestion,
+			GTSModel:       vote,
+			Receiving:      receiver,
+			Requesting:     requester,
+		})
+	} else {
+		// Create new poll vote and enqueue create to fedi API worker.
+		f.state.Workers.Federator.Queue.Push(&messages.FromFediAPI{
+			APActivityType: ap.ActivityCreate,
+			APObjectType:   ap.ActivityQuestion,
+			GTSModel: &gtsmodel.PollVote{
+				ID:        id.NewULID(),
+				Choices:   choices,
+				AccountID: requester.ID,
+				Account:   requester,
+				PollID:    inReplyTo.PollID,
+				Poll:      inReplyTo.Poll,
+			},
+			Receiving:  receiver,
+			Requesting: requester,
+		})
+	}
 
 	return nil
 }