From 2b9570f135c5dc38c9bdf43c35074c541fd62253 Mon Sep 17 00:00:00 2001
From: Noam Preil <noam@pixelhero.dev>
Date: Wed, 21 Oct 2020 23:04:47 +0000
Subject: [PATCH] Add /chart.png path

---
 main.go | 37 ++++++++++++++++++++++---------------
 1 file changed, 22 insertions(+), 15 deletions(-)

diff --git a/main.go b/main.go
index 2bce2d7..0eab06f 100644
--- a/main.go
+++ b/main.go
@@ -126,19 +126,8 @@ func metricName(metric map[string]string) string {
 	return out + "{" + strings.Join(inner, ",") + "}"
 }
 
-func main() {
-	plotutil.DefaultDashes = [][]vg.Length{{}}
-
-	if len(os.Args) < 2 {
-		fmt.Printf("Usage: %s server\n", os.Args[0])
-		os.Exit(1)
-	}
-	upstream = os.Args[1]
-	router := chi.NewRouter()
-	router.Use(middleware.RealIP)
-	router.Use(middleware.Logger)
-
-	router.Get("/chart.svg", func(w http.ResponseWriter, r *http.Request) {
+func registerExtension(router chi.Router, extension string, mime string) {
+	router.Get("/chart."+extension, func(w http.ResponseWriter, r *http.Request) {
 		args := r.URL.Query()
 		var query string
 		if q, ok := args["query"]; !ok {
@@ -261,17 +250,35 @@ func main() {
 			p.Add(plotters[i])
 		}
 
-		writer, err := p.WriterTo(width, height, "svg")
+		writer, err := p.WriterTo(width, height, extension)
 		if err != nil {
 			w.WriteHeader(400)
 			w.Write([]byte(fmt.Sprintf("%v", err)))
 			return
 		}
 
-		w.Header().Add("Content-Type", "image/svg+xml")
+		w.Header().Add("Content-Type", mime)
 		writer.WriteTo(w)
 	})
 
+}
+
+func main() {
+	plotutil.DefaultDashes = [][]vg.Length{{}}
+
+	if len(os.Args) < 2 {
+		fmt.Printf("Usage: %s server\n", os.Args[0])
+		os.Exit(1)
+	}
+	upstream = os.Args[1]
+	router := chi.NewRouter()
+
+	router.Use(middleware.RealIP)
+	router.Use(middleware.Logger)
+
+	registerExtension(router, "svg", "image/svg+xml")
+	registerExtension(router, "png", "image/png")
+
 	addr := ":8142"
 	if len(os.Args) > 2 {
 		addr = os.Args[2]