matrix-go-test/vendor/maunium.net/go/mautrix/appservice/statestore.go

// Copyright (c) 2020 Tulir Asokan
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package appservice

import (
	"sync"

	"maunium.net/go/mautrix/event"
	"maunium.net/go/mautrix/id"
)

type StateStore interface {
	IsRegistered(userID id.UserID) bool
	MarkRegistered(userID id.UserID)

	IsInRoom(roomID id.RoomID, userID id.UserID) bool
	IsInvited(roomID id.RoomID, userID id.UserID) bool
	IsMembership(roomID id.RoomID, userID id.UserID, allowedMemberships ...event.Membership) bool
	GetMember(roomID id.RoomID, userID id.UserID) *event.MemberEventContent
	TryGetMember(roomID id.RoomID, userID id.UserID) (*event.MemberEventContent, bool)
	SetMembership(roomID id.RoomID, userID id.UserID, membership event.Membership)
	SetMember(roomID id.RoomID, userID id.UserID, member *event.MemberEventContent)

	SetPowerLevels(roomID id.RoomID, levels *event.PowerLevelsEventContent)
	GetPowerLevels(roomID id.RoomID) *event.PowerLevelsEventContent
	GetPowerLevel(roomID id.RoomID, userID id.UserID) int
	GetPowerLevelRequirement(roomID id.RoomID, eventType event.Type) int
	HasPowerLevel(roomID id.RoomID, userID id.UserID, eventType event.Type) bool
}

func (as *AppService) UpdateState(evt *event.Event) {
	switch content := evt.Content.Parsed.(type) {
	case *event.MemberEventContent:
		as.StateStore.SetMember(evt.RoomID, id.UserID(evt.GetStateKey()), content)
	case *event.PowerLevelsEventContent:
		as.StateStore.SetPowerLevels(evt.RoomID, content)
	}
}

type BasicStateStore struct {
	Registrations map[id.UserID]bool                                    `json:"registrations"`
	Members       map[id.RoomID]map[id.UserID]*event.MemberEventContent `json:"memberships"`
	PowerLevels   map[id.RoomID]*event.PowerLevelsEventContent          `json:"power_levels"`

	registrationsLock sync.RWMutex
	membersLock       sync.RWMutex
	powerLevelsLock   sync.RWMutex
}

func NewBasicStateStore() StateStore {
	return &BasicStateStore{
		Registrations: make(map[id.UserID]bool),
		Members:       make(map[id.RoomID]map[id.UserID]*event.MemberEventContent),
		PowerLevels:   make(map[id.RoomID]*event.PowerLevelsEventContent),
	}
}

func (store *BasicStateStore) IsRegistered(userID id.UserID) bool {
	store.registrationsLock.RLock()
	defer store.registrationsLock.RUnlock()
	registered, ok := store.Registrations[userID]
	return ok && registered
}

func (store *BasicStateStore) MarkRegistered(userID id.UserID) {
	store.registrationsLock.Lock()
	defer store.registrationsLock.Unlock()
	store.Registrations[userID] = true
}

func (store *BasicStateStore) GetRoomMembers(roomID id.RoomID) map[id.UserID]*event.MemberEventContent {
	store.membersLock.RLock()
	members, ok := store.Members[roomID]
	store.membersLock.RUnlock()
	if !ok {
		members = make(map[id.UserID]*event.MemberEventContent)
		store.membersLock.Lock()
		store.Members[roomID] = members
		store.membersLock.Unlock()
	}
	return members
}

func (store *BasicStateStore) GetMembership(roomID id.RoomID, userID id.UserID) event.Membership {
	return store.GetMember(roomID, userID).Membership
}

func (store *BasicStateStore) GetMember(roomID id.RoomID, userID id.UserID) *event.MemberEventContent {
	member, ok := store.TryGetMember(roomID, userID)
	if !ok {
		member = &event.MemberEventContent{Membership: event.MembershipLeave}
	}
	return member
}

func (store *BasicStateStore) TryGetMember(roomID id.RoomID, userID id.UserID) (member *event.MemberEventContent, ok bool) {
	store.membersLock.RLock()
	defer store.membersLock.RUnlock()
	members, membersOk := store.Members[roomID]
	if !membersOk {
		return
	}
	member, ok = members[userID]
	return
}

func (store *BasicStateStore) IsInRoom(roomID id.RoomID, userID id.UserID) bool {
	return store.IsMembership(roomID, userID, "join")
}

func (store *BasicStateStore) IsInvited(roomID id.RoomID, userID id.UserID) bool {
	return store.IsMembership(roomID, userID, "join", "invite")
}

func (store *BasicStateStore) IsMembership(roomID id.RoomID, userID id.UserID, allowedMemberships ...event.Membership) bool {
	membership := store.GetMembership(roomID, userID)
	for _, allowedMembership := range allowedMemberships {
		if allowedMembership == membership {
			return true
		}
	}
	return false
}

func (store *BasicStateStore) SetMembership(roomID id.RoomID, userID id.UserID, membership event.Membership) {
	store.membersLock.Lock()
	members, ok := store.Members[roomID]
	if !ok {
		members = map[id.UserID]*event.MemberEventContent{
			userID: {Membership: membership},
		}
	} else {
		member, ok := members[userID]
		if !ok {
			members[userID] = &event.MemberEventContent{Membership: membership}
		} else {
			member.Membership = membership
			members[userID] = member
		}
	}
	store.Members[roomID] = members
	store.membersLock.Unlock()
}

func (store *BasicStateStore) SetMember(roomID id.RoomID, userID id.UserID, member *event.MemberEventContent) {
	store.membersLock.Lock()
	members, ok := store.Members[roomID]
	if !ok {
		members = map[id.UserID]*event.MemberEventContent{
			userID: member,
		}
	} else {
		members[userID] = member
	}
	store.Members[roomID] = members
	store.membersLock.Unlock()
}

func (store *BasicStateStore) SetPowerLevels(roomID id.RoomID, levels *event.PowerLevelsEventContent) {
	store.powerLevelsLock.Lock()
	store.PowerLevels[roomID] = levels
	store.powerLevelsLock.Unlock()
}

func (store *BasicStateStore) GetPowerLevels(roomID id.RoomID) (levels *event.PowerLevelsEventContent) {
	store.powerLevelsLock.RLock()
	levels = store.PowerLevels[roomID]
	store.powerLevelsLock.RUnlock()
	return
}

func (store *BasicStateStore) GetPowerLevel(roomID id.RoomID, userID id.UserID) int {
	return store.GetPowerLevels(roomID).GetUserLevel(userID)
}

func (store *BasicStateStore) GetPowerLevelRequirement(roomID id.RoomID, eventType event.Type) int {
	return store.GetPowerLevels(roomID).GetEventLevel(eventType)
}

func (store *BasicStateStore) HasPowerLevel(roomID id.RoomID, userID id.UserID, eventType event.Type) bool {
	return store.GetPowerLevel(roomID, userID) >= store.GetPowerLevelRequirement(roomID, eventType)
}