@@ -1,17 +1,19 @@
 package app
 
 import (
+	"fmt"
 	"log"
 	"net/http"
 	"regexp"
-	"strconv"
 	"text/template"
 
 	g "github.com/birabittoh/gopipe/src/globals"
+	"golang.org/x/time/rate"
 )
 
 const (
 	fmtYouTubeURL = "https://www.youtube.com/watch?v=%s"
+	err404        = "Not Found"
 	err500        = "Internal Server Error"
 )
 
@@ -20,6 +22,17 @@ templates      = template.Must(template.ParseGlob("templates/*.html"))
 	userAgentRegex = regexp.MustCompile(`(?i)bot|facebook|embed|got|firefox\/92|firefox\/38|curl|wget|go-http|yahoo|generator|whatsapp|preview|link|proxy|vkshare|images|analyzer|index|crawl|spider|python|cfnetwork|node`)
 	videoRegex     = regexp.MustCompile(`(?i)^[a-z0-9_-]{11}$`)
 )
+
+func limit(limiter *rate.Limiter, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !limiter.Allow() {
+			status := http.StatusTooManyRequests
+			http.Error(w, http.StatusText(status), status)
+			return
+		}
+		next.ServeHTTP(w, r)
+	})
+}
 
 func indexHandler(w http.ResponseWriter, r *http.Request) {
 	err := templates.ExecuteTemplate(w, "index.html", nil)
@@ -55,12 +68,9 @@ http.Error(w, "Invalid video ID.", http.StatusBadRequest)
 		return
 	}
 
-	formatID, err := strconv.ParseUint(r.PathValue("formatID"), 10, 64)
-	if err != nil {
-		formatID = 0
-	}
+	formatID := getFormatID(r.PathValue("formatID"))
 
-	video, format, err := getVideo(videoID, int(formatID))
+	video, format, err := getVideo(videoID, formatID)
 	if err != nil {
 		http.Error(w, err500, http.StatusInternalServerError)
 		return
@@ -76,9 +86,14 @@ if len(video.Thumbnails) > 0 {
 		thumbnail = video.Thumbnails[len(video.Thumbnails)-1].URL
 	}
 
+	videoURL := format.URL
+	if g.Proxy {
+		videoURL = fmt.Sprintf("/proxy/%s/%d", videoID, formatID)
+	}
+
 	data := map[string]interface{}{
 		"VideoID":     videoID,
-		"VideoURL":    format.URL,
+		"VideoURL":    videoURL,
 		"Uploader":    video.Author,
 		"Title":       video.Title,
 		"Description": video.Description,

@@ -1,45 +1,30 @@
 package app
 
 import (
+	"bytes"
 	"log"
 	"net/http"
 	"os"
-	"strings"
+	"time"
 
 	g "github.com/birabittoh/gopipe/src/globals"
+	"github.com/birabittoh/myks"
 	"github.com/joho/godotenv"
 	"golang.org/x/time/rate"
 )
 
-func limit(limiter *rate.Limiter, next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if !limiter.Allow() {
-			status := http.StatusTooManyRequests
-			http.Error(w, http.StatusText(status), status)
-			return
-		}
-		next.ServeHTTP(w, r)
-	})
-}
-
-func getEnvDefault(key string, def string) string {
-	res := os.Getenv(key)
-	if res == "" {
-		return def
-	}
-	return res
-}
-
-func parseBool(s string) bool {
-	return strings.ToLower(s) == "true" || s == "1"
-}
-
 func Main() {
 	godotenv.Load()
 
 	g.Debug = parseBool(os.Getenv("APP_DEBUG"))
 	if g.Debug {
 		log.Println("Debug mode enabled.")
+	}
+
+	g.Proxy = parseBool(os.Getenv("APP_PROXY"))
+	if g.Proxy {
+		g.PKS = myks.New[bytes.Buffer](3 * time.Minute)
+		log.Println("Proxy mode enabled.")
 	}
 
 	g.Port = getEnvDefault("APP_PORT", "3000")
@@ -61,6 +46,9 @@ r.HandleFunc("GET /watch", videoHandler)
 	r.HandleFunc("GET /shorts/{videoID}", videoHandler)
 	r.HandleFunc("GET /{videoID}", videoHandler)
 	r.HandleFunc("GET /{videoID}/{formatID}", videoHandler)
+
+	r.HandleFunc("GET /proxy/{videoID}", proxyHandler)
+	r.HandleFunc("GET /proxy/{videoID}/{formatID}", proxyHandler)
 
 	// r.HandleFunc("GET /robots.txt", robotsHandler)
 

@@ -0,0 +1,82 @@
+package app
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"strconv"
+	"time"
+
+	g "github.com/birabittoh/gopipe/src/globals"
+)
+
+func proxyHandler(w http.ResponseWriter, r *http.Request) {
+	if !g.Proxy {
+		http.Error(w, err404, http.StatusNotFound)
+		return
+	}
+
+	videoID := r.PathValue("videoID")
+	formatID := getFormatID(r.PathValue("formatID"))
+
+	video, err := g.KS.Get(videoID)
+	if err != nil || video == nil {
+		http.Error(w, err404, http.StatusNotFound)
+		return
+	}
+
+	if video.ID == emptyVideo.ID {
+		http.Error(w, err500, http.StatusInternalServerError)
+		return
+	}
+
+	format := getFormat(*video, formatID)
+	if format == nil {
+		http.Error(w, err500, http.StatusInternalServerError)
+		return
+	}
+
+	key := fmt.Sprintf("%s:%d", videoID, formatID)
+	content, err := g.PKS.Get(key)
+	if err == nil && content != nil {
+		log.Println("Using cached content for ", key)
+		w.Header().Set("Content-Type", "video/mp4")
+		w.Header().Set("Content-Length", strconv.Itoa(content.Len()))
+		w.Write(content.Bytes())
+		return
+	}
+
+	res, err := g.C.Get(format.URL)
+	if err != nil {
+		http.Error(w, err500, http.StatusInternalServerError)
+		return
+	}
+	defer res.Body.Close()
+
+	w.Header().Set("Content-Type", res.Header.Get("Content-Type"))
+	w.Header().Set("Content-Length", res.Header.Get("Content-Length"))
+
+	pr, pw := io.Pipe()
+
+	// Save the content to the cache asynchronously
+	go func() {
+		var buf bytes.Buffer
+		_, err := io.Copy(&buf, pr)
+		if err != nil {
+			log.Println("Error while copying to buffer for cache:", err)
+			return
+		}
+
+		g.PKS.Set(key, buf, 5*time.Minute)
+		pw.Close()
+	}()
+
+	// Stream the content to the client while it's being downloaded and piped
+	_, err = io.Copy(io.MultiWriter(w, pw), res.Body)
+	if err != nil {
+		http.Error(w, err500, http.StatusInternalServerError)
+		return
+	}
+}

@@ -0,0 +1,27 @@
+package app
+
+import (
+	"os"
+	"strconv"
+	"strings"
+)
+
+func getEnvDefault(key string, def string) string {
+	res := os.Getenv(key)
+	if res == "" {
+		return def
+	}
+	return res
+}
+
+func parseBool(s string) bool {
+	return strings.ToLower(s) == "true" || s == "1"
+}
+
+func getFormatID(s string) int {
+	formatID, err := strconv.ParseUint(s, 10, 64)
+	if err != nil {
+		formatID = 0
+	}
+	return int(formatID)
+}

@@ -1,6 +1,8 @@
 package globals
 
 import (
+	"bytes"
+	"net/http"
 	"time"
 
 	"github.com/birabittoh/myks"
@@ -9,8 +11,12 @@ )
 
 var (
 	Debug bool
+	Proxy bool
 	Port  string
 
+	C  = http.DefaultClient
 	YT = youtube.Client{}
-	KS = myks.New[youtube.Video](3 * time.Hour)
+
+	KS  = myks.New[youtube.Video](3 * time.Hour)
+	PKS *myks.KeyStore[bytes.Buffer]
 )