From e6ebac5e13d2fc27ab1744f54b35aff25db4b673 Mon Sep 17 00:00:00 2001
From: Hongbo Wu <hongbo@omnivore.app>
Date: Sun, 18 Aug 2024 11:52:21 +0800
Subject: [PATCH] fix: special handler for youtube shorts

---
 .../src/websites/youtube-handler.ts           | 73 ++++++++++++-------
 1 file changed, 47 insertions(+), 26 deletions(-)

diff --git a/packages/content-handler/src/websites/youtube-handler.ts b/packages/content-handler/src/websites/youtube-handler.ts
index 3779241dd..d83262176 100644
--- a/packages/content-handler/src/websites/youtube-handler.ts
+++ b/packages/content-handler/src/websites/youtube-handler.ts
@@ -1,21 +1,13 @@
-import { ContentHandler, PreHandleResult } from '../content-handler'
 import axios from 'axios'
 import _ from 'underscore'
+import { ContentHandler, PreHandleResult } from '../content-handler'
 
 const YOUTUBE_URL_MATCH =
-  /^((?:https?:)?\/\/)?((?:www|m)\.)?((?:youtube\.com|youtu.be))(\/(?:[\w-]+\?v=|embed\/|v\/)?)([\w-]+)(\S+)?$/
+  /^((?:https?:)?\/\/)?((?:www|m)\.)?((?:youtube\.com|youtu.be))(\/(?:[\w-]+\?v=|embed\/|v\/|shorts\/|playlist\?list=)?)([\w-]+)(\S+)?$/
 
 export const getYoutubeVideoId = (url: string) => {
   const u = new URL(url)
-  const videoId = u.searchParams.get('v')
-  if (!videoId) {
-    const match = url.toString().match(YOUTUBE_URL_MATCH)
-    if (match === null || match.length < 6 || !match[5]) {
-      return undefined
-    }
-    return match[5]
-  }
-  return videoId
+  return u.searchParams.get('v')
 }
 
 export const getYoutubePlaylistId = (url: string) => {
@@ -23,6 +15,41 @@ export const getYoutubePlaylistId = (url: string) => {
   return u.searchParams.get('list')
 }
 
+const getEmbedData = (url: string) => {
+  const BaseUrl = 'https://www.youtube.com'
+  const embedBaseUrl = 'https://www.youtube.com/embed'
+
+  const match = url.match(YOUTUBE_URL_MATCH)
+  if (match === null || match.length < 6) {
+    console.error('Invalid youtube url', url)
+    throw new Error(`Invalid youtube url: ${url}`)
+  }
+
+  const type = match[4]
+  const id = match[5]
+
+  if (type === '/playlist?list=') {
+    const playlistId = getYoutubePlaylistId(url) || id
+
+    return {
+      src: `${embedBaseUrl}/videoseries?list=${playlistId}`,
+      url: `${BaseUrl}/playlist?list=${playlistId}`,
+    }
+  } else if (type === '/shorts/') {
+    return {
+      src: `${embedBaseUrl}/${id}`,
+      url: `${BaseUrl}/shorts/${id}`,
+    }
+  }
+
+  const videoId = getYoutubeVideoId(url) || id
+
+  return {
+    src: `${embedBaseUrl}/${videoId}`,
+    url: `${BaseUrl}/watch?v=${videoId}`,
+  }
+}
+
 export const escapeTitle = (title: string) => {
   return _.escape(title)
 }
@@ -38,21 +65,15 @@ export class YoutubeHandler extends ContentHandler {
   }
 
   async preHandle(url: string): Promise<PreHandleResult> {
-    const BaseUrl = 'https://www.youtube.com'
-    const embedBaseUrl = 'https://www.youtube.com/embed'
-    let urlToEncode: string
-    let src: string
-    const playlistId = getYoutubePlaylistId(url)
-    if (playlistId) {
-      urlToEncode = `${BaseUrl}/playlist?list=${playlistId}`
-      src = `${embedBaseUrl}/videoseries?list=${playlistId}`
-    } else {
-      const videoId = getYoutubeVideoId(url)
-      if (!videoId) {
-        return {}
-      }
-      urlToEncode = `${BaseUrl}/watch?v=${videoId}`
-      src = `${embedBaseUrl}/${videoId}`
+    let src, urlToEncode
+
+    try {
+      const embedData = getEmbedData(url)
+      src = embedData.src
+      urlToEncode = embedData.url
+    } catch (error) {
+      console.error('Error getting embed data', error)
+      return {}
     }
 
     const oembedUrl =