Tulir Asokan
237 lines
8.4 KiB
Go
237 lines
8.4 KiB
Go
// mautrix-telegram - A Matrix-Telegram puppeting bridge.
|
|
// Copyright (C) 2025 Sumner Evans
|
|
//
|
|
// This program is free software: you can redistribute it and/or modify
|
|
// it under the terms of the GNU Affero General Public License as published by
|
|
// the Free Software Foundation, either version 3 of the License, or
|
|
// (at your option) any later version.
|
|
//
|
|
// This program is distributed in the hope that it will be useful,
|
|
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
// GNU Affero General Public License for more details.
|
|
//
|
|
// You should have received a copy of the GNU Affero General Public License
|
|
// along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
package store
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"fmt"
|
|
|
|
"go.mau.fi/util/dbutil"
|
|
|
|
"go.mau.fi/mautrix-telegram/pkg/connector/ids"
|
|
"go.mau.fi/mautrix-telegram/pkg/gotd/telegram/updates"
|
|
)
|
|
|
|
// ScopedStore is a wrapper around a database that implements
|
|
// [session.Storage] scoped to a specific Telegram user ID.
|
|
type ScopedStore struct {
|
|
db *dbutil.Database
|
|
telegramUserID int64
|
|
}
|
|
|
|
const (
|
|
// State Storage Queries
|
|
allChannelsQuery = "SELECT channel_id, pts FROM telegram_channel_state WHERE user_id=$1"
|
|
getChannelPtsQuery = "SELECT pts FROM telegram_channel_state WHERE user_id=$1 AND channel_id=$2"
|
|
setChannelPtsQuery = `
|
|
INSERT INTO telegram_channel_state (user_id, channel_id, pts)
|
|
VALUES ($1, $2, $3)
|
|
ON CONFLICT (user_id, channel_id) DO UPDATE SET pts=excluded.pts
|
|
`
|
|
getStateQuery = "SELECT pts, qts, date, seq from telegram_user_state WHERE user_id=$1"
|
|
setStateQuery = `
|
|
INSERT INTO telegram_user_state (user_id, pts, qts, date, seq)
|
|
VALUES ($1, $2, $3, $4, $5)
|
|
ON CONFLICT (user_id) DO UPDATE SET
|
|
pts=excluded.pts,
|
|
qts=excluded.qts,
|
|
date=excluded.date,
|
|
seq=excluded.seq
|
|
`
|
|
setPtsQuery = "UPDATE telegram_user_state SET pts=$1 WHERE user_id=$2"
|
|
setQtsQuery = "UPDATE telegram_user_state SET qts=$1 WHERE user_id=$2"
|
|
setDateQuery = "UPDATE telegram_user_state SET date=$1 WHERE user_id=$2"
|
|
setSeqQuery = "UPDATE telegram_user_state SET seq=$1 WHERE user_id=$2"
|
|
setDateSeqQuery = "UPDATE telegram_user_state SET date=$1, seq=$2 WHERE user_id=$3"
|
|
|
|
deleteChannelStateForUserQuery = "DELETE FROM telegram_channel_state WHERE user_id=$1"
|
|
deleteUserStateForUserQuery = "DELETE FROM telegram_user_state WHERE user_id=$1"
|
|
|
|
getAccessHashQuery = "SELECT access_hash FROM telegram_access_hash WHERE user_id=$1 AND entity_type=$2 AND entity_id=$3"
|
|
setAccessHashQuery = `
|
|
INSERT INTO telegram_access_hash (user_id, entity_type, entity_id, access_hash)
|
|
VALUES ($1, $2, $3, $4)
|
|
ON CONFLICT (user_id, entity_type, entity_id) DO UPDATE SET access_hash=excluded.access_hash
|
|
`
|
|
deleteAccessHashesForUserQuery = "DELETE FROM telegram_access_hash WHERE user_id=$1"
|
|
)
|
|
|
|
var _ updates.StateStorage = (*ScopedStore)(nil)
|
|
|
|
type channelIDPtsTuple struct {
|
|
ChannelID int64
|
|
Pts int
|
|
}
|
|
|
|
var ciptScanner = dbutil.ConvertRowFn[channelIDPtsTuple](func(row dbutil.Scannable) (cipt channelIDPtsTuple, err error) {
|
|
err = row.Scan(&cipt.ChannelID, &cipt.Pts)
|
|
return
|
|
})
|
|
|
|
func (s *ScopedStore) ForEachChannels(ctx context.Context, userID int64, f func(ctx context.Context, channelID int64, pts int) error) error {
|
|
s.assertUserIDMatches(userID)
|
|
items, err := ciptScanner.NewRowIter(s.db.Query(ctx, allChannelsQuery, userID)).AsList()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, item := range items {
|
|
err = f(ctx, item.ChannelID, item.Pts)
|
|
if err != nil {
|
|
return fmt.Errorf("iteration error for channel %d: %w", item.ChannelID, err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *ScopedStore) GetChannelPts(ctx context.Context, userID int64, channelID int64) (pts int, found bool, err error) {
|
|
s.assertUserIDMatches(userID)
|
|
err = s.db.QueryRow(ctx, getChannelPtsQuery, userID, channelID).Scan(&pts)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return 0, false, nil
|
|
}
|
|
return pts, err == nil, err
|
|
}
|
|
|
|
func (s *ScopedStore) SetChannelPts(ctx context.Context, userID int64, channelID int64, pts int) (err error) {
|
|
s.assertUserIDMatches(userID)
|
|
_, err = s.db.Exec(ctx, setChannelPtsQuery, userID, channelID, pts)
|
|
return
|
|
}
|
|
|
|
func (s *ScopedStore) GetState(ctx context.Context, userID int64) (state updates.State, found bool, err error) {
|
|
s.assertUserIDMatches(userID)
|
|
err = s.db.QueryRow(ctx, getStateQuery, userID).Scan(&state.Pts, &state.Qts, &state.Date, &state.Seq)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return state, false, nil
|
|
}
|
|
return state, err == nil, err
|
|
}
|
|
|
|
func (s *ScopedStore) SetState(ctx context.Context, userID int64, state updates.State) (err error) {
|
|
s.assertUserIDMatches(userID)
|
|
_, err = s.db.Exec(ctx, setStateQuery, userID, state.Pts, state.Qts, state.Date, state.Seq)
|
|
return
|
|
}
|
|
|
|
func (s *ScopedStore) SetPts(ctx context.Context, userID int64, pts int) (err error) {
|
|
s.assertUserIDMatches(userID)
|
|
_, err = s.db.Exec(ctx, setPtsQuery, userID, pts)
|
|
return
|
|
}
|
|
|
|
func (s *ScopedStore) SetQts(ctx context.Context, userID int64, qts int) (err error) {
|
|
s.assertUserIDMatches(userID)
|
|
_, err = s.db.Exec(ctx, setQtsQuery, userID, qts)
|
|
return
|
|
}
|
|
|
|
func (s *ScopedStore) SetSeq(ctx context.Context, userID int64, seq int) (err error) {
|
|
s.assertUserIDMatches(userID)
|
|
_, err = s.db.Exec(ctx, setSeqQuery, userID, seq)
|
|
return
|
|
}
|
|
|
|
func (s *ScopedStore) SetDate(ctx context.Context, userID int64, date int) (err error) {
|
|
s.assertUserIDMatches(userID)
|
|
_, err = s.db.Exec(ctx, setDateQuery, userID, date)
|
|
return
|
|
}
|
|
|
|
func (s *ScopedStore) SetDateSeq(ctx context.Context, userID int64, date int, seq int) (err error) {
|
|
s.assertUserIDMatches(userID)
|
|
_, err = s.db.Exec(ctx, setDateSeqQuery, userID, date, seq)
|
|
return
|
|
}
|
|
|
|
func (s *ScopedStore) DeleteUserState(ctx context.Context) (err error) {
|
|
_, err = s.db.Exec(ctx, deleteUserStateForUserQuery, s.telegramUserID)
|
|
return
|
|
}
|
|
|
|
func (s *ScopedStore) DeleteChannelStateForUser(ctx context.Context) (err error) {
|
|
_, err = s.db.Exec(ctx, deleteChannelStateForUserQuery, s.telegramUserID)
|
|
return
|
|
}
|
|
|
|
var _ updates.AccessHasher = (*ScopedStore)(nil)
|
|
|
|
// Deprecated: only for interface, don't use directly. Use [GetAccessHash]
|
|
// instead.
|
|
func (s *ScopedStore) GetChannelAccessHash(ctx context.Context, forUserID, channelID int64) (int64, bool, error) {
|
|
s.assertUserIDMatches(forUserID)
|
|
accessHash, err := s.GetAccessHash(ctx, ids.PeerTypeChannel, channelID)
|
|
if errors.Is(err, ErrNoAccessHash) {
|
|
return 0, false, nil
|
|
}
|
|
return accessHash, true, err
|
|
}
|
|
|
|
// Deprecated: only for interface, don't use directly. Use [SetAccessHash]
|
|
// instead.
|
|
func (s *ScopedStore) SetChannelAccessHash(ctx context.Context, forUserID, channelID, accessHash int64) (err error) {
|
|
s.assertUserIDMatches(forUserID)
|
|
return s.SetAccessHash(ctx, ids.PeerTypeChannel, channelID, accessHash)
|
|
}
|
|
|
|
// Deprecated: only for interface, don't use directly. Use [GetAccessHash]
|
|
// instead.
|
|
func (s *ScopedStore) GetUserAccessHash(ctx context.Context, forUserID int64, userID int64) (int64, bool, error) {
|
|
s.assertUserIDMatches(forUserID)
|
|
accessHash, err := s.GetAccessHash(ctx, ids.PeerTypeUser, userID)
|
|
if errors.Is(err, ErrNoAccessHash) {
|
|
return 0, false, nil
|
|
}
|
|
return accessHash, true, err
|
|
}
|
|
|
|
// Deprecated: only for interface, don't use directly. Use [SetAccessHash]
|
|
// instead.
|
|
func (s *ScopedStore) SetUserAccessHash(ctx context.Context, forUserID int64, userID int64, accessHash int64) error {
|
|
s.assertUserIDMatches(forUserID)
|
|
return s.SetAccessHash(ctx, ids.PeerTypeUser, userID, accessHash)
|
|
}
|
|
|
|
var ErrNoAccessHash = errors.New("access hash not found")
|
|
|
|
func (s *ScopedStore) GetAccessHash(ctx context.Context, entityType ids.PeerType, entityID int64) (accessHash int64, err error) {
|
|
err = s.db.QueryRow(ctx, getAccessHashQuery, s.telegramUserID, entityType, entityID).Scan(&accessHash)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
err = fmt.Errorf("%w for %d", ErrNoAccessHash, entityID)
|
|
}
|
|
return
|
|
}
|
|
|
|
func (s *ScopedStore) SetAccessHash(ctx context.Context, entityType ids.PeerType, entityID, accessHash int64) (err error) {
|
|
_, err = s.db.Exec(ctx, setAccessHashQuery, s.telegramUserID, entityType, entityID, accessHash)
|
|
return
|
|
}
|
|
|
|
func (s *ScopedStore) DeleteAccessHashesForUser(ctx context.Context) (err error) {
|
|
_, err = s.db.Exec(ctx, deleteAccessHashesForUserQuery, s.telegramUserID)
|
|
return
|
|
}
|
|
|
|
// Helper Functions
|
|
|
|
func (s *ScopedStore) assertUserIDMatches(userID int64) {
|
|
if s.telegramUserID != userID {
|
|
panic(fmt.Sprintf("scoped store for %d function called with user ID %d", s.telegramUserID, userID))
|
|
}
|
|
}
|