diff --git a/main.go b/main.go index 6ea3392..02d0799 100644 --- a/main.go +++ b/main.go @@ -24,7 +24,7 @@ func main() { r.Use(securityHeaders) r.HandleFunc("/{id}-lyrics", lyricsHandler) - r.HandleFunc("/images/{filename}.jpg", proxyHandler) + r.HandleFunc("/images/{filename}.{ext}", proxyHandler) r.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.FileServer(http.Dir("static")))) server := &http.Server{ diff --git a/proxy.go b/proxy.go index ccb3666..76bbb66 100644 --- a/proxy.go +++ b/proxy.go @@ -4,13 +4,33 @@ import ( "fmt" "io" "net/http" + "strings" "github.com/gorilla/mux" ) +func isValidExt(ext string) bool { + valid := []string{"jpg", "jpeg", "png", "gif"} + for _, c := range valid { + if strings.ToLower(ext) == c { + return true + } + } + + return false +} + func proxyHandler(w http.ResponseWriter, r *http.Request) { - f := mux.Vars(r)["filename"] - url := fmt.Sprintf("https://images.genius.com/%s.jpg", f) + v := mux.Vars(r) + f := v["filename"] + ext := v["ext"] + + if !isValidExt(ext) { + write(w, http.StatusBadRequest, []byte("not an image :/")) + return + } + + url := fmt.Sprintf("https://images.genius.com/%s.%s", f, ext) res, err := http.Get(url) if err != nil { @@ -23,6 +43,6 @@ func proxyHandler(w http.ResponseWriter, r *http.Request) { return } - w.Header().Add("Content-type", "image/jpeg") + w.Header().Add("Content-type", fmt.Sprintf("image/%s", ext)) io.Copy(w, res.Body) }