From fb085cbb36d49989c01873050ae14a498b103b0f Mon Sep 17 00:00:00 2001
From: Daniel Morrison <daniel@collectiveidea.com>
Date: Fri, 4 Oct 2024 10:48:27 -0400
Subject: [PATCH] Allow handlers to use rescue_from in a way familiar to
 ActionController users

ActiveSupport::Rescuable does close to what we need, but rather than calling render or redirect, handlers return an object (which may be a Twirp::Error). We refine ActiveSupport::Rescuable to return the result of the rescue_from block/method.
---
 lib/twirp/rails.rb                            |  1 +
 lib/twirp/rails/handler.rb                    |  8 ++++-
 lib/twirp/rails/rescuable.rb                  | 31 +++++++++++++++++++
 .../app/handlers/haberdasher_handler.rb       |  4 +++
 spec/requests/haberdasher_spec.rb             | 19 ++++++++++++
 5 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 lib/twirp/rails/rescuable.rb

diff --git a/lib/twirp/rails.rb b/lib/twirp/rails.rb
index 00ca11f..ac52d0d 100644
--- a/lib/twirp/rails.rb
+++ b/lib/twirp/rails.rb
@@ -15,6 +15,7 @@ class Error < StandardError; end
 require_relative "rails/configuration"
 require_relative "rails/dispatcher"
 require_relative "rails/engine"
+require_relative "rails/rescuable"
 require_relative "rails/handler"
 
 module Twirp
diff --git a/lib/twirp/rails/handler.rb b/lib/twirp/rails/handler.rb
index 4ed90da..089f574 100644
--- a/lib/twirp/rails/handler.rb
+++ b/lib/twirp/rails/handler.rb
@@ -4,6 +4,8 @@ module Twirp
   module Rails
     class Handler
       include Twirp::Rails::Callbacks
+      include ActiveSupport::Rescuable
+      using Twirp::Rails::Rescuable
 
       attr_reader :request, :env
       attr_reader :action_name
@@ -34,7 +36,11 @@ def process_action(name)
         ActiveSupport::Notifications.instrument("handler_run_callbacks.twirp_rails", handler: self.class.name, action: action_name, env: @env, request: @request) do
           run_callbacks(:process_action) do
             ActiveSupport::Notifications.instrument("handler_run.twirp_rails", handler: self.class.name, action: action_name, env: @env, request: @request) do |payload|
-              payload[:response] = send_action(name)
+              payload[:response] = begin
+                send_action(name)
+              rescue => exception
+                rescue_with_handler_and_return(exception) || raise
+              end
             end
           end
         end
diff --git a/lib/twirp/rails/rescuable.rb b/lib/twirp/rails/rescuable.rb
new file mode 100644
index 0000000..9fe076b
--- /dev/null
+++ b/lib/twirp/rails/rescuable.rb
@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Twirp
+  module Rails
+    module Rescuable
+      refine ::ActiveSupport::Rescuable::ClassMethods do
+        # A slightly altered version of ActiveSupport::Rescuable#rescue_with_handler
+        # that returns the result rather than the handled exception
+        def rescue_with_handler_and_return(exception, object: self, visited_exceptions: [])
+          visited_exceptions << exception
+
+          if (handler = handler_for_rescue(exception, object: object))
+            handler.call exception
+          elsif exception
+            if visited_exceptions.include?(exception.cause)
+              nil
+            else
+              rescue_with_handler(exception.cause, object: object, visited_exceptions: visited_exceptions)
+            end
+          end
+        end
+      end
+
+      refine ::ActiveSupport::Rescuable do
+        def rescue_with_handler_and_return(exception)
+          self.class.rescue_with_handler_and_return exception, object: self
+        end
+      end
+    end
+  end
+end
diff --git a/spec/rails_app/app/handlers/haberdasher_handler.rb b/spec/rails_app/app/handlers/haberdasher_handler.rb
index 948372e..89222f3 100644
--- a/spec/rails_app/app/handlers/haberdasher_handler.rb
+++ b/spec/rails_app/app/handlers/haberdasher_handler.rb
@@ -3,6 +3,10 @@ class HaberdasherHandler < Twirp::Rails::Handler
   before_action :track_request_ip
   before_action :reject_giant_hats
 
+  rescue_from "ArgumentError" do |error|
+    Twirp::Error.invalid_argument(error.message)
+  end
+
   def make_hat
     # We can return a Twirp::Error when appropriate
     if request.inches < 12
diff --git a/spec/requests/haberdasher_spec.rb b/spec/requests/haberdasher_spec.rb
index 7ae32b2..1060786 100644
--- a/spec/requests/haberdasher_spec.rb
+++ b/spec/requests/haberdasher_spec.rb
@@ -145,4 +145,23 @@ def make_hat_success_request
       expect(response.status).to eq(304)
     end
   end
+
+  describe "Rescuable" do
+    it "rescues from ArgumentError" do
+      size = Twirp::Example::Haberdasher::Size.new(inches: 100)
+
+      # Fake an exception
+      expect(Twirp::Example::Haberdasher::Hat).to receive(:new).and_raise(ArgumentError.new("is way too large"))
+
+      post "/twirp/twirp.example.haberdasher.Haberdasher/MakeHat",
+        params: size.to_proto, headers: {
+          :accept => "application/protobuf",
+          "Content-Type" => "application/protobuf"
+        }
+
+      expect(response.status).to eq(400)
+      expect(response.content_type).to eq("application/json")
+      expect(response.body).to eq('{"code":"invalid_argument","msg":"is way too large"}')
+    end
+  end
 end