From 0994194049c8162490da3c6a99a636bab03349e6 Mon Sep 17 00:00:00 2001 From: Jan Galek Date: Mon, 23 Dec 2024 00:46:00 +0100 Subject: [PATCH] [Improve] Editable MIME types --- diago.go | 71 ++++++++++++++++++++++++++++-- diagoMiddleware.go | 4 +- tests/diago_test.go | 18 ++++++++ tests/template_test.go | 99 +++++++++++++++++++++++++++++++++++++++--- 4 files changed, 182 insertions(+), 10 deletions(-) diff --git a/diago.go b/diago.go index a256ce2..0b5d76f 100644 --- a/diago.go +++ b/diago.go @@ -1,10 +1,28 @@ 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 { @@ -12,9 +30,54 @@ 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 } diff --git a/diagoMiddleware.go b/diagoMiddleware.go index 5333427..c4ab3b6 100644 --- a/diagoMiddleware.go +++ b/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) } } diff --git a/tests/diago_test.go b/tests/diago_test.go index 6438bc7..63310fb 100644 --- a/tests/diago_test.go +++ b/tests/diago_test.go @@ -131,4 +131,22 @@ func TestDiago(t *testing.T) { assert.Equal(t, "", 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")) + }) } diff --git a/tests/template_test.go b/tests/template_test.go index 75c8839..4e39b35 100644 --- a/tests/template_test.go +++ b/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("
OK
")}) + }) + + 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(), `
OK
`) + }) +} + +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, `
OK
Error generating Diago panel HTML`, w.Body.String()) + assert.Equal(t, 200, w.Code) + assert.Contains(t, w.Body.String(), `
OK
`) + }) +} + +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("
OK
")}) + }) + + 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, `
OK
`) + 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) }) }