MM-69609: Recover from GitHub SAML SSO expiry without forcing a full reconnect (#1036)

* MM-69609: Recover from GitHub SAML SSO expiry without forcing a full reconnect

* linT

* Coderabbit suggestions

---------

Co-authored-by: Nevyana Angelova <nevyangelova@192.168.100.47>

Author: nang2049

server/plugin/api.go

@@ -799,7 +799,7 @@ func (p *Plugin) getPrsDetails(c *UserContext, w http.ResponseWriter, r *http.Re
 	wg.Wait()
 
 	if isGitHubAuthFailure(fetchErr) {
-		p.handleRevokedToken(c.GHInfo)
+		p.handleAuthFailure(c.GHInfo, fetchErr)
 		p.writeAPIError(w, &APIErrorResponse{ID: "", Message: "Not authorized.", StatusCode: http.StatusUnauthorized})
 		return
 	}
@@ -1124,7 +1124,7 @@ func (p *Plugin) getLHSData(c *UserContext) (reviewResp []*graphql.GithubPRDetai
 
 	reviewResp, assignmentResp, openPRResp, err = graphQLClient.GetLHSData(c.Ctx)
 	if isGitHubAuthFailure(err) {
-		p.handleRevokedToken(c.GHInfo)
+		p.handleAuthFailure(c.GHInfo, err)
 	}
 	if err != nil {
 		return reviewResp, assignmentResp, openPRResp, err

server/plugin/plugin.go

@@ -1350,7 +1350,7 @@ func (p *Plugin) useGitHubClient(info *GitHubUserInfo, toRun func(info *GitHubUs
 	}
 
 	if isGitHubAuthFailure(err) {
-		p.handleRevokedToken(info)
+		p.handleAuthFailure(info, err)
 	}
 
 	return err
@@ -1424,3 +1424,57 @@ func (p *Plugin) handleRevokedToken(info *GitHubUserInfo) {
 	p.disconnectGitHubAccount(info.UserID)
 	p.CreateBotDMPost(info.UserID, "Your Github account was disconnected due to an invalid or revoked authorization token. Reconnect your account using the `/github connect` command.", "custom_git_revoked_token")
 }
+
+const (
+	authFailureDMKey      = "auth_failure_dm_"
+	authFailureDMCooldown = time.Hour
+)
+
+func (p *Plugin) handleAuthFailure(info *GitHubUserInfo, err error) {
+	if ssoURL := extractSSOAuthorizeURL(err); ssoURL != "" {
+		p.dmSAMLReauthorize(info, ssoURL)
+		return
+	}
+	p.handleRevokedToken(info)
+}
+
+func (p *Plugin) dmSAMLReauthorize(info *GitHubUserInfo, ssoURL string) {
+	// Throttle only DM once per cooldown window per user.
+	ok, err := p.store.Set(authFailureDMKey+info.UserID, []byte("1"),
+		pluginapi.SetExpiry(authFailureDMCooldown),
+		pluginapi.SetAtomic(nil),
+	)
+	if err != nil {
+		p.client.Log.Warn("Failed to set SAML reauthorize DM cooldown", "error", err.Error())
+	}
+	if !ok {
+		return
+	}
+	msg := fmt.Sprintf(
+		"Your GitHub SAML SSO session expired, so GitHub temporarily blocked this plugin from acting on your behalf. "+
+			"Re-authorize SSO for your organization here: %s\n\n"+
+			"You do not need to run `/github connect` again — once SSO is re-authorized, this account will keep working.",
+		ssoURL,
+	)
+	p.CreateBotDMPost(info.UserID, msg, "custom_git_saml_reauth")
+}
+
+// extractSSOAuthorizeURL pulls the SSO reauthorization url that GitHub
+// includes in the X-GitHub-SSO header on SAML enforced 403 responses
+func extractSSOAuthorizeURL(err error) string {
+	var ghErr *github.ErrorResponse
+	if !errors.As(err, &ghErr) || ghErr.Response == nil {
+		return ""
+	}
+	header := ghErr.Response.Header.Get("X-GitHub-SSO")
+	if header == "" {
+		return ""
+	}
+	for part := range strings.SplitSeq(header, ";") {
+		part = strings.TrimSpace(part)
+		if u, ok := strings.CutPrefix(part, "url="); ok {
+			return strings.TrimSpace(u)
+		}
+	}
+	return ""
+}

server/plugin/plugin_test.go

@@ -470,6 +470,36 @@ func TestIsGitHubAuthFailure(t *testing.T) {
 	})
 }
 
+func TestExtractSSOAuthorizeURL(t *testing.T) {
+	makeErr := func(header string) error {
+		h := http.Header{}
+		if header != "" {
+			h.Set("X-GitHub-SSO", header)
+		}
+		return &github.ErrorResponse{
+			Response: &http.Response{StatusCode: http.StatusForbidden, Header: h, Request: &http.Request{}},
+		}
+	}
+	tests := []struct {
+		name string
+		err  error
+		want string
+	}{
+		{"nil", nil, ""},
+		{"non-github error", errors.New("boom"), ""},
+		{"no header", makeErr(""), ""},
+		{"required with url", makeErr("required; url=https://github.com/orgs/foo/sso?authorization_request=abc"), "https://github.com/orgs/foo/sso?authorization_request=abc"},
+		{"partial-results with url", makeErr("partial-results; url=https://github.com/orgs/foo/sso"), "https://github.com/orgs/foo/sso"},
+		{"header without url", makeErr("required"), ""},
+		{"whitespace tolerant", makeErr("required ;  url=https://github.com/orgs/foo/sso  "), "https://github.com/orgs/foo/sso"},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.want, extractSSOAuthorizeURL(tc.err))
+		})
+	}
+}
+
 func connectedGitHubUserInfo(t *testing.T) *GitHubUserInfo {
 	t.Helper()
 	encryptedToken, err := encrypt([]byte(testNewKey), MockAccessToken)
@@ -510,21 +540,39 @@ func expectRevokedTokenNotification(api *plugintest.API, mockKvStore *mocks.Mock
 	})).Return(&model.Post{}, nil).Once()
 }
 
+func expectSAMLReauthorizeNotification(api *plugintest.API, mockKvStore *mocks.MockKvStore, userInfo *GitHubUserInfo) {
+	mockKvStore.EXPECT().Set(authFailureDMKey+userInfo.UserID, gomock.Any(), gomock.Any(), gomock.Any()).Return(true, nil)
+	api.On("GetDirectChannel", userInfo.UserID, MockBotID).Return(&model.Channel{Id: "dmchannel"}, nil)
+	api.On("CreatePost", mock.MatchedBy(func(post *model.Post) bool {
+		return post.UserId == MockBotID &&
+			post.ChannelId == "dmchannel" &&
+			post.Type == "custom_git_saml_reauth" &&
+			strings.Contains(post.Message, "https://github.com/orgs/foo/sso")
+	})).Return(&model.Post{}, nil).Once()
+}
+
 func TestUseGitHubClient_AuthFailureNotifiesUser(t *testing.T) {
 	samlGraphQLErr := errors.New("error in executing query: GraphQL: Resource protected by organization SAML enforcement. You must grant your OAuth token access to this organization.")
 
+	const (
+		expectNone     = ""
+		expectRevoked  = "revoked"
+		expectReauthDM = "reauth"
+	)
+
 	tests := []struct {
 		name   string
 		err    error
-		notify bool
+		expect string
 	}{
 		{
 			name:   "401 bad credentials",
 			err:    errors.New(invalidTokenError),
-			notify: true,
+			expect: expectRevoked,
 		},
 		{
-			name: "403 SAML REST",
+			// SSO URL is available: keep account connected, DM the reauthorize link.
+			name: "403 SAML REST with SSO URL",
 			err: &github.ErrorResponse{
 				Message: "Resource protected by organization SAML enforcement. You must grant your OAuth token access to this organization.",
 				Response: &http.Response{
@@ -533,25 +581,26 @@ func TestUseGitHubClient_AuthFailureNotifiesUser(t *testing.T) {
 					Request:    &http.Request{},
 				},
 			},
-			notify: true,
+			expect: expectReauthDM,
 		},
 		{
+			// GraphQL error string carries no header we can parse — fall back to disconnect.
 			name:   "403 SAML graphql",
 			err:    samlGraphQLErr,
-			notify: true,
+			expect: expectRevoked,
 		},
 		{
 			name: "403 unrelated",
 			err: &github.ErrorResponse{
 				Message:  "Forbidden",
 				Response: &http.Response{StatusCode: http.StatusForbidden, Request: &http.Request{}},
 			},
-			notify: false,
+			expect: expectNone,
 		},
 		{
 			name:   "generic error",
 			err:    errors.New("connection reset"),
-			notify: false,
+			expect: expectNone,
 		},
 	}
 
@@ -561,8 +610,11 @@ func TestUseGitHubClient_AuthFailureNotifiesUser(t *testing.T) {
 			defer ctrl.Finish()
 
 			userInfo := connectedGitHubUserInfo(t)
-			if tc.notify {
+			switch tc.expect {
+			case expectRevoked:
 				expectRevokedTokenNotification(api, mockKvStore, userInfo)
+			case expectReauthDM:
+				expectSAMLReauthorizeNotification(api, mockKvStore, userInfo)
 			}
 
 			err := p.useGitHubClient(userInfo, func(_ *GitHubUserInfo, _ *oauth2.Token) error {