From 0ad2eff2a46c28a877de1960a2dc5c15c0f94444 Mon Sep 17 00:00:00 2001
From: syeopite <syeopite@syeopite.dev>
Date: Tue, 30 Jan 2024 15:25:45 -0800
Subject: [PATCH] WebVTT::Builder: Add logic to escape special chars

---
 spec/helpers/vtt/builder_spec.cr | 65 +++++++++++++++++++++-----------
 src/invidious/helpers/webvtt.cr  | 16 +++++++-
 2 files changed, 59 insertions(+), 22 deletions(-)

diff --git a/spec/helpers/vtt/builder_spec.cr b/spec/helpers/vtt/builder_spec.cr
index 7b543ddc..dc1f4613 100644
--- a/spec/helpers/vtt/builder_spec.cr
+++ b/spec/helpers/vtt/builder_spec.cr
@@ -1,34 +1,27 @@
 require "../../spec_helper.cr"
 
-MockLines = [
-  {
-    "start_time": Time::Span.new(seconds: 1),
-    "end_time":   Time::Span.new(seconds: 2),
-    "text":       "Line 1",
-  },
-
-  {
-    "start_time": Time::Span.new(seconds: 2),
-    "end_time":   Time::Span.new(seconds: 3),
-    "text":       "Line 2",
-  },
-]
+MockLines                       = ["Line 1", "Line 2"]
+MockLinesWithEscapableCharacter = ["<Line 1>", "&Line 2>", '\u200E' + "Line\u200F 3", "\u00A0Line 4"]
 
 Spectator.describe "WebVTT::Builder" do
   it "correctly builds a vtt file" do
     result = WebVTT.build do |vtt|
-      MockLines.each do |line|
-        vtt.cue(line["start_time"], line["end_time"], line["text"])
+      2.times do |i|
+        vtt.cue(
+          Time::Span.new(seconds: i),
+          Time::Span.new(seconds: i + 1),
+          MockLines[i]
+        )
       end
     end
 
     expect(result).to eq([
       "WEBVTT",
       "",
-      "00:00:01.000 --> 00:00:02.000",
+      "00:00:00.000 --> 00:00:01.000",
       "Line 1",
       "",
-      "00:00:02.000 --> 00:00:03.000",
+      "00:00:01.000 --> 00:00:02.000",
       "Line 2",
       "",
       "",
@@ -42,8 +35,12 @@ Spectator.describe "WebVTT::Builder" do
     }
 
     result = WebVTT.build(setting_fields) do |vtt|
-      MockLines.each do |line|
-        vtt.cue(line["start_time"], line["end_time"], line["text"])
+      2.times do |i|
+        vtt.cue(
+          Time::Span.new(seconds: i),
+          Time::Span.new(seconds: i + 1),
+          MockLines[i]
+        )
       end
     end
 
@@ -52,13 +49,39 @@ Spectator.describe "WebVTT::Builder" do
       "Kind: captions",
       "Language: en",
       "",
-      "00:00:01.000 --> 00:00:02.000",
+      "00:00:00.000 --> 00:00:01.000",
       "Line 1",
       "",
-      "00:00:02.000 --> 00:00:03.000",
+      "00:00:01.000 --> 00:00:02.000",
       "Line 2",
       "",
       "",
     ].join('\n'))
   end
+
+  it "properly escapes characters" do
+    result = WebVTT.build do |vtt|
+      4.times do |i|
+        vtt.cue(Time::Span.new(seconds: i), Time::Span.new(seconds: i + 1), MockLinesWithEscapableCharacter[i])
+      end
+    end
+
+    expect(result).to eq([
+      "WEBVTT",
+      "",
+      "00:00:00.000 --> 00:00:01.000",
+      "&lt;Line 1&gt;",
+      "",
+      "00:00:01.000 --> 00:00:02.000",
+      "&amp;Line 2&gt;",
+      "",
+      "00:00:02.000 --> 00:00:03.000",
+      "&lrm;Line&rlm; 3",
+      "",
+      "00:00:03.000 --> 00:00:04.000",
+      "&nbsp;Line 4",
+      "",
+      "",
+    ].join('\n'))
+  end
 end
diff --git a/src/invidious/helpers/webvtt.cr b/src/invidious/helpers/webvtt.cr
index 56f761ed..260d250f 100644
--- a/src/invidious/helpers/webvtt.cr
+++ b/src/invidious/helpers/webvtt.cr
@@ -4,13 +4,23 @@
 module WebVTT
   # A WebVTT builder generates WebVTT files
   private class Builder
+    # See https://developer.mozilla.org/en-US/docs/Web/API/WebVTT_API#cue_payload
+    private ESCAPE_SUBSTITUTIONS = {
+      '&'      => "&amp;",
+      '<'      => "&lt;",
+      '>'      => "&gt;",
+      '\u200E' => "&lrm;",
+      '\u200F' => "&rlm;",
+      '\u00A0' => "&nbsp;",
+    }
+
     def initialize(@io : IO)
     end
 
     # Writes an vtt cue with the specified time stamp and contents
     def cue(start_time : Time::Span, end_time : Time::Span, text : String)
       timestamp(start_time, end_time)
-      @io << text
+      @io << self.escape(text)
       @io << "\n\n"
     end
 
@@ -29,6 +39,10 @@ module WebVTT
       @io << '.' << timestamp.milliseconds.to_s.rjust(3, '0')
     end
 
+    private def escape(text : String) : String
+      return text.gsub(ESCAPE_SUBSTITUTIONS)
+    end
+
     def document(setting_fields : Hash(String, String)? = nil, &)
       @io << "WEBVTT\n"