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,9 @@
 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
+
 	def initialize(collection)
 		@collection = collection
 		@_headers = []
@@ -15,17 +18,21 @@ def initialize(collection)
 	attr_reader :collection
 
 	def call(buffer = +"", view_context: nil)
+		unless prevent_csv_injection? == true || prevent_csv_injection? == false
+			raise "You must define `prevent_csv_injection?` on #{self.class.inspect}, returning `true` or `false`. See https://owasp.org/www-community/attacks/CSV_Injection"
+		end
+
 		@_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 +55,56 @@ 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 prevent_csv_injection?
+		nil
 	end
 
-	def escape(value)
-		value = value.to_s
+	def helpers
+		@_view_context
+	end
 
-		if value.include?('"') || value.include?(",") || value.include?("\n")
+	def escape(value)
+		value = trim_whitespace? ? value.to_s.strip : value.to_s
+		first_char = value[0]
+		last_char = value[-1]
+
+		if prevent_csv_injection? && FORMULA_PREFIXES[first_char]
+			# Prefix a single quote to prevent Excel, Google Docs, etc. from interpreting the value as a formula.
+			# See https://owasp.org/www-community/attacks/CSV_Injection
+			%("'#{value.gsub('"', '""')}")
+		elsif (!trim_whitespace? && (SPACE_CHARACTERS[first_char] || SPACE_CHARACTERS[last_char])) || value.include?('"') || value.include?(",") || value.include?("\n")
 			%("#{value.gsub('"', '""')}")
 		else
 			value

test/phlex/csv.rb

@@ -3,6 +3,10 @@
 Product = Struct.new(:name, :price)
 
 class Example < Phlex::CSV
+	def prevent_csv_injection?
+		true
+	end
+
 	def view_template(product)
 		column("name", product.name)
 		column("price", product.price)