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	apiKey      string
15	debugSwitch = false
16	logger      = logrus.New()
17)
18
19func limit(limiter *rate.Limiter, next http.Handler) http.Handler {
20	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
21		if !limiter.Allow() {
22			status := http.StatusTooManyRequests
23			http.Error(w, http.StatusText(status), status)
24			return
25		}
26		next.ServeHTTP(w, r)
27	})
28}
29
30func getenvDefault(key string, def string) string {
31	res := os.Getenv(key)
32	if res == "" {
33		return def
34	}
35	return res
36}
37
38func getenvDefaultParse(key string, def string) float64 {
39	value := getenvDefault(key, def)
40	res, err := strconv.ParseFloat(value, 64)
41	if err != nil {
42		logger.Fatal(err)
43	}
44	return res
45}
46
47func main() {
48	err := godotenv.Load()
49	if err != nil {
50		logger.Info("No .env file provided.")
51	}
52
53	if os.Getenv("DEBUG") != "" {
54		logger.SetLevel(logrus.DebugLevel)
55		logger.Debug("Debug mode enabled (rate limiting is disabled)")
56		debugSwitch = true
57	}
58
59	apiKey = getenvDefault("API_KEY", "itsme")
60	port := getenvDefault("PORT", "3000")
61	burstTokens := getenvDefaultParse("BURST_TOKENS", "3")
62	rateLimit := getenvDefaultParse("RATE_LIMIT", "1")
63
64	r := http.NewServeMux()
65	r.HandleFunc("/", indexHandler)
66	r.HandleFunc("/clear", clearHandler)
67	r.HandleFunc("/watch", watchHandler)
68	r.HandleFunc("/proxy/{videoId}", proxyHandler)
69	r.HandleFunc("/{videoId}", shortHandler)
70
71	var serveMux http.Handler
72	if debugSwitch {
73		serveMux = r
74	} else {
75		limiter := rate.NewLimiter(rate.Limit(rateLimit), int(burstTokens))
76		serveMux = limit(limiter, r)
77	}
78	logger.Info("Serving on port ", port)
79	http.ListenAndServe(":"+port, serveMux)
80}