client_ips.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  1. # -*- coding: utf-8 -*-
  2. # Copyright 2016 OpenMarket Ltd
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. import logging
  16. from twisted.internet import defer
  17. from ._base import Cache
  18. from . import background_updates
  19. logger = logging.getLogger(__name__)
  20. # Number of msec of granularity to store the user IP 'last seen' time. Smaller
  21. # times give more inserts into the database even for readonly API hits
  22. # 120 seconds == 2 minutes
  23. LAST_SEEN_GRANULARITY = 120 * 1000
  24. class ClientIpStore(background_updates.BackgroundUpdateStore):
  25. def __init__(self, hs):
  26. self.client_ip_last_seen = Cache(
  27. name="client_ip_last_seen",
  28. keylen=4,
  29. )
  30. super(ClientIpStore, self).__init__(hs)
  31. self.register_background_index_update(
  32. "user_ips_device_index",
  33. index_name="user_ips_device_id",
  34. table="user_ips",
  35. columns=["user_id", "device_id", "last_seen"],
  36. )
  37. @defer.inlineCallbacks
  38. def insert_client_ip(self, user, access_token, ip, user_agent, device_id):
  39. now = int(self._clock.time_msec())
  40. key = (user.to_string(), access_token, ip)
  41. try:
  42. last_seen = self.client_ip_last_seen.get(key)
  43. except KeyError:
  44. last_seen = None
  45. # Rate-limited inserts
  46. if last_seen is not None and (now - last_seen) < LAST_SEEN_GRANULARITY:
  47. defer.returnValue(None)
  48. self.client_ip_last_seen.prefill(key, now)
  49. # It's safe not to lock here: a) no unique constraint,
  50. # b) LAST_SEEN_GRANULARITY makes concurrent updates incredibly unlikely
  51. yield self._simple_upsert(
  52. "user_ips",
  53. keyvalues={
  54. "user_id": user.to_string(),
  55. "access_token": access_token,
  56. "ip": ip,
  57. "user_agent": user_agent,
  58. "device_id": device_id,
  59. },
  60. values={
  61. "last_seen": now,
  62. },
  63. desc="insert_client_ip",
  64. lock=False,
  65. )
  66. @defer.inlineCallbacks
  67. def get_last_client_ip_by_device(self, devices):
  68. """For each device_id listed, give the user_ip it was last seen on
  69. Args:
  70. devices (iterable[(str, str)]): list of (user_id, device_id) pairs
  71. Returns:
  72. defer.Deferred: resolves to a dict, where the keys
  73. are (user_id, device_id) tuples. The values are also dicts, with
  74. keys giving the column names
  75. """
  76. res = yield self.runInteraction(
  77. "get_last_client_ip_by_device",
  78. self._get_last_client_ip_by_device_txn,
  79. retcols=(
  80. "user_id",
  81. "access_token",
  82. "ip",
  83. "user_agent",
  84. "device_id",
  85. "last_seen",
  86. ),
  87. devices=devices
  88. )
  89. ret = {(d["user_id"], d["device_id"]): d for d in res}
  90. defer.returnValue(ret)
  91. @classmethod
  92. def _get_last_client_ip_by_device_txn(cls, txn, devices, retcols):
  93. where_clauses = []
  94. bindings = []
  95. for (user_id, device_id) in devices:
  96. if device_id is None:
  97. where_clauses.append("(user_id = ? AND device_id IS NULL)")
  98. bindings.extend((user_id, ))
  99. else:
  100. where_clauses.append("(user_id = ? AND device_id = ?)")
  101. bindings.extend((user_id, device_id))
  102. inner_select = (
  103. "SELECT MAX(last_seen) mls, user_id, device_id FROM user_ips "
  104. "WHERE %(where)s "
  105. "GROUP BY user_id, device_id"
  106. ) % {
  107. "where": " OR ".join(where_clauses),
  108. }
  109. sql = (
  110. "SELECT %(retcols)s FROM user_ips "
  111. "JOIN (%(inner_select)s) ips ON"
  112. " user_ips.last_seen = ips.mls AND"
  113. " user_ips.user_id = ips.user_id AND"
  114. " (user_ips.device_id = ips.device_id OR"
  115. " (user_ips.device_id IS NULL AND ips.device_id IS NULL)"
  116. " )"
  117. ) % {
  118. "retcols": ",".join("user_ips." + c for c in retcols),
  119. "inner_select": inner_select,
  120. }
  121. txn.execute(sql, bindings)
  122. return cls.cursor_to_dict(txn)