1package main
 2
 3import (
 4	"net/http"
 5	"os"
 6	"strconv"
 7
 8	"github.com/joho/godotenv"
 9	"github.com/sirupsen/logrus"
10	"golang.org/x/time/rate"
11)
12
13var (
14	debugSwitch = false
15	logger      = logrus.New()
16)
17
18func limit(limiter *rate.Limiter, next http.Handler) http.Handler {
19	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
20		if !limiter.Allow() {
21			status := http.StatusTooManyRequests
22			http.Error(w, http.StatusText(status), status)
23			return
24		}
25		next.ServeHTTP(w, r)
26	})
27}
28
29func getenvDefault(key string, def string) string {
30	res := os.Getenv(key)
31	if res == "" {
32		return def
33	}
34	return res
35}
36
37func getenvDefaultParse(key string, def string) float64 {
38	value := getenvDefault(key, def)
39	res, err := strconv.ParseFloat(value, 64)
40	if err != nil {
41		logger.Fatal(err)
42	}
43	return res
44}
45
46func main() {
47	err := godotenv.Load()
48	if err != nil {
49		logger.Info("No .env file provided.")
50	}
51
52	if os.Getenv("DEBUG") != "" {
53		logger.SetLevel(logrus.DebugLevel)
54		logger.Debug("Debug mode enabled (rate limiting is disabled)")
55		debugSwitch = true
56	}
57
58	port := getenvDefault("PORT", "3000")
59	burstTokens := getenvDefaultParse("BURST_TOKENS", "3")
60	rateLimit := getenvDefaultParse("RATE_LIMIT", "1")
61
62	r := http.NewServeMux()
63	r.HandleFunc("GET /", indexHandler)
64	r.HandleFunc("GET /watch", watchHandler)
65	r.HandleFunc("GET /shorts/{videoID}", shortHandler)
66	r.HandleFunc("GET /proxy/{videoID}", proxyHandler)
67	r.HandleFunc("GET /refresh/{videoID}", refreshHandler)
68	r.HandleFunc("GET /{videoID}", shortHandler)
69
70	r.HandleFunc("POST /download", downloadHandler)
71
72	var serveMux http.Handler
73	if debugSwitch {
74		serveMux = r
75	} else {
76		limiter := rate.NewLimiter(rate.Limit(rateLimit), int(burstTokens))
77		serveMux = limit(limiter, r)
78	}
79	logger.Info("Serving on port ", port)
80	http.ListenAndServe(":"+port, serveMux)
81}