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:
@@ -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')
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user