iOS AWS Cognito Custom Login Provider 2.4+

usage.m

    _identityProviderManager = [YOLOAuthenticationIdentityProviderManager new];
    _identityProviderManager.token = /* current token */
    _identityProvider = [[YOLOIdentityProvider alloc] initWithRegionType:/* current region type */
                                                          identityPoolId:/* current pool id */
                                                         useEnhancedFlow:NO
                                                 identityProviderManager:_identityProviderManager];

    _credentialsProvider = [[AWSCognitoCredentialsProvider alloc] initWithRegionType:regionType
                                                                    identityProvider:_identityProvider];

    _configuration = [[AWSServiceConfiguration alloc] initWithRegion:regionType
                                                 credentialsProvider:self.credentialsProvider];

    [AWSServiceManager defaultServiceManager].defaultServiceConfiguration = self.configuration;

YOLOAuthenticationIdentityProviderManager.m

@interface YOLOAuthenticationIdentityProviderManager : NSObject <AWSIdentityProviderManager>

@property (nonatomic, strong, nullable) NSString *token;

@end

@implementation YOLOAuthenticationIdentityProviderManager

- (AWSTask<NSDictionary<NSString *, NSString *> *> *)logins {
    if (self.token == nil) {
        return [AWSTask taskWithResult:nil];
    }
    
    return [AWSTask taskWithResult:@{ /* YOLO-PROVIDER-NAME-IN-COGNITO-CONSOLE */: self.token }];
}

@end

YOLOIdentityProvider.m

@protocol YOLOIdentityResponse <NSObject>

@property (nonatomic, strong) NSString *identityId;
@property (nonatomic, strong) NSString *token;

@end


@protocol YOLOAuthenticationClient <NSObject>

- (AWSTask<id<YOLOIdentityResponse>> *)getToken:(nullable NSString *)identityId
                                         logins:(nonnull NSDictionary *)logins;
                                         
@end


@interface YOLOIdentityProvider : AWSCognitoCredentialsProviderHelper

@property (nonatomic, strong) id<YOLOAuthenticationClient> *authClient;
@property (nonatomic, strong) NSDictionary<NSString *, NSString *> *authorizedLogins;

@end


@implementation YOLOIdentityProvider

- (instancetype)initWithRegionType:(AWSRegionType)regionType
                    identityPoolId:(NSString *)identityPoolId
                   useEnhancedFlow:(BOOL)useEnhancedFlow
           identityProviderManager:(nullable id<AWSIdentityProviderManager>)identityProviderManager {
    
    if (self = [super initWithRegionType:regionType
                          identityPoolId:identityPoolId
                         useEnhancedFlow:useEnhancedFlow
                 identityProviderManager:identityProviderManager]) {
        _authClient = /* YOU CUSTOM AUTH CLIENT */;
    }

    return self;
}

- (AWSTask<NSString *> *)token {
    return [[self getLogins] continueWithSuccessBlock:^id _Nullable(AWSTask<NSDictionary<NSString *,NSString *> *> * _Nonnull task) {
        NSDictionary<NSString *,NSString *> *logins = task.result;

        if (logins.count == 0) {
            return [self superToken];
        }

        return [[self.authClient getToken:self.identityId logins:logins] continueWithSuccessBlock:^id _Nullable(AWSTask * _Nonnull task) {
            id<YOLOIdentityResponse> *response = task.result;
            return response.token;
        }];
    }];
}

- (AWSTask<NSString *> *)superToken {
    return [super token];
}

- (AWSTask<NSString *> *)getIdentityId {
    return [[self getLogins] continueWithSuccessBlock:^id _Nullable(AWSTask<NSDictionary<NSString *,NSString *> *> * _Nonnull task) {
        NSDictionary<NSString *,NSString *> *logins = task.result;

        if (logins.count == 0) {
            return [self superGetIdentityId];
        }

        if (self.identityId != nil && [self.authorizedLogins isEqualToDictionary:logins]) {
            return self.identityId;
        }

        self.authorizedLogins = logins;

        return [[self.authClient getToken:self.identityId logins:logins] continueWithSuccessBlock:^id _Nullable(AWSTask * _Nonnull task) {
            id<YOLOIdentityResponse> *response = task.result;
            self.identityId = response.identityId;
            return response.identityId;
        }];
    }];
}

- (AWSTask<NSString *> *)superGetIdentityId {
    return [super getIdentityId];
}

- (AWSTask<NSDictionary<NSString *, NSString *> *> *)getLogins {
    return [self.identityProviderManager logins];
}

@end