From 7fba7b946ec0d382de9ef13fa10f6cdfb779a63e Mon Sep 17 00:00:00 2001
From: Shivam Mishra <scm.mymail@gmail.com>
Date: Tue, 13 Jan 2026 18:22:35 +0530
Subject: [PATCH] fix(captain): add operation whitelist to prevent method
 injection

- Add ALLOWED_OPERATIONS whitelist to validate operation parameter
- Prevent arbitrary method invocation via send()
- Refactor tone methods using metaprogramming to reduce duplication
- Add specs for invalid operations and injection attempts
---
 lib/captain/rewrite_service.rb           | 34 +++++++++---------------
 spec/lib/captain/rewrite_service_spec.rb | 24 +++++++++++++++++
 2 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/lib/captain/rewrite_service.rb b/lib/captain/rewrite_service.rb
index 653b7465f..3a217d3c6 100644
--- a/lib/captain/rewrite_service.rb
+++ b/lib/captain/rewrite_service.rb
@@ -1,8 +1,20 @@
 class Captain::RewriteService < Captain::BaseTaskService
   pattr_initialize [:account!, :content!, :operation!, { conversation_display_id: nil }]
 
+  TONE_OPERATIONS = %i[casual professional friendly confident straightforward].freeze
+  ALLOWED_OPERATIONS = (%i[fix_spelling_grammar improve] + TONE_OPERATIONS).freeze
+
   def perform
-    send(operation)
+    operation_sym = operation.to_sym
+    raise ArgumentError, "Invalid operation: #{operation}" unless ALLOWED_OPERATIONS.include?(operation_sym)
+
+    send(operation_sym)
+  end
+
+  TONE_OPERATIONS.each do |tone|
+    define_method(tone) do
+      call_llm_with_prompt(tone_rewrite_prompt(tone.to_s))
+    end
   end
 
   private
@@ -11,26 +23,6 @@ class Captain::RewriteService < Captain::BaseTaskService
     call_llm_with_prompt(prompt_from_file('fix_spelling_grammar'))
   end
 
-  def casual
-    call_llm_with_prompt(tone_rewrite_prompt('casual'))
-  end
-
-  def professional
-    call_llm_with_prompt(tone_rewrite_prompt('professional'))
-  end
-
-  def friendly
-    call_llm_with_prompt(tone_rewrite_prompt('friendly'))
-  end
-
-  def confident
-    call_llm_with_prompt(tone_rewrite_prompt('confident'))
-  end
-
-  def straightforward
-    call_llm_with_prompt(tone_rewrite_prompt('straightforward'))
-  end
-
   def improve
     template = prompt_from_file('improve')
 
diff --git a/spec/lib/captain/rewrite_service_spec.rb b/spec/lib/captain/rewrite_service_spec.rb
index cf53fc8f1..d7900d3cb 100644
--- a/spec/lib/captain/rewrite_service_spec.rb
+++ b/spec/lib/captain/rewrite_service_spec.rb
@@ -139,4 +139,28 @@ RSpec.describe Captain::RewriteService do
       expect(result[:message]).to eq('Rewritten text')
     end
   end
+
+  describe '#perform with invalid operation' do
+    it 'raises ArgumentError for unknown operation' do
+      invalid_service = described_class.new(
+        account: account,
+        content: content,
+        operation: 'invalid_operation',
+        conversation_display_id: conversation.display_id
+      )
+
+      expect { invalid_service.perform }.to raise_error(ArgumentError, /Invalid operation/)
+    end
+
+    it 'prevents method injection attacks' do
+      dangerous_service = described_class.new(
+        account: account,
+        content: content,
+        operation: 'perform',
+        conversation_display_id: conversation.display_id
+      )
+
+      expect { dangerous_service.perform }.to raise_error(ArgumentError, /Invalid operation/)
+    end
+  end
 end