Merge pull request #2805 from nobuyo/skip-redirection-to-approval-end…

…point

Skip redirection to approval when it is not required (#2686)

server/handlers.go

@@ -372,13 +372,24 @@ func (s *Server) handlePasswordLogin(w http.ResponseWriter, r *http.Request) {
 			}
 			return
 		}
-		redirectURL, err := s.finalizeLogin(identity, authReq, conn.Connector)
+		redirectURL, canSkipApproval, err := s.finalizeLogin(identity, authReq, conn.Connector)
 		if err != nil {
 			s.logger.Errorf("Failed to finalize login: %v", err)
 			s.renderError(r, w, http.StatusInternalServerError, "Login error.")
 			return
 		}
 
+		if canSkipApproval {
+			authReq, err = s.storage.GetAuthRequest(authReq.ID)
+			if err != nil {
+				s.logger.Errorf("Failed to get finalized auth request: %v", err)
+				s.renderError(r, w, http.StatusInternalServerError, "Login error.")
+				return
+			}
+			s.sendCodeResponse(w, r, authReq)
+			return
+		}
+
 		http.Redirect(w, r, redirectURL, http.StatusSeeOther)
 	default:
 		s.renderError(r, w, http.StatusBadRequest, "Unsupported request method.")
@@ -460,19 +471,30 @@ func (s *Server) handleConnectorCallback(w http.ResponseWriter, r *http.Request)
 		return
 	}
 
-	redirectURL, err := s.finalizeLogin(identity, authReq, conn.Connector)
+	redirectURL, canSkipApproval, err := s.finalizeLogin(identity, authReq, conn.Connector)
 	if err != nil {
 		s.logger.Errorf("Failed to finalize login: %v", err)
 		s.renderError(r, w, http.StatusInternalServerError, "Login error.")
 		return
 	}
 
+	if canSkipApproval {
+		authReq, err = s.storage.GetAuthRequest(authReq.ID)
+		if err != nil {
+			s.logger.Errorf("Failed to get finalized auth request: %v", err)
+			s.renderError(r, w, http.StatusInternalServerError, "Login error.")
+			return
+		}
+		s.sendCodeResponse(w, r, authReq)
+		return
+	}
+
 	http.Redirect(w, r, redirectURL, http.StatusSeeOther)
 }
 
 // finalizeLogin associates the user's identity with the current AuthRequest, then returns
 // the approval page's path.
-func (s *Server) finalizeLogin(identity connector.Identity, authReq storage.AuthRequest, conn connector.Connector) (string, error) {
+func (s *Server) finalizeLogin(identity connector.Identity, authReq storage.AuthRequest, conn connector.Connector) (string, bool, error) {
 	claims := storage.Claims{
 		UserID:            identity.UserID,
 		Username:          identity.Username,
@@ -489,7 +511,7 @@ func (s *Server) finalizeLogin(identity connector.Identity, authReq storage.Auth
 		return a, nil
 	}
 	if err := s.storage.UpdateAuthRequest(authReq.ID, updater); err != nil {
-		return "", fmt.Errorf("failed to update auth request: %v", err)
+		return "", false, fmt.Errorf("failed to update auth request: %v", err)
 	}
 
 	email := claims.Email
@@ -500,7 +522,10 @@ func (s *Server) finalizeLogin(identity connector.Identity, authReq storage.Auth
 	s.logger.Infof("login successful: connector %q, username=%q, preferred_username=%q, email=%q, groups=%q",
 		authReq.ConnectorID, claims.Username, claims.PreferredUsername, email, claims.Groups)
 
-	// TODO: if s.skipApproval or !authReq.ForceApprovalPrompt, we can skip the redirect to /approval and go ahead and send code
+	// we can skip the redirect to /approval and go ahead and send code if it's not required
+	if s.skipApproval || !authReq.ForceApprovalPrompt {
+		return "", true, nil
+	}
 
 	// an HMAC is used here to ensure that the request ID is unpredictable, ensuring that an attacker who intercepted the original
 	// flow would be unable to poll for the result at the /approval endpoint
@@ -511,15 +536,15 @@ func (s *Server) finalizeLogin(identity connector.Identity, authReq storage.Auth
 	returnURL := path.Join(s.issuerURL.Path, "/approval") + "?req=" + authReq.ID + "&hmac=" + base64.RawURLEncoding.EncodeToString(mac)
 	_, ok := conn.(connector.RefreshConnector)
 	if !ok {
-		return returnURL, nil
+		return returnURL, false, nil
 	}
 
 	// Try to retrieve an existing OfflineSession object for the corresponding user.
 	session, err := s.storage.GetOfflineSessions(identity.UserID, authReq.ConnectorID)
 	if err != nil {
 		if err != storage.ErrNotFound {
 			s.logger.Errorf("failed to get offline session: %v", err)
-			return "", err
+			return "", false, err
 		}
 		offlineSessions := storage.OfflineSessions{
 			UserID:        identity.UserID,
@@ -532,10 +557,10 @@ func (s *Server) finalizeLogin(identity connector.Identity, authReq storage.Auth
 		// the newly received refreshtoken.
 		if err := s.storage.CreateOfflineSessions(offlineSessions); err != nil {
 			s.logger.Errorf("failed to create offline session: %v", err)
-			return "", err
+			return "", false, err
 		}
 
-		return returnURL, nil
+		return returnURL, false, nil
 	}
 
 	// Update existing OfflineSession obj with new RefreshTokenRef.
@@ -546,10 +571,10 @@ func (s *Server) finalizeLogin(identity connector.Identity, authReq storage.Auth
 		return old, nil
 	}); err != nil {
 		s.logger.Errorf("failed to update offline session: %v", err)
-		return "", err
+		return "", false, err
 	}
 
-	return returnURL, nil
+	return returnURL, false, nil
 }
 
 func (s *Server) handleApproval(w http.ResponseWriter, r *http.Request) {
@@ -588,6 +613,8 @@ func (s *Server) handleApproval(w http.ResponseWriter, r *http.Request) {
 
 	switch r.Method {
 	case http.MethodGet:
+		// TODO: `finalizeLogin()` now sends code directly to client without going through this endpoint,
+		//		 the `if skipApproval { ... }` block needs to be removed after a grace period.
 		if s.skipApproval {
 			s.sendCodeResponse(w, r, authReq)
 			return

server/handlers_test.go

@@ -5,10 +5,12 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"net/url"
 	"path"
+	"strings"
 	"testing"
 	"time"
 
@@ -323,3 +325,181 @@ func TestHandlePassword(t *testing.T) {
 		require.Equal(t, `{"test": "true"}`, string(newSess.ConnectorData))
 	}
 }
+
+func TestHandlePasswordLoginWithSkipApproval(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	connID := "mockPw"
+	authReqID := "test"
+	expiry := time.Now().Add(100 * time.Second)
+	resTypes := []string{"code"}
+
+	tests := []struct {
+		name         string
+		skipApproval bool
+		authReq      storage.AuthRequest
+	}{
+		{
+			name:         "Force approval",
+			skipApproval: false,
+			authReq: storage.AuthRequest{
+				ID:                  authReqID,
+				ConnectorID:         connID,
+				RedirectURI:         "cb",
+				Expiry:              expiry,
+				ResponseTypes:       resTypes,
+				ForceApprovalPrompt: true,
+			},
+		},
+		{
+			name:         "Skip approval by server config",
+			skipApproval: true,
+			authReq: storage.AuthRequest{
+				ID:                  authReqID,
+				ConnectorID:         connID,
+				RedirectURI:         "cb",
+				Expiry:              expiry,
+				ResponseTypes:       resTypes,
+				ForceApprovalPrompt: true,
+			},
+		},
+		{
+			name:         "Skip approval by auth request",
+			skipApproval: false,
+			authReq: storage.AuthRequest{
+				ID:                  authReqID,
+				ConnectorID:         connID,
+				RedirectURI:         "cb",
+				Expiry:              expiry,
+				ResponseTypes:       resTypes,
+				ForceApprovalPrompt: false,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		httpServer, s := newTestServer(ctx, t, func(c *Config) {
+			c.SkipApprovalScreen = tc.skipApproval
+			c.Now = time.Now
+		})
+		defer httpServer.Close()
+
+		sc := storage.Connector{
+			ID:              connID,
+			Type:            "mockPassword",
+			Name:            "MockPassword",
+			ResourceVersion: "1",
+			Config:          []byte("{\"username\": \"foo\", \"password\": \"password\"}"),
+		}
+		if err := s.storage.CreateConnector(sc); err != nil {
+			t.Fatalf("create connector: %v", err)
+		}
+		if _, err := s.OpenConnector(sc); err != nil {
+			t.Fatalf("open connector: %v", err)
+		}
+		if err := s.storage.CreateAuthRequest(tc.authReq); err != nil {
+			t.Fatalf("failed to create AuthRequest: %v", err)
+		}
+
+		rr := httptest.NewRecorder()
+
+		path := fmt.Sprintf("/auth/%s/login?state=%s&back=&login=foo&password=password", connID, authReqID)
+		s.handlePasswordLogin(rr, httptest.NewRequest("POST", path, nil))
+
+		require.Equal(t, 303, rr.Code)
+
+		resp := rr.Result()
+		cbPath := strings.Split(resp.Header.Get("Location"), "?")[0]
+
+		if tc.skipApproval || !tc.authReq.ForceApprovalPrompt {
+			require.Equal(t, "/auth/mockPw/cb", cbPath)
+		} else {
+			require.Equal(t, "/approval", cbPath)
+		}
+
+		resp.Body.Close()
+	}
+}
+
+func TestHandleConnectorCallbackWithSkipApproval(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	connID := "mock"
+	authReqID := "test"
+	expiry := time.Now().Add(100 * time.Second)
+	resTypes := []string{"code"}
+
+	tests := []struct {
+		name         string
+		skipApproval bool
+		authReq      storage.AuthRequest
+	}{
+		{
+			name:         "Force approval",
+			skipApproval: false,
+			authReq: storage.AuthRequest{
+				ID:                  authReqID,
+				ConnectorID:         connID,
+				RedirectURI:         "cb",
+				Expiry:              expiry,
+				ResponseTypes:       resTypes,
+				ForceApprovalPrompt: true,
+			},
+		},
+		{
+			name:         "Skip approval by server config",
+			skipApproval: true,
+			authReq: storage.AuthRequest{
+				ID:                  authReqID,
+				ConnectorID:         connID,
+				RedirectURI:         "cb",
+				Expiry:              expiry,
+				ResponseTypes:       resTypes,
+				ForceApprovalPrompt: true,
+			},
+		},
+		{
+			name:         "Skip approval by auth request",
+			skipApproval: false,
+			authReq: storage.AuthRequest{
+				ID:                  authReqID,
+				ConnectorID:         connID,
+				RedirectURI:         "cb",
+				Expiry:              expiry,
+				ResponseTypes:       resTypes,
+				ForceApprovalPrompt: false,
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		httpServer, s := newTestServer(ctx, t, func(c *Config) {
+			c.SkipApprovalScreen = tc.skipApproval
+			c.Now = time.Now
+		})
+		defer httpServer.Close()
+
+		if err := s.storage.CreateAuthRequest(tc.authReq); err != nil {
+			t.Fatalf("failed to create AuthRequest: %v", err)
+		}
+		rr := httptest.NewRecorder()
+
+		path := fmt.Sprintf("/callback/%s?state=%s", connID, authReqID)
+		s.handleConnectorCallback(rr, httptest.NewRequest("GET", path, nil))
+
+		require.Equal(t, 303, rr.Code)
+
+		resp := rr.Result()
+		cbPath := strings.Split(resp.Header.Get("Location"), "?")[0]
+
+		if tc.skipApproval || !tc.authReq.ForceApprovalPrompt {
+			require.Equal(t, "/callback/cb", cbPath)
+		} else {
+			require.Equal(t, "/approval", cbPath)
+		}
+
+		resp.Body.Close()
+	}
+}

server/server_test.go

@@ -98,6 +98,7 @@ func newTestServer(ctx context.Context, t *testing.T, updateConfig func(c *Confi
 		Logger:             logger,
 		PrometheusRegistry: prometheus.NewRegistry(),
 		HealthChecker:      gosundheit.New(),
+		SkipApprovalScreen: true, // Don't prompt for approval, just immediately redirect with code.
 	}
 	if updateConfig != nil {
 		updateConfig(&config)
@@ -118,7 +119,6 @@ func newTestServer(ctx context.Context, t *testing.T, updateConfig func(c *Confi
 	if server, err = newServer(ctx, config, staticRotationStrategy(testKey)); err != nil {
 		t.Fatal(err)
 	}
-	server.skipApproval = true // Don't prompt for approval, just immediately redirect with code.
 
 	// Default rotation policy
 	if server.refreshTokenPolicy == nil {