From 9acabe93d2dede7b588dc03ee57e4c27269a183d Mon Sep 17 00:00:00 2001
From: Mark Haines <mjark@negativecurvature.net>
Date: Fri, 7 Jul 2017 11:17:59 +0100
Subject: [PATCH] More stuff

---
 .../dendrite/roomserver/input/membership.go   | 132 ++++++++++++------
 .../dendrite/roomserver/storage/storage.go    |   5 +
 .../dendrite/roomserver/types/types.go        |   2 +
 3 files changed, 97 insertions(+), 42 deletions(-)

diff --git a/src/github.com/matrix-org/dendrite/roomserver/input/membership.go b/src/github.com/matrix-org/dendrite/roomserver/input/membership.go
index bd4353554..650586a7a 100644
--- a/src/github.com/matrix-org/dendrite/roomserver/input/membership.go
+++ b/src/github.com/matrix-org/dendrite/roomserver/input/membership.go
@@ -19,6 +19,96 @@ import (
 	"github.com/matrix-org/gomatrixserverlib"
 )
 
+func updateMemberships(
+	db RoomEventDatabase, updater types.RoomRecentEventsUpdater, removed, added []types.StateEntry,
+) error {
+	changes := membershipChanges(removed, added)
+	var eventNIDs []types.EventNID
+	for _, change := range changes {
+		if change.added.EventNID != 0 {
+			eventNIDs = append(eventNIDs, change.added.EventNID)
+		}
+		if change.removed.EventNID != 0 {
+			eventNIDs = append(eventNIDs, change.removed.EventNID)
+		}
+	}
+	events, err := db.Events(eventNIDs)
+	if err != nil {
+		return err
+	}
+
+	for _, change := range changes {
+		var ae *gomatrixserverlib.Event
+		var re *gomatrixserverlib.Event
+		var targetNID types.EventStateKeyNID
+		if change.removed.EventNID != 0 {
+			ev, _ := eventMap(events).lookup(change.removed.EventNID)
+			if ev != nil {
+				re = &ev.Event
+			}
+			targetNID = change.removed.EventStateKeyNID
+		}
+		if change.added.EventNID != 0 {
+			ev, _ := eventMap(events).lookup(change.added.EventNID)
+			if ae != nil {
+				ae = &ev.Event
+			}
+			targetNID = change.added.EventStateKeyNID
+		}
+		if err := updateMembership(updater, targetNID, re, ae); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func updateMembership(
+	updater types.RoomRecentEventsUpdater, targetNID types.EventStateKeyNID, remove *gomatrixserverlib.Event, add *gomatrixserverlib.Event,
+) error {
+	old := "leave"
+	new := "leave"
+
+	if remove != nil {
+		membership, err := remove.Membership()
+		if err != nil {
+			return err
+		}
+		old = *membership
+	}
+	if add != nil {
+		membership, err := add.Membership()
+		if err != nil {
+			return err
+		}
+		new = *membership
+	}
+	if old == new {
+		return nil
+	}
+
+	mu, err := updater.MembershipUpdater(targetNID)
+	if err != nil {
+		return err
+	}
+
+	switch new {
+	case "invite":
+		_, err := mu.SetToInvite(*add)
+		if err != nil {
+			return err
+		}
+	case "join":
+		if !mu.IsJoin() {
+			mu.SetToJoin(add.Sender())
+		}
+	case "leave":
+		if !mu.IsLeave() {
+			mu.SetToLeave(add.Sender())
+		}
+	}
+	return nil
+}
+
 type stateChange struct {
 	removed types.StateEntry
 	added   types.StateEntry
@@ -64,45 +154,3 @@ func membershipChanges(removed, added []types.StateEntry) []stateChange {
 		}
 	}
 }
-
-func updateMemberships(
-	db RoomEventDatabase, updater types.RoomRecentEventsUpdater, removed, added []types.StateEntry,
-) error {
-	changes := membershipChanges(removed, added)
-	var eventNIDs []types.EventNID
-	for _, change := range changes {
-		if change.added.EventNID != 0 {
-			eventNIDs = append(eventNIDs, change.added.EventNID)
-		}
-		if change.removed.EventNID != 0 {
-			eventNIDs = append(eventNIDs, change.removed.EventNID)
-		}
-	}
-	events, err := db.Events(eventNIDs)
-	if err != nil {
-		return err
-	}
-
-	for _, change := range changes {
-		var ae *gomatrixserverlib.Event
-		var re *gomatrixserverlib.Event
-		if change.removed.EventNID != 0 {
-			ev, _ := eventMap(events).lookup(change.removed.EventNID)
-			if ev != nil {
-				re = &ev.Event
-			}
-		}
-		if change.added.EventNID != 0 {
-			ev, _ := eventMap(events).lookup(change.added.EventNID)
-			if ae != nil {
-				ae = &ev.Event
-			}
-		}
-		updateMembership(re, ae)
-	}
-	return nil
-}
-
-func updateMembership(remove *gomatrixserverlib.Event, add *gomatrixserverlib.Event) {
-
-}
diff --git a/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go b/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go
index 5070289a6..84b2523a8 100644
--- a/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go
+++ b/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go
@@ -407,6 +407,11 @@ func (u *membershipUpdater) IsJoin() bool {
 	return u.membership == membershipStateJoin
 }
 
+// IsLeave implements types.MembershipUpdater
+func (u *membershipUpdater) IsLeave() bool {
+	return u.membership == membershipStateLeaveOrBan
+}
+
 // SetToInvite implements types.MembershipUpdater
 func (u *membershipUpdater) SetToInvite(event gomatrixserverlib.Event) (bool, error) {
 	// TODO: assign the state key inside the transaction.
diff --git a/src/github.com/matrix-org/dendrite/roomserver/types/types.go b/src/github.com/matrix-org/dendrite/roomserver/types/types.go
index ba12fe492..5b7d48a2f 100644
--- a/src/github.com/matrix-org/dendrite/roomserver/types/types.go
+++ b/src/github.com/matrix-org/dendrite/roomserver/types/types.go
@@ -186,6 +186,8 @@ type MembershipUpdater interface {
 	IsInvite() bool
 	// True if the target user is joined to the room.
 	IsJoin() bool
+	// True if the target user is not invited or joined to the room.
+	IsLeave() bool
 	// Set the state to invite.
 	// Returns whether this invite needs to be sent
 	SetToInvite(event gomatrixserverlib.Event) (needsSending bool, err error)