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
This commit is contained in:
Shivam Mishra
2026-01-13 18:22:35 +05:30
parent f63cef20bf
commit 7fba7b946e
2 changed files with 37 additions and 21 deletions
+13 -21
View File
@@ -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')
+24
View File
@@ -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