M fixyoutube.go → fixyoutube.go

@@ -1,153 +1,32 @@
 package main
 
 import (
-	"bytes"
-	"embed"
-	"html/template"
-	"io"
 	"net/http"
-	"net/url"
 	"os"
-	"regexp"
 	"strconv"
 	"time"
 
 	"github.com/BiRabittoh/fixyoutube-go/invidious"
 	"github.com/joho/godotenv"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/time/rate"
 )
-
-const templatesDirectory = "templates/"
 
 var (
-	//go:embed templates/index.html templates/video.html
-	templates      embed.FS
-	apiKey         string
-	logger         = logrus.New()
-	indexTemplate  = template.Must(template.ParseFS(templates, templatesDirectory+"index.html"))
-	videoTemplate  = template.Must(template.ParseFS(templates, templatesDirectory+"video.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}$`)
+	apiKey      string
+	debugSwitch = false
+	logger      = logrus.New()
 )
 
-func indexHandler(w http.ResponseWriter, r *http.Request) {
-	buf := &bytes.Buffer{}
-	err := indexTemplate.Execute(buf, nil)
-	if err != nil {
-		logger.Error("Failed to fill index template.")
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-		return
-	}
-
-	buf.WriteTo(w)
-}
-
-func clearHandler(w http.ResponseWriter, r *http.Request) {
-	if r.Method != http.MethodPost {
-		http.Redirect(w, r, "/", http.StatusMovedPermanently)
-		return
-	}
-
-	err := r.ParseForm()
-	if err != nil {
-		logger.Error("Failed to parse form in /clear.")
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-		return
-	}
-
-	providedKey := r.PostForm.Get("apiKey")
-	if providedKey != apiKey {
-		logger.Debug("Wrong API key: ", providedKey)
-		http.Error(w, "Wrong or missing API key.", http.StatusForbidden)
-		return
-	}
-
-	invidious.ClearDB()
-	logger.Info("Cache cleared.")
-	http.Error(w, "Done.", http.StatusOK)
-}
-
-func videoHandler(videoId string, invidiousClient *invidious.Client, w http.ResponseWriter, r *http.Request) {
-	url := "https://www.youtube.com/watch?v=" + videoId
-	userAgent := r.UserAgent()
-	res := userAgentRegex.MatchString(userAgent)
-	if !res {
-		logger.Debug("Regex did not match. Redirecting. UA:", userAgent)
-		http.Redirect(w, r, url, http.StatusMovedPermanently)
-		return
-	}
-
-	if !videoRegex.MatchString(videoId) {
-		logger.Info("Invalid video ID: ", videoId)
-		http.Error(w, "Invalid video ID.", http.StatusBadRequest)
-		return
-	}
-
-	video, err := invidiousClient.GetVideo(videoId, true)
-	if err != nil {
-		logger.Info("Wrong video ID: ", videoId)
-		http.Error(w, "Wrong video ID.", http.StatusNotFound)
-		return
-	}
-
-	if video.Url == "" {
-		logger.Debug("No URL available. Redirecting.")
-		http.Redirect(w, r, url, http.StatusMovedPermanently)
-		return
-	}
-
-	buf := &bytes.Buffer{}
-	err = videoTemplate.Execute(buf, video)
-	if err != nil {
-		logger.Error("Failed to fill video template.")
-		http.Error(w, err.Error(), http.StatusInternalServerError)
-		return
-	}
-	buf.WriteTo(w)
-}
-
-func watchHandler(invidiousClient *invidious.Client) http.HandlerFunc {
-	return func(w http.ResponseWriter, r *http.Request) {
-		u, err := url.Parse(r.URL.String())
-		if err != nil {
-			logger.Error("Failed to parse URL: ", r.URL.String())
-			http.Error(w, err.Error(), http.StatusInternalServerError)
+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
 		}
-		q := u.Query()
-		videoId := q.Get("v")
-		videoHandler(videoId, invidiousClient, w, r)
-	}
-}
-
-func shortHandler(invidiousClient *invidious.Client) http.HandlerFunc {
-	return func(w http.ResponseWriter, r *http.Request) {
-		videoId := r.PathValue("videoId")
-		videoHandler(videoId, invidiousClient, w, r)
-	}
-}
-
-func proxyHandler(invidiousClient *invidious.Client) http.HandlerFunc {
-	return func(w http.ResponseWriter, r *http.Request) {
-		videoId := r.PathValue("videoId")
-
-		vb, s := invidiousClient.ProxyVideoId(videoId)
-		if s != http.StatusOK {
-			logger.Error("proxyHandler() failed. Final code: ", s)
-			http.Error(w, "Something went wrong.", s)
-			return
-		}
-		if !vb.ValidateLength() {
-			logger.Error("Buffer length is inconsistent.")
-			http.Error(w, "Something went wrong.", http.StatusInternalServerError)
-			return
-		}
-		h := w.Header()
-		h.Set("Status", "200")
-		h.Set("Content-Type", "video/mp4")
-		h.Set("Content-Length", strconv.FormatInt(vb.Length, 10))
-		io.Copy(w, vb.Buffer)
-	}
+		next.ServeHTTP(w, r)
+	})
 }
 
 func getenvDefault(key string, def string) string {
@@ -168,14 +47,21 @@ return res
 }
 
 func main() {
-	logger.SetLevel(logrus.DebugLevel)
 	err := godotenv.Load()
 	if err != nil {
 		logger.Info("No .env file provided.")
 	}
 
+	if os.Getenv("DEBUG") != "" {
+		logger.SetLevel(logrus.DebugLevel)
+		logger.Debug("Debug mode enabled (rate limiting is disabled)")
+		debugSwitch = true
+	}
+
 	apiKey = getenvDefault("API_KEY", "itsme")
 	port := getenvDefault("PORT", "3000")
+	burstTokens := getenvDefaultParse("BURST_TOKENS", "3")
+	rateLimit := getenvDefaultParse("RATE_LIMIT", "1")
 	cacheDuration := getenvDefaultParse("CACHE_DURATION_MINUTES", "5")
 	timeoutDuration := getenvDefaultParse("TIMEOUT_DURATION_MINUTES", "10")
 	cleanupInterval := getenvDefaultParse("CLEANUP_INTERVAL_SECONDS", "30")
@@ -196,13 +82,14 @@ r.HandleFunc("/clear", clearHandler)
 	r.HandleFunc("/watch", watchHandler(videoapi))
 	r.HandleFunc("/proxy/{videoId}", proxyHandler(videoapi))
 	r.HandleFunc("/{videoId}", shortHandler(videoapi))
-	/*
-		// native go implementation
-		r := http.NewServeMux()
-		r.HandleFunc("/watch", watchHandler(videoapi))
-		r.HandleFunc("/{videoId}/", shortHandler(videoapi))
-		r.HandleFunc("/", indexHandler)
-	*/
+
+	var serveMux http.Handler
+	if debugSwitch {
+		serveMux = r
+	} else {
+		limiter := rate.NewLimiter(rate.Limit(rateLimit), int(burstTokens))
+		serveMux = limit(limiter, r)
+	}
 	logger.Info("Serving on port ", port)
-	http.ListenAndServe(":"+port, r)
+	http.ListenAndServe(":"+port, serveMux)
 }

M go.mod → go.mod

@@ -6,6 +6,8 @@ require github.com/joho/godotenv v1.5.1
 
 require github.com/mattn/go-sqlite3 v1.14.22
 
+require golang.org/x/time v0.5.0
+
 require (
 	github.com/sirupsen/logrus v1.9.3
 	golang.org/x/sys v0.17.0 // indirect

M go.sum → go.sum

@@ -15,6 +15,8 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y=
 golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
+golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

A handlers.go

@@ -0,0 +1,146 @@
+package main
+
+import (
+	"bytes"
+	"embed"
+	"io"
+	"net/http"
+	"net/url"
+	"regexp"
+	"strconv"
+	"text/template"
+
+	"github.com/BiRabittoh/fixyoutube-go/invidious"
+)
+
+const templatesDirectory = "templates/"
+
+var (
+	//go:embed templates/index.html templates/video.html
+	templates      embed.FS
+	indexTemplate  = template.Must(template.ParseFS(templates, templatesDirectory+"index.html"))
+	videoTemplate  = template.Must(template.ParseFS(templates, templatesDirectory+"video.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 indexHandler(w http.ResponseWriter, r *http.Request) {
+	buf := &bytes.Buffer{}
+	err := indexTemplate.Execute(buf, nil)
+	if err != nil {
+		logger.Error("Failed to fill index template.")
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	buf.WriteTo(w)
+}
+
+func clearHandler(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodPost {
+		http.Redirect(w, r, "/", http.StatusMovedPermanently)
+		return
+	}
+
+	err := r.ParseForm()
+	if err != nil {
+		logger.Error("Failed to parse form in /clear.")
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	providedKey := r.PostForm.Get("apiKey")
+	if providedKey != apiKey {
+		logger.Debug("Wrong API key: ", providedKey)
+		http.Error(w, "Wrong or missing API key.", http.StatusForbidden)
+		return
+	}
+
+	invidious.ClearDB()
+	logger.Info("Cache cleared.")
+	http.Error(w, "Done.", http.StatusOK)
+}
+
+func videoHandler(videoId string, invidiousClient *invidious.Client, w http.ResponseWriter, r *http.Request) {
+	url := "https://www.youtube.com/watch?v=" + videoId
+	userAgent := r.UserAgent()
+	res := userAgentRegex.MatchString(userAgent)
+	if !res {
+		logger.Debug("Regex did not match. Redirecting. UA:", userAgent)
+		http.Redirect(w, r, url, http.StatusMovedPermanently)
+		return
+	}
+
+	if !videoRegex.MatchString(videoId) {
+		logger.Info("Invalid video ID: ", videoId)
+		http.Error(w, "Invalid video ID.", http.StatusBadRequest)
+		return
+	}
+
+	video, err := invidiousClient.GetVideo(videoId, true)
+	if err != nil {
+		logger.Info("Wrong video ID: ", videoId)
+		http.Error(w, "Wrong video ID.", http.StatusNotFound)
+		return
+	}
+
+	if video.Url == "" {
+		logger.Debug("No URL available. Redirecting.")
+		http.Redirect(w, r, url, http.StatusMovedPermanently)
+		return
+	}
+
+	buf := &bytes.Buffer{}
+	err = videoTemplate.Execute(buf, video)
+	if err != nil {
+		logger.Error("Failed to fill video template.")
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+	buf.WriteTo(w)
+}
+
+func watchHandler(invidiousClient *invidious.Client) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		u, err := url.Parse(r.URL.String())
+		if err != nil {
+			logger.Error("Failed to parse URL: ", r.URL.String())
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		q := u.Query()
+		videoId := q.Get("v")
+		videoHandler(videoId, invidiousClient, w, r)
+	}
+}
+
+func shortHandler(invidiousClient *invidious.Client) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		videoId := r.PathValue("videoId")
+		videoHandler(videoId, invidiousClient, w, r)
+	}
+}
+
+func proxyHandler(invidiousClient *invidious.Client) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		videoId := r.PathValue("videoId")
+
+		vb, s := invidiousClient.ProxyVideoId(videoId)
+		if s != http.StatusOK {
+			logger.Error("proxyHandler() failed. Final code: ", s)
+			http.Error(w, http.StatusText(s), s)
+			return
+		}
+		if !vb.ValidateLength() {
+			logger.Error("Buffer length is inconsistent.")
+			status := http.StatusInternalServerError
+			http.Error(w, http.StatusText(status), status)
+			return
+		}
+		h := w.Header()
+		h.Set("Status", "200")
+		h.Set("Content-Type", "video/mp4")
+		h.Set("Content-Length", strconv.FormatInt(vb.Length, 10))
+		io.Copy(w, vb.Buffer)
+	}
+}