From bd37c4b135d4da774c3f8db1e26f96e4233438c0 Mon Sep 17 00:00:00 2001
From: Bogdan Ungureanu <ubogdan@gmail.com>
Date: Tue, 12 Oct 2021 22:22:17 +0300
Subject: [PATCH] add content-type response headers

---
 swagger.go      | 15 ++++++++++++++-
 swagger_test.go | 18 ++++++++++++++----
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/swagger.go b/swagger.go
index a7dcba0..bf07690 100644
--- a/swagger.go
+++ b/swagger.go
@@ -3,6 +3,7 @@ package httpSwagger
 import (
 	"html/template"
 	"net/http"
+	"path/filepath"
 	"regexp"
 	"sync"
 
@@ -120,6 +121,19 @@ func Handler(configFns ...func(*Config)) http.HandlerFunc {
 			h.Prefix = matches[1]
 		})
 
+		switch filepath.Ext(path) {
+		case ".html":
+			w.Header().Set("Content-Type", "text/html; charset=utf-8")
+		case ".css":
+			w.Header().Set("Content-Type", "text/css; charset=utf-8")
+		case ".js":
+			w.Header().Set("Content-Type", "application/javascript")
+		case ".png":
+			w.Header().Set("Content-Type", "image/png")
+		case ".json":
+			w.Header().Set("Content-Type", "application/json; charset=utf-8")
+		}
+
 		switch path {
 		case "index.html":
 			_ = index.Execute(w, config)
@@ -130,7 +144,6 @@ func Handler(configFns ...func(*Config)) http.HandlerFunc {
 
 				return
 			}
-			w.Header().Set("Content-Type", "application/json; charset=utf-8")
 			_, _ = w.Write([]byte(doc))
 		case "":
 			http.Redirect(w, r, h.Prefix+"index.html", 301)
diff --git a/swagger_test.go b/swagger_test.go
index 39fee6f..57aa259 100644
--- a/swagger_test.go
+++ b/swagger_test.go
@@ -45,6 +45,7 @@ func TestWrapHandler(t *testing.T) {
 
 	w1 := performRequest("GET", "/index.html", router)
 	assert.Equal(t, 200, w1.Code)
+	assert.Equal(t, w1.Header()["Content-Type"][0], "text/html; charset=utf-8")
 
 	w2 := performRequest("GET", "/doc.json", router)
 	assert.Equal(t, 500, w2.Code)
@@ -56,12 +57,21 @@ func TestWrapHandler(t *testing.T) {
 
 	w3 := performRequest("GET", "/favicon-16x16.png", router)
 	assert.Equal(t, 200, w3.Code)
+	assert.Equal(t, w3.Header()["Content-Type"][0], "image/png")
 
-	w4 := performRequest("GET", "/notfound", router)
-	assert.Equal(t, 404, w4.Code)
+	w4 := performRequest("GET", "/swagger-ui.css", router)
+	assert.Equal(t, 200, w4.Code)
+	assert.Equal(t, w4.Header()["Content-Type"][0], "text/css; charset=utf-8")
 
-	w5 := performRequest("GET", "/", router)
-	assert.Equal(t, 301, w5.Code)
+	w5 := performRequest("GET", "/swagger-ui-bundle.js", router)
+	assert.Equal(t, 200, w5.Code)
+	assert.Equal(t, w5.Header()["Content-Type"][0], "application/javascript")
+
+	w6 := performRequest("GET", "/notfound", router)
+	assert.Equal(t, 404, w6.Code)
+
+	w7 := performRequest("GET", "/", router)
+	assert.Equal(t, 301, w7.Code)
 }
 
 func performRequest(method, target string, h http.Handler) *httptest.ResponseRecorder {