Fix : Context timeout error handler changes (#6493)

* context timeout error handler changes

* changes updated

* oidc test case added

internal/daemon/controller/handlers/authmethods/oidc.go

@@ -269,6 +269,11 @@ func (s Service) authenticateOidcCallback(ctx context.Context, req *pbs.Authenti
 		attrs.GetState(),
 		attrs.GetCode())
 	if err != nil {
+		// Check if the error is due to context timeout.
+		// This typically happens when the OIDC provider takes too long to respond
+		if errors.Is(err, context.DeadlineExceeded) {
+			return nil, handlers.ApiErrorWithCodeAndMessage(codes.DeadlineExceeded, "The identity provider took too long to respond. Please try again.")
+		}
 		return errResponse(errors.New(ctx, errors.InvalidParameter, op, "Callback validation failed.", errors.WithWrap(err)))
 	}
 

internal/daemon/controller/handlers/authmethods/oidc_test.go

@@ -11,7 +11,9 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"regexp"
+	"runtime/debug"
 	"sort"
+	"strings"
 	"testing"
 	"time"
 
@@ -2048,3 +2050,45 @@ func TestAuthenticate_OIDC_Callback_ErrorRedirect(t *testing.T) {
 		})
 	}
 }
+
+func TestAuthenticate_OIDC_Callback_DeadlineExceeded(t *testing.T) {
+	s := getSetup(t)
+
+	oidcRepoFn := func() (*oidc.Repository, error) {
+		if strings.Contains(string(debug.Stack()), "internal/auth/oidc.Callback") {
+			return nil, context.DeadlineExceeded
+		}
+
+		return s.oidcRepoFn()
+	}
+
+	tested, err := authmethods.NewService(
+		s.ctx,
+		s.kmsCache,
+		s.pwRepoFn,
+		oidcRepoFn,
+		s.iamRepoFn,
+		s.atRepoFn,
+		s.ldapRepoFn,
+		s.authMethodRepoFn,
+		uint(s.maxPageSize),
+	)
+	require.NoError(t, err)
+
+	_, err = tested.Authenticate(auth.DisabledAuthTestContext(s.iamRepoFn, s.org.GetPublicId()),
+		&pbs.AuthenticateRequest{
+			AuthMethodId: s.authMethod.GetPublicId(),
+			Command:      "callback",
+			Attrs: &pbs.AuthenticateRequest_OidcAuthMethodAuthenticateCallbackRequest{
+				OidcAuthMethodAuthenticateCallbackRequest: &pb.OidcAuthMethodAuthenticateCallbackRequest{
+					Code:  "code",
+					State: "state",
+				},
+			},
+		},
+	)
+
+	require.Error(t, err)
+	assert.True(t, errors.Is(err, handlers.ApiErrorWithCode(codes.DeadlineExceeded)))
+	assert.Equal(t, "The identity provider took too long to respond. Please try again.", handlers.ToApiError(err).GetMessage())
+}

internal/daemon/controller/handlers/errors.go

@@ -198,6 +198,15 @@ func backendErrorToApiError(inErr error) *ApiError {
 	stErr := status.Convert(inErr)
 
 	switch {
+	case errors.Is(inErr, context.DeadlineExceeded):
+		// Context deadline exceeded (timeout) should return HTTP 504 Gateway Timeout
+		return &ApiError{
+			Status: http.StatusGatewayTimeout,
+			Inner: &pb.Error{
+				Kind:    codes.DeadlineExceeded.String(),
+				Message: http.StatusText(http.StatusGatewayTimeout),
+			},
+		}
 	case errors.Is(inErr, runtime.ErrNotMatch):
 		// grpc gateway uses this error when the path was not matched, but the error uses codes.Unimplemented which doesn't match the intention.
 		// Overwrite the error to match our expected behavior.

internal/daemon/controller/handlers/errors_test.go

@@ -85,6 +85,17 @@ func TestApiErrorHandler(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "Context Deadline Exceeded",
+			err:  context.DeadlineExceeded,
+			expected: ApiError{
+				Status: http.StatusGatewayTimeout,
+				Inner: &pb.Error{
+					Kind:    "DeadlineExceeded",
+					Message: http.StatusText(http.StatusGatewayTimeout),
+				},
+			},
+		},
 		{
 			name: "Unimplemented error",
 			err:  status.Error(codes.Unimplemented, "Test"),

internal/daemon/controller/handlers/targets/target_service.go

@@ -1207,12 +1207,12 @@ func (s Service) getFromRepo(ctx context.Context, id string) (target.Target, []t
 		}
 		return nil, nil, nil, err
 	}
-	hs := u.GetHostSources()
-	cl := u.GetCredentialSources()
-
 	if u == nil {
 		return nil, nil, nil, handlers.NotFoundErrorf("Target %q doesn't exist.", id)
 	}
+	hs := u.GetHostSources()
+	cl := u.GetCredentialSources()
+
 	return u, hs, cl, nil
 }