Repository: labstack/echo
Pull request #1335 (2019-05-09 09:29:23)

Fix for #1334
Commit the response code only when actually starting to write the payload 

codecov[bot]: # [Codecov](https://codecov.io/gh/labstack/echo/pull/1335?src=pr&el=h1) Report
> Merging [#1335](https://codecov.io/gh/labstack/echo/pull/1335?src=pr&el=desc) into [master](https://codecov.io/gh/labstack/echo/commit/5d2c33ad5dbb78540a56dd08d09d91f991bc3156?src=pr&el=desc) will **increase** coverage by `0.01%`.
> The diff coverage is `100%`.

[![Impacted file tree graph](https://codecov.io/gh/labstack/echo/pull/1335/graphs/tree.svg?width=650&token=vqnT76l79i&height=150&src=pr)](https://codecov.io/gh/labstack/echo/pull/1335?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #1335      +/-   ##
==========================================
+ Coverage   84.34%   84.35%   +0.01%     
==========================================
  Files          27       27              
  Lines        2012     2014       +2     
==========================================
+ Hits         1697     1699       +2     
  Misses        205      205              
  Partials      110      110
```


| [Impacted Files](https://codecov.io/gh/labstack/echo/pull/1335?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [response.go](https://codecov.io/gh/labstack/echo/pull/1335/diff?src=pr&el=tree#diff-cmVzcG9uc2UuZ28=) | `71.05% <100%> (+1.6%)` | :arrow_up: |
| [context.go](https://codecov.io/gh/labstack/echo/pull/1335/diff?src=pr&el=tree#diff-Y29udGV4dC5nbw==) | `90.71% <100%> (ø)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/labstack/echo/pull/1335?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/labstack/echo/pull/1335?src=pr&el=footer). Last update [5d2c33a...f45968e](https://codecov.io/gh/labstack/echo/pull/1335?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/README.md b/README.md
index 1db5e35..0da0312 100644
--- a/README.md
+++ b/README.md
@@ -56,30 +56,29 @@ Lower is better!
 package main
 
 import (
-	"net/http"
-
-	"github.com/labstack/echo/v4"
-	"github.com/labstack/echo/v4/middleware"
+  "net/http"
+  "github.com/labstack/echo/v4"
+  "github.com/labstack/echo/v4/middleware"
 )
 
 func main() {
-	// Echo instance
-	e := echo.New()
+  // Echo instance
+  e := echo.New()
 
-	// Middleware
-	e.Use(middleware.Logger())
-	e.Use(middleware.Recover())
+  // Middleware
+  e.Use(middleware.Logger())
+  e.Use(middleware.Recover())
 
-	// Routes
-	e.GET("/", hello)
+  // Routes
+  e.GET("/", hello)
 
-	// Start server
-	e.Logger.Fatal(e.Start(":1323"))
+  // Start server
+  e.Logger.Fatal(e.Start(":1323"))
 }
 
 // Handler
 func hello(c echo.Context) error {
-	return c.String(http.StatusOK, "Hello, World!")
+  return c.String(http.StatusOK, "Hello, World!")
 }
 ```
 
diff --git a/context.go b/context.go
index 98cf50b..065f581 100644
--- a/context.go
+++ b/context.go
@@ -437,7 +437,7 @@ func (c *context) json(code int, i interface{}, indent string) error {
 		enc.SetIndent("", indent)
 	}
 	c.writeContentType(MIMEApplicationJSONCharsetUTF8)
-	c.response.WriteHeader(code)
+	c.response.Status = code
 	return enc.Encode(i)
 }
 
diff --git a/middleware/jwt.go b/middleware/jwt.go
index 861d314..279cb6f 100644
--- a/middleware/jwt.go
+++ b/middleware/jwt.go
@@ -26,10 +26,14 @@ type (
 		// It may be used to define a custom JWT error.
 		ErrorHandler JWTErrorHandler
 
-		// Signing key to validate token.
-		// Required.
+		// Signing key to validate token. Used as fallback if SigningKeys has length 0.
+		// Required. This or SigningKeys.
 		SigningKey interface{}
 
+		// Map of signing keys to validate token with kid field usage.
+		// Required. This or SigningKey.
+		SigningKeys map[string]interface{}
+
 		// Signing method, used to check token signing method.
 		// Optional. Default value HS256.
 		SigningMethod string
@@ -110,7 +114,7 @@ func JWTWithConfig(config JWTConfig) echo.MiddlewareFunc {
 	if config.Skipper == nil {
 		config.Skipper = DefaultJWTConfig.Skipper
 	}
-	if config.SigningKey == nil {
+	if config.SigningKey == nil && len(config.SigningKeys) == 0 {
 		panic("echo: jwt middleware requires signing key")
 	}
 	if config.SigningMethod == "" {
@@ -133,6 +137,15 @@ func JWTWithConfig(config JWTConfig) echo.MiddlewareFunc {
 		if t.Method.Alg() != config.SigningMethod {
 			return nil, fmt.Errorf("unexpected jwt signing method=%v", t.Header["alg"])
 		}
+		if len(config.SigningKeys) > 0 {
+			if kid, ok := t.Header["kid"].(string); ok {
+				if key, ok := config.SigningKeys[kid]; ok {
+					return key, nil
+				}
+			}
+			return nil, fmt.Errorf("unexpected jwt key id=%v", t.Header["kid"])
+		}
+
 		return config.SigningKey, nil
 	}
 
diff --git a/response.go b/response.go
index eb2d988..ca7405c 100644
--- a/response.go
+++ b/response.go
@@ -67,7 +67,10 @@ func (r *Response) WriteHeader(code int) {
 // Write writes the data to the connection as part of an HTTP reply.
 func (r *Response) Write(b []byte) (n int, err error) {
 	if !r.Committed {
-		r.WriteHeader(http.StatusOK)
+		if r.Status == 0 {
+			r.Status = http.StatusOK
+		}
+		r.WriteHeader(r.Status)
 	}
 	n, err = r.Writer.Write(b)
 	r.Size += int64(n)


Test patch:
diff --git a/context_test.go b/context_test.go
index 6caa270..fe542c8 100644
--- a/context_test.go
+++ b/context_test.go
@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"mime/multipart"
 	"net/http"
 	"net/http/httptest"
@@ -374,6 +375,34 @@ func TestContext(t *testing.T) {
 	assert.Equal(0, len(c.QueryParams()))
 }
 
+func TestContext_JSON_CommitsCustomResponseCode(t *testing.T) {
+	e := New()
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec).(*context)
+	err := c.JSON(http.StatusCreated, user{1, "Jon Snow"})
+
+	assert := testify.New(t)
+	if assert.NoError(err) {
+		assert.Equal(http.StatusCreated, rec.Code)
+		assert.Equal(MIMEApplicationJSONCharsetUTF8, rec.Header().Get(HeaderContentType))
+		assert.Equal(userJSON+"\n", rec.Body.String())
+	}
+}
+
+func TestContext_JSON_DoesntCommitResponseCodePrematurely(t *testing.T) {
+	e := New()
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec).(*context)
+	err := c.JSON(http.StatusCreated, map[string]float64{"a": math.NaN()})
+
+	assert := testify.New(t)
+	if assert.Error(err) {
+		assert.False(c.response.Committed)
+	}
+}
+
 func TestContextCookie(t *testing.T) {
 	e := New()
 	req := httptest.NewRequest(http.MethodGet, "/", nil)
diff --git a/middleware/jwt_test.go b/middleware/jwt_test.go
index fc90073..8165c2d 100644
--- a/middleware/jwt_test.go
+++ b/middleware/jwt_test.go
@@ -224,3 +224,93 @@ func TestJWT(t *testing.T) {
 		}
 	}
 }
+
+func TestJWTwithKID(t *testing.T) {
+	test := assert.New(t)
+
+	e := echo.New()
+	handler := func(c echo.Context) error {
+		return c.String(http.StatusOK, "test")
+	}
+	firstToken := "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiIsImtpZCI6ImZpcnN0T25lIn0.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9.w5VGpHOe0jlNgf7jMVLHzIYH_XULmpUlreJnilwSkWk"
+	secondToken := "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiIsImtpZCI6InNlY29uZE9uZSJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9.sdghDYQ85jdh0hgQ6bKbMguLI_NSPYWjkhVJkee-yZM"
+	wrongToken := "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiIsImtpZCI6InNlY29uZE9uZSJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9.RyhLybtVLpoewF6nz9YN79oXo32kAtgUxp8FNwTkb90"
+	staticToken := "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiYWRtaW4iOnRydWV9.1_-XFYUPpJfgsaGwYhgZEt7hfySMg-a3GN-nfZmbW7o"
+	validKeys := map[string]interface{}{"firstOne": []byte("first_secret"), "secondOne": []byte("second_secret")}
+	invalidKeys := map[string]interface{}{"thirdOne": []byte("third_secret")}
+	staticSecret := []byte("static_secret")
+	invalidStaticSecret := []byte("invalid_secret")
+
+	for _, tc := range []struct {
+		expErrCode int // 0 for Success
+		config     JWTConfig
+		hdrAuth    string
+		info       string
+	}{
+		{
+			hdrAuth: DefaultJWTConfig.AuthScheme + " " + firstToken,
+			config:  JWTConfig{SigningKeys: validKeys},
+			info:    "First token valid",
+		},
+		{
+			hdrAuth: DefaultJWTConfig.AuthScheme + " " + secondToken,
+			config:  JWTConfig{SigningKeys: validKeys},
+			info:    "Second token valid",
+		},
+		{
+			expErrCode: http.StatusUnauthorized,
+			hdrAuth:    DefaultJWTConfig.AuthScheme + " " + wrongToken,
+			config:     JWTConfig{SigningKeys: validKeys},
+			info:       "Wrong key id token",
+		},
+		{
+			hdrAuth: DefaultJWTConfig.AuthScheme + " " + staticToken,
+			config:  JWTConfig{SigningKey: staticSecret},
+			info:    "Valid static secret token",
+		},
+		{
+			expErrCode: http.StatusUnauthorized,
+			hdrAuth:    DefaultJWTConfig.AuthScheme + " " + staticToken,
+			config:     JWTConfig{SigningKey: invalidStaticSecret},
+			info:       "Invalid static secret",
+		},
+		{
+			expErrCode: http.StatusUnauthorized,
+			hdrAuth:    DefaultJWTConfig.AuthScheme + " " + firstToken,
+			config:     JWTConfig{SigningKeys: invalidKeys},
+			info:       "Invalid keys first token",
+		},
+		{
+			expErrCode: http.StatusUnauthorized,
+			hdrAuth:    DefaultJWTConfig.AuthScheme + " " + secondToken,
+			config:     JWTConfig{SigningKeys: invalidKeys},
+			info:       "Invalid keys second token",
+		},
+	} {
+		req := httptest.NewRequest(http.MethodGet, "/", nil)
+		res := httptest.NewRecorder()
+		req.Header.Set(echo.HeaderAuthorization, tc.hdrAuth)
+		c := e.NewContext(req, res)
+
+		if tc.expErrCode != 0 {
+			h := JWTWithConfig(tc.config)(handler)
+			he := h(c).(*echo.HTTPError)
+			test.Equal(tc.expErrCode, he.Code, tc.info)
+			continue
+		}
+
+		h := JWTWithConfig(tc.config)(handler)
+		if test.NoError(h(c), tc.info) {
+			user := c.Get("user").(*jwt.Token)
+			switch claims := user.Claims.(type) {
+			case jwt.MapClaims:
+				test.Equal(claims["name"], "John Doe", tc.info)
+			case *jwtCustomClaims:
+				test.Equal(claims.Name, "John Doe", tc.info)
+				test.Equal(claims.Admin, true, tc.info)
+			default:
+				panic("unexpected type of claims")
+			}
+		}
+	}
+}
diff --git a/response_test.go b/response_test.go
index 8f08250..bc570a5 100644
--- a/response_test.go
+++ b/response_test.go
@@ -22,3 +22,22 @@ func TestResponse(t *testing.T) {
 	res.Write([]byte("test"))
 	assert.Equal(t, "echo", rec.Header().Get(HeaderServer))
 }
+
+func TestResponse_Write_FallsBackToDefaultStatus(t *testing.T) {
+	e := New()
+	rec := httptest.NewRecorder()
+	res := &Response{echo: e, Writer: rec}
+
+	res.Write([]byte("test"))
+	assert.Equal(t, http.StatusOK, rec.Code)
+}
+
+func TestResponse_Write_UsesSetResponseCode(t *testing.T) {
+	e := New()
+	rec := httptest.NewRecorder()
+	res := &Response{echo: e, Writer: rec}
+
+	res.Status = http.StatusBadRequest
+	res.Write([]byte("test"))
+	assert.Equal(t, http.StatusBadRequest, rec.Code)
+}