[Improve] Editable MIME types

diago.go

@@ -1,20 +1,83 @@
 package diago
 
-type Diago struct {
-	Extensions []Extension
+import (
+	"github.com/gouef/utils"
+	"strings"
+)
+
+type ContentType struct {
+	Types    []string
+	Charsets []string
+}
 
-	TemplateProvider TemplateProvider
-	PanelGenerator   PanelGenerator
+const (
+	ContentType_HTML  = "text/html"
+	ContentType_PLAIN = "text/plain"
+
+	Charset_UTF8 = "utf-8"
+	Charset_ALL  = "*"
+)
+
+type Diago struct {
+	Extensions          []Extension
+	TemplateProvider    TemplateProvider
+	PanelGenerator      PanelGenerator
+	AllowedContentTypes ContentType
 }
 
 func NewDiago() *Diago {
 
 	return &Diago{
 		TemplateProvider: NewDefaultTemplateProvider(),
 		PanelGenerator:   NewDefaultPanelGenerator(),
+		AllowedContentTypes: ContentType{
+			Types: []string{
+				ContentType_HTML,
+				ContentType_PLAIN,
+			},
+			Charsets: []string{
+				Charset_ALL,
+			},
+		},
 	}
 }
 
+func (d *Diago) SetAllowedContentTypes(contentType ContentType) *Diago {
+	d.AllowedContentTypes = contentType
+	return d
+}
+
+func (d *Diago) AddContentType(typeString string) *Diago {
+	d.AllowedContentTypes.Types = append(d.AllowedContentTypes.Types, typeString)
+	return d
+}
+
+func (d *Diago) AddContentCharset(charset string) *Diago {
+	d.AllowedContentTypes.Types = append(d.AllowedContentTypes.Charsets, charset)
+	return d
+}
+
+func (d *Diago) ContainsMIME(header string) bool {
+	parts := strings.Split(header, ";")
+	if len(parts) < 1 {
+		return false
+	}
+	contentType := strings.TrimSpace(parts[0])
+
+	charset := ""
+	if len(parts) > 1 {
+		for _, part := range parts[1:] {
+			if strings.HasPrefix(strings.TrimSpace(part), "charset=") {
+				part = strings.TrimSpace(part)
+				charset = strings.TrimSpace(strings.TrimPrefix(part, "charset="))
+				break
+			}
+		}
+	}
+
+	return utils.InListArray([]string{"*", charset}, d.AllowedContentTypes.Charsets) && utils.InArray(contentType, d.AllowedContentTypes.Types)
+}
+
 func (d *Diago) GetExtensions() []Extension {
 	return d.Extensions
 }

diagoMiddleware.go

@@ -45,7 +45,7 @@ func DiagoMiddleware(r *router.Router, d *Diago) gin.HandlerFunc {
 
 		contentType := writer.Header().Get("Content-Type")
 
-		if contentType == "text/html; charset=utf-8" {
+		if d.ContainsMIME(contentType) {
 			var extensionsHtml []template.HTML
 			var extensionsPanelHtml []template.HTML
 			var extensionsJSHtml []template.HTML
@@ -74,8 +74,10 @@ func DiagoMiddleware(r *router.Router, d *Diago) gin.HandlerFunc {
 			writer.buffer.WriteString(diagoPanelHTML)
 		}
 
+		status := c.Writer.Status()
 		c.Writer = originalWriter
 		c.Writer.Write(responseBuffer.Bytes())
+		c.Status(status)
 	}
 }
 

tests/diago_test.go

@@ -131,4 +131,22 @@ func TestDiago(t *testing.T) {
 
 		assert.Equal(t, "<script>console.log('test');</script>", newDiago.GetExtensions()[0].GetJSHtml(nil))
 	})
+
+	t.Run("Test ContainsMIME", func(t *testing.T) {
+		newDiago := diago.NewDiago()
+
+		assert.True(t, newDiago.ContainsMIME(diago.ContentType_PLAIN))
+	})
+
+	t.Run("Test ContainsMIME false", func(t *testing.T) {
+		newDiago := diago.NewDiago()
+
+		assert.False(t, newDiago.ContainsMIME("application/json; charset=test"))
+	})
+
+	t.Run("Test ContainsMIME false not contain charset", func(t *testing.T) {
+		newDiago := diago.NewDiago()
+
+		assert.False(t, newDiago.ContainsMIME("application/json"))
+	})
 }

tests/template_test.go

@@ -30,15 +30,49 @@ func (m *MockPanelGenerator) GenerateHTML(name string, templateProvider diago.Te
 	return "", errors.New("template parsing error")
 }
 
-func TestDiagoMiddleware_GenerateHTML_Error(t *testing.T) {
+func TestDiagoMiddleware_ContentTypeAndCharset(t *testing.T) {
 	gin.SetMode(gin.TestMode)
-	mockPanelGenerator := new(MockPanelGenerator)
 
 	r := router.NewRouter()
 	n := r.GetNativeRouter()
 	n.LoadHTMLGlob("templates/*")
+	d := diago.NewDiago()
 
-	middleware := diago.DiagoMiddleware(r, &diago.Diago{PanelGenerator: mockPanelGenerator, TemplateProvider: diago.NewDefaultTemplateProvider()})
+	middleware := diago.DiagoMiddleware(r, d)
+
+	n.Use(middleware)
+
+	r.AddRouteGet("notfound", "/notfound", func(c *gin.Context) {
+		c.HTML(http.StatusOK, "status.gohtml", gin.H{"content": template.HTML("<div>OK</div>")})
+	})
+
+	t.Run("Test Custom 404 Handler", func(t *testing.T) {
+		req := httptest.NewRequest(http.MethodGet, "/notfound", nil)
+		w := httptest.NewRecorder()
+
+		r.GetNativeRouter().ServeHTTP(w, req)
+
+		assert.Equal(t, 200, w.Code)
+		assert.Contains(t, w.Body.String(), `<div>OK</div>`)
+	})
+}
+
+func TestDiagoMiddleware_ContentTypeAndCharsetAdd(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	r := router.NewRouter()
+	n := r.GetNativeRouter()
+	n.LoadHTMLGlob("templates/*")
+	d := diago.NewDiago()
+	d.SetAllowedContentTypes(diago.ContentType{
+		Types:    []string{},
+		Charsets: []string{},
+	})
+
+	d.AddContentCharset("utf-8")
+	d.AddContentType(diago.ContentType_HTML)
+
+	middleware := diago.DiagoMiddleware(r, d)
 
 	n.Use(middleware)
 
@@ -52,8 +86,63 @@ func TestDiagoMiddleware_GenerateHTML_Error(t *testing.T) {
 
 		r.GetNativeRouter().ServeHTTP(w, req)
 
-		assert.Equal(t, 500, w.Code)
-		assert.Equal(t, `<div id="content"><div>OK</div></div>Error generating Diago panel HTML`, w.Body.String())
+		assert.Equal(t, 200, w.Code)
+		assert.Contains(t, w.Body.String(), `<div>OK</div>`)
+	})
+}
+
+func TestDiagoMiddleware_ContentTypeAndCharset_NotAllowed(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	r := router.NewRouter()
+	n := r.GetNativeRouter()
+	n.LoadHTMLGlob("templates/*")
+	d := diago.NewDiago()
+	d.SetAllowedContentTypes(diago.ContentType{
+		Types:    []string{},
+		Charsets: []string{},
+	})
+
+	middleware := diago.DiagoMiddleware(r, d)
+
+	n.Use(middleware)
+
+	r.AddRouteGet("notfound", "/notfound", func(c *gin.Context) {
+		c.HTML(http.StatusOK, "status.gohtml", gin.H{"content": template.HTML("<div>OK</div>")})
+	})
+
+	t.Run("Test Custom 404 Handler", func(t *testing.T) {
+		req := httptest.NewRequest(http.MethodGet, "/notfound", nil)
+		w := httptest.NewRecorder()
+
+		r.GetNativeRouter().ServeHTTP(w, req)
+
+		responseResult := w.Body.String()
+		assert.Equal(t, 200, w.Code)
+		assert.Contains(t, responseResult, `<div id="content"><div>OK</div></div>`)
+		assert.NotContains(t, responseResult, "Error generating Diago panel HTML")
+	})
+}
+
+func TestDiagoMiddleware_GenerateHTML_Error(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	mockPanelGenerator := new(MockPanelGenerator)
+
+	r := router.NewRouter()
+	n := r.GetNativeRouter()
+	n.LoadHTMLGlob("templates/*")
+
+	middleware := diago.DiagoMiddleware(r, &diago.Diago{PanelGenerator: mockPanelGenerator, TemplateProvider: diago.NewDefaultTemplateProvider()})
+
+	n.Use(middleware)
+
+	t.Run("Test Custom 404 Handler", func(t *testing.T) {
+		req := httptest.NewRequest(http.MethodGet, "/notfound", nil)
+		w := httptest.NewRecorder()
+
+		r.GetNativeRouter().ServeHTTP(w, req)
+
+		assert.Equal(t, 404, w.Code)
 	})
 }