Merge branch 'coderhaoxin-middleware-secure-header'

context.go

@@ -87,7 +87,7 @@ type (
 
 		// Cookie returns the named cookie provided in the request.
 		// It is an alias for `engine.Request#Cookie()`.
-		Cookie(string) engine.Cookie
+		Cookie(string) (engine.Cookie, error)
 
 		// SetCookie adds a `Set-Cookie` header in HTTP response.
 		// It is an alias for `engine.Response#SetCookie()`.
@@ -295,7 +295,7 @@ func (c *context) MultipartForm() (*multipart.Form, error) {
 	return c.request.MultipartForm()
 }
 
-func (c *context) Cookie(name string) engine.Cookie {
+func (c *context) Cookie(name string) (engine.Cookie, error) {
 	return c.request.Cookie(name)
 }
 

context_test.go

@@ -186,9 +186,11 @@ func TestContextCookie(t *testing.T) {
 	c := e.NewContext(req, rec).(*context)
 
 	// Read single
-	cookie := c.Cookie("theme")
-	assert.Equal(t, "theme", cookie.Name())
-	assert.Equal(t, "light", cookie.Value())
+	cookie, err := c.Cookie("theme")
+	if assert.NoError(t, err) {
+		assert.Equal(t, "theme", cookie.Name())
+		assert.Equal(t, "light", cookie.Value())
+	}
 
 	// Read multiple
 	for _, cookie := range c.Cookies() {

echo.go

@@ -166,6 +166,13 @@ const (
 	HeaderAccessControlAllowCredentials = "Access-Control-Allow-Credentials"
 	HeaderAccessControlExposeHeaders    = "Access-Control-Expose-Headers"
 	HeaderAccessControlMaxAge           = "Access-Control-Max-Age"
+
+	// Security
+	HeaderStrictTransportSecurity = "Strict-Transport-Security"
+	HeaderXContentTypeOptions     = "X-Content-Type-Options"
+	HeaderXXSSProtection          = "X-XSS-Protection"
+	HeaderXFrameOptions           = "X-Frame-Options"
+	HeaderContentSecurityPolicy   = "Content-Security-Policy"
 )
 
 var (
@@ -191,6 +198,7 @@ var (
 	ErrStatusRequestEntityTooLarge = NewHTTPError(http.StatusRequestEntityTooLarge)
 	ErrRendererNotRegistered       = errors.New("renderer not registered")
 	ErrInvalidRedirectCode         = errors.New("invalid redirect status code")
+	ErrCookieNotFound              = errors.New("cookie not found")
 )
 
 // Error handlers

engine/engine.go

@@ -85,7 +85,7 @@ type (
 		MultipartForm() (*multipart.Form, error)
 
 		// Cookie returns the named cookie provided in the request.
-		Cookie(string) Cookie
+		Cookie(string) (Cookie, error)
 
 		// Cookies returns the HTTP cookies sent with the request.
 		Cookies() []Cookie

engine/fasthttp/request.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"mime/multipart"
 
+	"github.com/labstack/echo"
 	"github.com/labstack/echo/engine"
 	"github.com/labstack/gommon/log"
 	"github.com/valyala/fasthttp"
@@ -128,11 +129,15 @@ func (r *Request) MultipartForm() (*multipart.Form, error) {
 }
 
 // Cookie implements `engine.Request#Cookie` function.
-func (r *Request) Cookie(name string) engine.Cookie {
+func (r *Request) Cookie(name string) (engine.Cookie, error) {
 	c := new(fasthttp.Cookie)
 	c.SetKey(name)
-	c.ParseBytes(r.Request.Header.Cookie(name))
-	return &Cookie{c}
+	b := r.Request.Header.Cookie(name)
+	if b == nil {
+		return nil, echo.ErrCookieNotFound
+	}
+	c.ParseBytes(b)
+	return &Cookie{c}, nil
 }
 
 // Cookies implements `engine.Request#Cookies` function.

engine/standard/request.go

@@ -153,9 +153,12 @@ func (r *Request) MultipartForm() (*multipart.Form, error) {
 }
 
 // Cookie implements `engine.Request#Cookie` function.
-func (r *Request) Cookie(name string) engine.Cookie {
-	c, _ := r.Request.Cookie(name)
-	return &Cookie{c}
+func (r *Request) Cookie(name string) (engine.Cookie, error) {
+	c, err := r.Request.Cookie(name)
+	if err != nil {
+		return nil, echo.ErrCookieNotFound
+	}
+	return &Cookie{c}, nil
 }
 
 // Cookies implements `engine.Request#Cookies` function.

middleware/secure.go

@@ -0,0 +1,55 @@
+package middleware
+
+import (
+	"fmt"
+
+	"github.com/labstack/echo"
+)
+
+type (
+	SecureConfig struct {
+		DisableXSSProtection         bool
+		DisableContentTypeNosniff    bool
+		XFrameOptions                string
+		DisableHSTSIncludeSubdomains bool
+		HSTSMaxAge                   int
+		ContentSecurityPolicy        string
+	}
+)
+
+var (
+	DefaultSecureConfig = SecureConfig{
+		XFrameOptions: "SAMEORIGIN",
+	}
+)
+
+func Secure() echo.MiddlewareFunc {
+	return SecureWithConfig(DefaultSecureConfig)
+}
+
+func SecureWithConfig(config SecureConfig) echo.MiddlewareFunc {
+	return func(next echo.HandlerFunc) echo.HandlerFunc {
+		return func(c echo.Context) error {
+			if !config.DisableXSSProtection {
+				c.Response().Header().Set(echo.HeaderXXSSProtection, "1; mode=block")
+			}
+			if !config.DisableContentTypeNosniff {
+				c.Response().Header().Set(echo.HeaderXContentTypeOptions, "nosniff")
+			}
+			if config.XFrameOptions != "" {
+				c.Response().Header().Set(echo.HeaderXFrameOptions, config.XFrameOptions)
+			}
+			if config.HSTSMaxAge != 0 {
+				subdomains := ""
+				if !config.DisableHSTSIncludeSubdomains {
+					subdomains = "; includeSubdomains"
+				}
+				c.Response().Header().Set(echo.HeaderStrictTransportSecurity, fmt.Sprintf("max-age=%d%s", config.HSTSMaxAge, subdomains))
+			}
+			if config.ContentSecurityPolicy != "" {
+				c.Response().Header().Set(echo.HeaderContentSecurityPolicy, config.ContentSecurityPolicy)
+			}
+			return next(c)
+		}
+	}
+}

middleware/secure_test.go

@@ -0,0 +1,32 @@
+package middleware
+
+// func TestSecureWithConfig(t *testing.T) {
+// 	e := echo.New()
+//
+// 	config := SecureConfig{
+// 		STSMaxAge:             100,
+// 		STSIncludeSubdomains:  true,
+// 		FrameDeny:             true,
+// 		FrameOptionsValue:     "",
+// 		ContentTypeNosniff:    true,
+// 		XssProtected:          true,
+// 		XssProtectionValue:    "",
+// 		ContentSecurityPolicy: "default-src 'self'",
+// 		DisableProdCheck:      true,
+// 	}
+// 	secure := SecureWithConfig(config)
+// 	h := secure(func(c echo.Context) error {
+// 		return c.String(http.StatusOK, "test")
+// 	})
+//
+// 	rq := test.NewRequest(echo.GET, "/", nil)
+// 	rc := test.NewResponseRecorder()
+// 	c := e.NewContext(rq, rc)
+// 	h(c)
+//
+// 	assert.Equal(t, "max-age=100; includeSubdomains", rc.Header().Get(stsHeader))
+// 	assert.Equal(t, "DENY", rc.Header().Get(frameOptionsHeader))
+// 	assert.Equal(t, "nosniff", rc.Header().Get(contentTypeHeader))
+// 	assert.Equal(t, xssProtectionValue, rc.Header().Get(xssProtectionHeader))
+// 	assert.Equal(t, "default-src 'self'", rc.Header().Get(cspHeader))
+// }

test/request.go

@@ -1,6 +1,7 @@
 package test
 
 import (
+	"errors"
 	"io"
 	"io/ioutil"
 	"mime/multipart"
@@ -130,9 +131,12 @@ func (r *Request) MultipartForm() (*multipart.Form, error) {
 	return r.request.MultipartForm, err
 }
 
-func (r *Request) Cookie(name string) engine.Cookie {
-	c, _ := r.request.Cookie(name)
-	return &Cookie{c}
+func (r *Request) Cookie(name string) (engine.Cookie, error) {
+	c, err := r.request.Cookie(name)
+	if err != nil {
+		return nil, errors.New("cookie not found")
+	}
+	return &Cookie{c}, nil
 }
 
 // Cookies implements `engine.Request#Cookies` function.