Prevent CSV injection attacks

.rubocop.yml

@@ -22,3 +22,6 @@ Naming/AsciiIdentifiers:
 
 Naming/MethodName:
   Enabled: false
+
+Style/ReturnNilInPredicateMethodDefinition:
+  Enabled: false

lib/phlex/csv.rb

@@ -3,6 +3,25 @@
 class Phlex::CSV
 	include Phlex::Callable
 
+	FORMULA_PREFIXES = ["=", "+", "-", "@", "\t", "\r"].to_h { |prefix| [prefix, true] }.freeze
+	SPACE_CHARACTERS = [" ", "\t", "\r"].to_h { |char| [char, true] }.freeze
+
+	module InjectionPolicy
+		Raise = lambda do |value|
+			raise "CSV injection detected: #{value.inspect}"
+		end
+
+		Ignore = lambda do |value|
+			value
+		end
+
+		Escape = lambda do |value|
+			%("'#{value.gsub('"', '""')}")
+		end
+	end
+
+	include InjectionPolicy
+
 	def initialize(collection)
 		@collection = collection
 		@_headers = []
@@ -18,14 +37,14 @@ def call(buffer = +"", view_context: nil)
 		@_view_context = view_context
 
 		each_item do |record|
-			collection_yielder(record) do |*args, **kwargs|
+			yielder(record) do |*args, **kwargs|
 				view_template(*args, **kwargs)
 
 				if @_first && render_headers?
-					buffer << @_headers.map! { |value| escape(value) }.join(",") << "\n"
+					buffer << @_headers.join(",") << "\n"
 				end
 
-				buffer << @_current_row.map! { |value| escape(value) }.join(",") << "\n"
+				buffer << @_current_row.join(",") << "\n"
 				@_current_column_index = 0
 				@_current_row.clear
 			end
@@ -48,43 +67,54 @@ def content_type
 
 	def column(header = nil, value)
 		if @_first
-			@_headers << header
+			@_headers << escape(header)
 		elsif header != @_headers[@_current_column_index]
 			raise "Inconsistent header."
 		end
 
-		@_current_row << value
+		@_current_row << escape(value)
 		@_current_column_index += 1
 	end
 
 	def each_item(&block)
 		collection.each(&block)
 	end
 
-	def collection_yielder(record)
+	def yielder(record)
 		yield(record)
 	end
 
 	def template(...)
 		nil
 	end
 
+	# Override and set to `false` to disable rendering headers.
 	def render_headers?
 		true
 	end
 
-	def helpers
-		@_view_context
+	# Override and set to `true` to strip leading and trailing whitespace from values.
+	def trim_whitespace?
+		false
 	end
 
-	def render(renderable)
-		renderable.call(view_context: @_view_context)
+	# Override and set to `false` to disable CSV injection escapes or `true` to enable.
+	def in_case_of_injection
+		Raise
+	end
+
+	def helpers
+		@_view_context
 	end
 
 	def escape(value)
-		value = value.to_s
+		value = trim_whitespace? ? value.to_s.strip : value.to_s
+		first_char = value[0]
+		last_char = value[-1]
 
-		if value.include?('"') || value.include?(",") || value.include?("\n")
+		if InjectionPolicy::Ignore != in_case_of_injection && FORMULA_PREFIXES[first_char]
+			in_case_of_injection.call(value)
+		elsif (!trim_whitespace? && (SPACE_CHARACTERS[first_char] || SPACE_CHARACTERS[last_char])) || value.include?('"') || value.include?(",") || value.include?("\n")
 			%("#{value.gsub('"', '""')}")
 		else
 			value