diff --git a/Makefile b/Makefile index 17d0bf9..3dba678 100644 --- a/Makefile +++ b/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 \ No newline at end of file + go test -v -covermode=set ./... -coverprofile=coverage.txt && go tool cover -func=coverage.txt \ No newline at end of file diff --git a/diago.go b/diago.go index a256ce2..fff981d 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 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 } 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/go.mod b/go.mod index 78002cc..d8956fe 100644 --- a/go.mod +++ b/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 ) diff --git a/go.sum b/go.sum index 55204ad..4842563 100644 --- a/go.sum +++ b/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= diff --git a/tests/diago_test.go b/tests/diago_test.go index 6438bc7..da1628b 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 empty", func(t *testing.T) { + newDiago := diago.NewDiago() + + assert.False(t, newDiago.ContainsMIME("")) + }) } diff --git a/tests/template_test.go b/tests/template_test.go index 75c8839..f88112f 100644 --- a/tests/template_test.go +++ b/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("
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) + + 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_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_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("
OK
")}) + 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, `
OK
Error generating Diago panel HTML`, w.Body.String()) + assert.Contains(t, w.Body.String(), "Error generating Diago panel HTML") }) }