From 0d1e3759ebbaf74644b32a96c791461d0c0abdce Mon Sep 17 00:00:00 2001
From: Marcus Efraimsson <marcus.efraimsson@gmail.com>
Date: Mon, 21 Jan 2019 10:20:06 +0100
Subject: [PATCH] mixor fixes

---
 pkg/middleware/middleware.go         |  2 +-
 pkg/services/auth/auth_token.go      | 23 +++++++++++++----------
 pkg/services/auth/auth_token_test.go |  3 +++
 3 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/pkg/middleware/middleware.go b/pkg/middleware/middleware.go
index 109def9ff2c..a6800971f4f 100644
--- a/pkg/middleware/middleware.go
+++ b/pkg/middleware/middleware.go
@@ -132,7 +132,7 @@ func initContextWithToken(ctx *m.ReqContext, orgID int64, ts *auth.UserAuthToken
 			Path:     setting.AppSubUrl + "/",
 		}
 
-		ctx.Resp.Header().Add("Set-Cookie", cookie.String())
+		http.SetCookie(ctx.Resp, &cookie)
 	}
 
 	return true
diff --git a/pkg/services/auth/auth_token.go b/pkg/services/auth/auth_token.go
index db5b938e0fb..aefacd7788d 100644
--- a/pkg/services/auth/auth_token.go
+++ b/pkg/services/auth/auth_token.go
@@ -19,7 +19,11 @@ func init() {
 	registry.RegisterService(&UserAuthTokenService{})
 }
 
-var now = time.Now
+var (
+	now              = time.Now
+	RotateTime       = 10 * time.Second
+	UrgentRotateTime = 5 * time.Second
+)
 
 // UserAuthTokenService are used for generating and validating user auth tokens
 type UserAuthTokenService struct {
@@ -50,7 +54,7 @@ func (s *UserAuthTokenService) UserAuthenticatedHook(user *models.User, c *model
 		Path:     setting.AppSubUrl + "/",
 	}
 
-	c.Resp.Header().Add("Set-Cookie", cookie.String())
+	http.SetCookie(c.Resp, &cookie)
 
 	return nil
 }
@@ -61,12 +65,10 @@ func (s *UserAuthTokenService) UserSignedOutHook(c *models.ReqContext) {
 		Name:     sessionCookieKey,
 		Value:    "",
 		HttpOnly: true,
-		MaxAge:   -1,
 		Domain:   setting.Domain,
 		Path:     setting.AppSubUrl + "/",
 	}
-
-	c.Resp.Header().Add("Set-Cookie", cookie.String())
+	http.SetCookie(c.Resp, &cookie)
 }
 
 func (s *UserAuthTokenService) CreateToken(userId int64, clientIP, userAgent string) (*models.UserAuthToken, error) {
@@ -115,7 +117,7 @@ func (s *UserAuthTokenService) LookupToken(unhashedToken string) (*models.UserAu
 
 	if userToken.AuthToken != hashedToken && userToken.PrevAuthToken == hashedToken && userToken.AuthTokenSeen {
 		userToken.AuthTokenSeen = false
-		expireBefore := now().Add(-1 * time.Minute).Unix()
+		expireBefore := now().Add(-RotateTime).Unix()
 		affectedRows, err := s.SQLStore.NewSession().Where("id = ? AND prev_auth_token = ? AND rotated_at < ?", userToken.Id, userToken.PrevAuthToken, expireBefore).AllCols().Update(&userToken)
 		if err != nil {
 			return nil, err
@@ -158,12 +160,12 @@ func (s *UserAuthTokenService) RefreshToken(token *models.UserAuthToken, clientI
 		return false, nil
 	}
 
-	var needsRotation = false
+	needsRotation := false
 	rotatedAt := time.Unix(token.RotatedAt, 0)
 	if token.AuthTokenSeen {
-		needsRotation = rotatedAt.Before(now().Add(time.Duration(-1) * time.Minute))
+		needsRotation = rotatedAt.Before(now().Add(-RotateTime))
 	} else {
-		needsRotation = rotatedAt.Before(now().Add(time.Duration(-30) * time.Second))
+		needsRotation = rotatedAt.Before(now().Add(-UrgentRotateTime))
 	}
 
 	s.log.Debug("refresh token", "needs rotation?", needsRotation, "auth_token_seen", token.AuthTokenSeen, "rotated_at", rotatedAt, "token.Id", token.Id)
@@ -171,6 +173,7 @@ func (s *UserAuthTokenService) RefreshToken(token *models.UserAuthToken, clientI
 		return false, nil
 	}
 
+	clientIP = util.ParseIPAddress(clientIP)
 	newToken, _ := util.RandomHex(16)
 	hashedToken := hashToken(newToken)
 
@@ -186,7 +189,7 @@ func (s *UserAuthTokenService) RefreshToken(token *models.UserAuthToken, clientI
 			rotated_at = ?
 		WHERE id = ? AND (auth_token_seen or rotated_at < ?)`
 
-	res, err := s.SQLStore.NewSession().Exec(sql, userAgent, clientIP, hashedToken, now().Unix(), token.Id, now().Add(time.Duration(-30)*time.Second))
+	res, err := s.SQLStore.NewSession().Exec(sql, userAgent, clientIP, hashedToken, now().Unix(), token.Id, now().Add(-UrgentRotateTime))
 	if err != nil {
 		return false, err
 	}
diff --git a/pkg/services/auth/auth_token_test.go b/pkg/services/auth/auth_token_test.go
index 2ee7e2d67be..bb146252fa4 100644
--- a/pkg/services/auth/auth_token_test.go
+++ b/pkg/services/auth/auth_token_test.go
@@ -170,6 +170,9 @@ func createTestContext(t *testing.T) *testContext {
 		log:      log.New("test-logger"),
 	}
 
+	RotateTime = 10 * time.Minute
+	UrgentRotateTime = time.Minute
+
 	return &testContext{
 		sqlstore:     sqlstore,
 		tokenService: tokenService,