[MM-45817] Set cloud cookies whenever user gets logged in (#20692)

api4/user.go

@@ -1881,11 +1881,6 @@ func login(c *Context, w http.ResponseWriter, r *http.Request) {
 		c.App.AttachSessionCookies(c.AppContext, w, r)
 	}
 
-	// For context see: https://mattermost.atlassian.net/browse/MM-39583
-	if c.App.Channels().License() != nil && *c.App.Channels().License().Features.Cloud {
-		c.App.AttachCloudSessionCookie(c.AppContext, w, r)
-	}
-
 	userTermsOfService, err := c.App.GetUserTermsOfService(user.Id)
 	if err != nil && err.StatusCode != http.StatusNotFound {
 		c.Err = err

api4/user_test.go

@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"net/url"
 	"os"
 	"regexp"
 	"strings"
@@ -3610,14 +3611,61 @@ func TestLoginCookies(t *testing.T) {
 
 		th.App.Srv().SetLicense(model.NewTestLicense("cloud"))
 
+		th.Client.HTTPHeader[model.HeaderRequestedWith] = model.HeaderRequestedWithXML
 		_, resp, _ := th.Client.Login(th.BasicUser.Email, th.BasicUser.Password)
 
-		val := strings.Split(resp.Header["Set-Cookie"][0], ";")
-		cloudSessionCookie := strings.Split(val[0], "=")[1]
-		domain := strings.Split(val[2], "=")[1]
+		found := false
+		cookies := resp.Header.Values("Set-Cookie")
+		for i := range cookies {
+			if strings.Contains(cookies[i], "MMCLOUDURL") {
+				found = true
+				assert.Contains(t, cookies[i], "MMCLOUDURL=testchips;", "should contain MMCLOUDURL")
+				assert.Contains(t, cookies[i], "Domain=mattermost.com;", "should contain Domain=mattermost.com")
+				break
+			}
+		}
+		assert.True(t, found, "Did not find MMCLOUDURL cookie")
+	})
 
-		assert.Equal(t, "testchips", cloudSessionCookie)
-		assert.Equal(t, "mattermost.com", domain)
+	t.Run("should return cookie with MMCLOUDURL for cloud installations when doing cws login", func(t *testing.T) {
+		token := model.NewRandomString(64)
+		os.Setenv("CWS_CLOUD_TOKEN", token)
+
+		updateConfig := func(cfg *model.Config) {
+			*cfg.ServiceSettings.SiteURL = "https://testchips.cloud.mattermost.com"
+		}
+		th := SetupAndApplyConfigBeforeLogin(t, updateConfig).InitBasic()
+		defer th.TearDown()
+
+		th.App.Srv().SetLicense(model.NewTestLicense("cloud"))
+
+		form := url.Values{}
+		form.Add("login_id", th.SystemAdminUser.Email)
+		form.Add("cws_token", token)
+
+		th.Client.HTTPClient.CheckRedirect = func(req *http.Request, via []*http.Request) error {
+			return http.ErrUseLastResponse
+		}
+
+		r, _ := th.Client.DoAPIRequestWithHeaders(
+			http.MethodPost,
+			th.Client.APIURL+"/users/login/cws",
+			form.Encode(),
+			map[string]string{
+				"Content-Type": "application/x-www-form-urlencoded",
+			},
+		)
+		defer closeBody(r)
+
+		cookies := r.Cookies()
+		found := false
+		for i := range cookies {
+			if cookies[i].Name == model.SessionCookieCloudUrl {
+				found = true
+				assert.Equal(t, "testchips", cookies[i].Value)
+			}
+		}
+		assert.True(t, found, "should have found cookie")
 	})
 
 	t.Run("should NOT return cookie with MMCLOUDURL for cloud installations without expected format of cloud URL", func(t *testing.T) {

app/login.go

@@ -331,6 +331,11 @@ func (a *App) AttachSessionCookies(c *request.Context, w http.ResponseWriter, r
 	http.SetCookie(w, sessionCookie)
 	http.SetCookie(w, userCookie)
 	http.SetCookie(w, csrfCookie)
+
+	// For context see: https://mattermost.atlassian.net/browse/MM-39583
+	if a.Channels().License() != nil && *a.Channels().License().Features.Cloud {
+		a.AttachCloudSessionCookie(c, w, r)
+	}
 }
 
 func GetProtocol(r *http.Request) string {