diff --git a/sgx_jmp.rb b/sgx_jmp.rb index 5cc1d1da03d39d9f53a8be967e623f80c9797053..2025184bcea69dcf1efc13d9100286060b035a11 100644 --- a/sgx_jmp.rb +++ b/sgx_jmp.rb @@ -311,13 +311,21 @@ def billable_message(m) (m.body && !m.body.empty?) || m.find("ns:x", ns: OOB.registered_ns).first end -def notify_admin_of_usage(customer, usage, today) - ExpiringLock.new("jmp_usage_notify-#{customer.customer_id}").with do - BLATHER.join(CONFIG[:notify_admin], "sgx-jmp") - BLATHER.say( - CONFIG[:notify_admin], "#{customer.customer_id} has used " \ - "#{usage} messages since #{today - 30}", :groupchat - ) +class OverLimit < StandardError + def initialize(customer, usage) + super("Please contact support") + @customer = customer + @usage = usage + end + + def notify_admin + ExpiringLock.new("jmp_usage_notify-#{@customer.customer_id}").with do + BLATHER.join(CONFIG[:notify_admin], "sgx-jmp") + BLATHER.say( + CONFIG[:notify_admin], "#{@customer.customer_id} has used " \ + "#{@usage} messages today", :groupchat + ) + end end end @@ -329,19 +337,19 @@ message do |m| CustomerRepo .new(set_user: sentry_hub.current_scope.method(:set_user)) .find_by_jid(m.from.stripped).then { |customer| - EMPromise.all([ - (customer.incr_message_usage if billable_message(m)), - customer.message_usage((today..(today - 30))).then do |usage| - if usage < 4500 - customer.stanza_from(m) - else - BLATHER << m.as_error( - "policy-violation", :wait, "Please contact support" - ) - end - notify_admin_of_usage(customer, usage, today) if usage > 900 - end - ]) + next customer.stanza_from(m) unless billable_message(m) + + customer.message_usage((today..today)).then { |usage| + raise OverLimit.new(customer, usage) if usage > 500 + }.then do + EMPromise.all([ + customer.incr_message_usage, + customer.stanza_from(m) + ]) + end + }.catch_only(OverLimit) { |e| + e.notify_admin + BLATHER << m.as_error("policy-violation", :wait, e.message) }.catch_only(CustomerRepo::NotFound) { |e| BLATHER << m.as_error("forbidden", :auth, e.message) }.catch { |e| panic(e, sentry_hub) }