From e26ee335f6d06fe2dcf52085288d1b74a121d875 Mon Sep 17 00:00:00 2001
From: Andrew Plaza <github@andrewplaza.dev>
Date: Thu, 25 Apr 2024 15:20:29 -0400
Subject: [PATCH] revocation for removed members

---
 pkg/mls/store/queries.sql            | 10 ++++++++++
 pkg/mls/store/queries/queries.sql.go | 22 ++++++++++++++++++++++
 pkg/mls/store/store.go               | 11 ++++++++++-
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/pkg/mls/store/queries.sql b/pkg/mls/store/queries.sql
index 57b697c9..a5cfa94c 100644
--- a/pkg/mls/store/queries.sql
+++ b/pkg/mls/store/queries.sql
@@ -41,6 +41,16 @@ INSERT INTO inbox_log (inbox_id, server_timestamp_ns, identity_update_proto)
 VALUES ($1, $2, $3)
 RETURNING sequence_id;
 
+-- name: RevokeAddressFromLog :exec
+UPDATE address_log
+SET revocation_sequence_id = $1
+WHERE (address, inbox_id, association_sequence_id) = (
+    SELECT address, inbox_id, MAX(association_sequence_id)
+    FROM address_log AS a
+    WHERE a.address = $2 AND a.inbox_id = $3
+    GROUP BY address, inbox_id
+);
+
 -- name: CreateInstallation :exec
 INSERT INTO installations (id, wallet_address, created_at, updated_at, credential_identity, key_package, expiration)
 VALUES ($1, $2, $3, $3, $4, $5, $6);
diff --git a/pkg/mls/store/queries/queries.sql.go b/pkg/mls/store/queries/queries.sql.go
index acbdcebb..14aa9425 100644
--- a/pkg/mls/store/queries/queries.sql.go
+++ b/pkg/mls/store/queries/queries.sql.go
@@ -511,6 +511,28 @@ func (q *Queries) QueryGroupMessagesWithCursorDesc(ctx context.Context, arg Quer
 	return items, nil
 }
 
+const revokeAddressFromLog = `-- name: RevokeAddressFromLog :exec
+UPDATE address_log
+SET revocation_sequence_id = $1
+WHERE (address, inbox_id, association_sequence_id) = (
+    SELECT address, inbox_id, MAX(association_sequence_id)
+    FROM address_log AS a
+    WHERE a.address = $2 AND a.inbox_id = $3
+    GROUP BY address, inbox_id
+)
+`
+
+type RevokeAddressFromLogParams struct {
+	RevocationSequenceID sql.NullInt64
+	Address              string
+	InboxID              string
+}
+
+func (q *Queries) RevokeAddressFromLog(ctx context.Context, arg RevokeAddressFromLogParams) error {
+	_, err := q.db.ExecContext(ctx, revokeAddressFromLog, arg.RevocationSequenceID, arg.Address, arg.InboxID)
+	return err
+}
+
 const revokeInstallation = `-- name: RevokeInstallation :exec
 UPDATE installations
 SET revoked_at = $1
diff --git a/pkg/mls/store/store.go b/pkg/mls/store/store.go
index d1850151..23030ed2 100644
--- a/pkg/mls/store/store.go
+++ b/pkg/mls/store/store.go
@@ -162,7 +162,16 @@ func (s *Store) PublishIdentityUpdate(ctx context.Context, req *identity.Publish
 
 		// Update address log table
 		for _, removed_member := range state.StateDiff.RemovedMembers {
-
+			if address, ok := removed_member.Kind.(*associations.MemberIdentifier_Address); ok {
+				err = txQueries.RevokeAddressFromLog(ctx, queries.RevokeAddressFromLogParams{
+					Address:              address.Address,
+					InboxID:              state.AssociationState.InboxId,
+					RevocationSequenceID: sql.NullInt64{Valid: true, Int64: sequence_id},
+				})
+				if err != nil {
+					return err
+				}
+			}
 		}
 
 		return nil