portal.go 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215
  1. // mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
  2. // Copyright (C) 2021 Tulir Asokan
  3. //
  4. // This program is free software: you can redistribute it and/or modify
  5. // it under the terms of the GNU Affero General Public License as published by
  6. // the Free Software Foundation, either version 3 of the License, or
  7. // (at your option) any later version.
  8. //
  9. // This program is distributed in the hope that it will be useful,
  10. // but WITHOUT ANY WARRANTY; without even the implied warranty of
  11. // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  12. // GNU Affero General Public License for more details.
  13. //
  14. // You should have received a copy of the GNU Affero General Public License
  15. // along with this program. If not, see <https://www.gnu.org/licenses/>.
  16. package database
  17. import (
  18. "database/sql"
  19. "fmt"
  20. log "maunium.net/go/maulogger/v2"
  21. "maunium.net/go/mautrix/id"
  22. "maunium.net/go/mautrix/util/dbutil"
  23. "go.mau.fi/whatsmeow/types"
  24. )
  25. type PortalKey struct {
  26. JID types.JID
  27. Receiver types.JID
  28. }
  29. func NewPortalKey(jid, receiver types.JID) PortalKey {
  30. if jid.Server == types.GroupServer {
  31. receiver = jid
  32. } else if jid.Server == types.LegacyUserServer {
  33. jid.Server = types.DefaultUserServer
  34. }
  35. return PortalKey{
  36. JID: jid.ToNonAD(),
  37. Receiver: receiver.ToNonAD(),
  38. }
  39. }
  40. func (key PortalKey) String() string {
  41. if key.Receiver == key.JID {
  42. return key.JID.String()
  43. }
  44. return key.JID.String() + "-" + key.Receiver.String()
  45. }
  46. type PortalQuery struct {
  47. db *Database
  48. log log.Logger
  49. }
  50. func (pq *PortalQuery) New() *Portal {
  51. return &Portal{
  52. db: pq.db,
  53. log: pq.log,
  54. }
  55. }
  56. const portalColumns = "jid, receiver, mxid, name, topic, avatar, avatar_url, encrypted, first_event_id, next_batch_id, relay_user_id, expiration_time"
  57. func (pq *PortalQuery) GetAll() []*Portal {
  58. return pq.getAll(fmt.Sprintf("SELECT %s FROM portal", portalColumns))
  59. }
  60. func (pq *PortalQuery) GetByJID(key PortalKey) *Portal {
  61. return pq.get(fmt.Sprintf("SELECT %s FROM portal WHERE jid=$1 AND receiver=$2", portalColumns), key.JID, key.Receiver)
  62. }
  63. func (pq *PortalQuery) GetByMXID(mxid id.RoomID) *Portal {
  64. return pq.get(fmt.Sprintf("SELECT %s FROM portal WHERE mxid=$1", portalColumns), mxid)
  65. }
  66. func (pq *PortalQuery) GetAllByJID(jid types.JID) []*Portal {
  67. return pq.getAll(fmt.Sprintf("SELECT %s FROM portal WHERE jid=$1", portalColumns), jid.ToNonAD())
  68. }
  69. func (pq *PortalQuery) FindPrivateChats(receiver types.JID) []*Portal {
  70. return pq.getAll(fmt.Sprintf("SELECT %s FROM portal WHERE receiver=$1 AND jid LIKE '%@s.whatsapp.net'", portalColumns), receiver.ToNonAD())
  71. }
  72. func (pq *PortalQuery) FindPrivateChatsNotInSpace(receiver types.JID) (keys []PortalKey) {
  73. receiver = receiver.ToNonAD()
  74. rows, err := pq.db.Query(`
  75. SELECT jid FROM portal
  76. LEFT JOIN user_portal ON portal.jid=user_portal.portal_jid AND portal.receiver=user_portal.portal_receiver
  77. WHERE mxid<>'' AND receiver=$1 AND (in_space=false OR in_space IS NULL)
  78. `, receiver)
  79. if err != nil || rows == nil {
  80. return
  81. }
  82. for rows.Next() {
  83. var key PortalKey
  84. key.Receiver = receiver
  85. err = rows.Scan(&key.JID)
  86. if err == nil {
  87. keys = append(keys, key)
  88. }
  89. }
  90. return
  91. }
  92. func (pq *PortalQuery) getAll(query string, args ...interface{}) (portals []*Portal) {
  93. rows, err := pq.db.Query(query, args...)
  94. if err != nil || rows == nil {
  95. return nil
  96. }
  97. defer rows.Close()
  98. for rows.Next() {
  99. portals = append(portals, pq.New().Scan(rows))
  100. }
  101. return
  102. }
  103. func (pq *PortalQuery) get(query string, args ...interface{}) *Portal {
  104. row := pq.db.QueryRow(query, args...)
  105. if row == nil {
  106. return nil
  107. }
  108. return pq.New().Scan(row)
  109. }
  110. type Portal struct {
  111. db *Database
  112. log log.Logger
  113. Key PortalKey
  114. MXID id.RoomID
  115. Name string
  116. Topic string
  117. Avatar string
  118. AvatarURL id.ContentURI
  119. Encrypted bool
  120. FirstEventID id.EventID
  121. NextBatchID id.BatchID
  122. RelayUserID id.UserID
  123. ExpirationTime uint32
  124. }
  125. func (portal *Portal) Scan(row dbutil.Scannable) *Portal {
  126. var mxid, avatarURL, firstEventID, nextBatchID, relayUserID sql.NullString
  127. err := row.Scan(&portal.Key.JID, &portal.Key.Receiver, &mxid, &portal.Name, &portal.Topic, &portal.Avatar, &avatarURL, &portal.Encrypted, &firstEventID, &nextBatchID, &relayUserID, &portal.ExpirationTime)
  128. if err != nil {
  129. if err != sql.ErrNoRows {
  130. portal.log.Errorln("Database scan failed:", err)
  131. }
  132. return nil
  133. }
  134. portal.MXID = id.RoomID(mxid.String)
  135. portal.AvatarURL, _ = id.ParseContentURI(avatarURL.String)
  136. portal.FirstEventID = id.EventID(firstEventID.String)
  137. portal.NextBatchID = id.BatchID(nextBatchID.String)
  138. portal.RelayUserID = id.UserID(relayUserID.String)
  139. return portal
  140. }
  141. func (portal *Portal) mxidPtr() *id.RoomID {
  142. if len(portal.MXID) > 0 {
  143. return &portal.MXID
  144. }
  145. return nil
  146. }
  147. func (portal *Portal) relayUserPtr() *id.UserID {
  148. if len(portal.RelayUserID) > 0 {
  149. return &portal.RelayUserID
  150. }
  151. return nil
  152. }
  153. func (portal *Portal) Insert() {
  154. _, err := portal.db.Exec("INSERT INTO portal (jid, receiver, mxid, name, topic, avatar, avatar_url, encrypted, first_event_id, next_batch_id, relay_user_id, expiration_time) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)",
  155. portal.Key.JID, portal.Key.Receiver, portal.mxidPtr(), portal.Name, portal.Topic, portal.Avatar, portal.AvatarURL.String(), portal.Encrypted, portal.FirstEventID.String(), portal.NextBatchID.String(), portal.relayUserPtr(), portal.ExpirationTime)
  156. if err != nil {
  157. portal.log.Warnfln("Failed to insert %s: %v", portal.Key, err)
  158. }
  159. }
  160. func (portal *Portal) Update(txn *sql.Tx) {
  161. query := `
  162. UPDATE portal
  163. SET mxid=$1, name=$2, topic=$3, avatar=$4, avatar_url=$5, encrypted=$6, first_event_id=$7, next_batch_id=$8, relay_user_id=$9, expiration_time=$10
  164. WHERE jid=$11 AND receiver=$12
  165. `
  166. args := []interface{}{
  167. portal.mxidPtr(), portal.Name, portal.Topic, portal.Avatar, portal.AvatarURL.String(), portal.Encrypted, portal.FirstEventID.String(), portal.NextBatchID.String(), portal.relayUserPtr(), portal.ExpirationTime, portal.Key.JID, portal.Key.Receiver,
  168. }
  169. var err error
  170. if txn != nil {
  171. _, err = txn.Exec(query, args...)
  172. } else {
  173. _, err = portal.db.Exec(query, args...)
  174. }
  175. if err != nil {
  176. portal.log.Warnfln("Failed to update %s: %v", portal.Key, err)
  177. }
  178. }
  179. func (portal *Portal) Delete() {
  180. _, err := portal.db.Exec("DELETE FROM portal WHERE jid=$1 AND receiver=$2", portal.Key.JID, portal.Key.Receiver)
  181. if err != nil {
  182. portal.log.Warnfln("Failed to delete %s: %v", portal.Key, err)
  183. }
  184. }