[Improve] Editable MIME types (#14)

* [Add] coverage to tests

* [Add] dependency gouef/utils

* [Improve] Editable MIME types

* [Test] Editable MIME types

Makefile

@@ -4,4 +4,4 @@ install:
 	go mod tidy && go mod vendor
 
 tests:
-	go test -covermode=set ./... -coverprofile=coverage.txt && go tool cover -func=coverage.txt
+	go test -v -covermode=set ./... -coverprofile=coverage.txt && go tool cover -func=coverage.txt

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 parts[0] == "" {
+		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)
 	}
 }
 

go.mod

@@ -5,6 +5,7 @@ go 1.23.4
 require (
 	github.com/gin-gonic/gin v1.10.0
 	github.com/gouef/router v1.1.0
+	github.com/gouef/utils v1.9.0
 	github.com/stretchr/testify v1.10.0
 )
 

go.sum

@@ -34,6 +34,8 @@ github.com/gouef/mode v1.0.3 h1:rdTXV2qF/YO3YzTWDrqy4connt+mXFaTMvKkcw+df4c=
 github.com/gouef/mode v1.0.3/go.mod h1:F6jloQRk8+CL9OxC4Nf5tgHiKvDXSCsONtt7RX4arss=
 github.com/gouef/router v1.1.0 h1:LQ5PLPY4BdM+2qinAwP15q3lcUw8yO8jPE35iwrGCzg=
 github.com/gouef/router v1.1.0/go.mod h1:ymO6Yo+jJkWxoesdsGa0uZJRx//6ceJOnFCXBt+3Sfs=
+github.com/gouef/utils v1.9.0 h1:yQBkkbZeJGMRfL3j89cJGTlFHoeY53znJN2iw3jMl0Y=
+github.com/gouef/utils v1.9.0/go.mod h1:7tAjRV4M5TdC2UubCcR3asVEvPRCEfSptfw3ZgitxiY=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
 github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=

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 empty", func(t *testing.T) {
+		newDiago := diago.NewDiago()
+
+		assert.False(t, newDiago.ContainsMIME(""))
+	})
 }

tests/template_test.go

@@ -30,7 +30,101 @@ 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)
+
+	r := router.NewRouter()
+	n := r.GetNativeRouter()
+	n.LoadHTMLGlob("templates/*")
+	d := diago.NewDiago()
+
+	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)
+
+	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_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_Handle_404(t *testing.T) {
 	gin.SetMode(gin.TestMode)
 	mockPanelGenerator := new(MockPanelGenerator)
 
@@ -42,8 +136,37 @@ func TestDiagoMiddleware_GenerateHTML_Error(t *testing.T) {
 
 	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)
+	})
+}
+
+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(),
+		AllowedContentTypes: diago.ContentType{
+			Types:    []string{diago.ContentType_PLAIN},
+			Charsets: []string{"*", "utf-8"},
+		},
+	})
+
+	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>")})
+		c.String(200, "OK")
 	})
 
 	t.Run("Test Custom 404 Handler", func(t *testing.T) {
@@ -53,7 +176,7 @@ 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.Contains(t, w.Body.String(), "Error generating Diago panel HTML")
 	})
 }