tech.ydb.core.auth.StaticCredentials Maven / Gradle / Ivy

Go to download
package tech.ydb.core.auth;

import java.time.Clock;
import java.time.Instant;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicReference;

import com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import tech.ydb.auth.AuthRpcProvider;
import tech.ydb.core.impl.auth.GrpcAuthRpc;
import tech.ydb.proto.auth.YdbAuth;

/**
 *
 * @author Aleksandr Gorshenin
 */
public class StaticCredentials implements AuthRpcProvider {
    private static final Logger logger = LoggerFactory.getLogger(StaticCredentials.class);

    private final Clock clock;
    private final YdbAuth.LoginRequest request;

    @VisibleForTesting
    StaticCredentials(Clock clock, String username, String password) {
        this.clock = clock;
        YdbAuth.LoginRequest.Builder builder = YdbAuth.LoginRequest.newBuilder()
                .setUser(username);
        if (password != null) {
            builder.setPassword(password);
        }
        this.request = builder.build();
    }

    public StaticCredentials(String username, String password) {
        this(Clock.systemUTC(), username, password);
    }

    @Override
    public tech.ydb.auth.AuthIdentity createAuthIdentity(GrpcAuthRpc rpc) {
        logger.info("create static identity for database {}", rpc.getDatabase());
        return new IdentityImpl(rpc);
    }

    private interface State {
        void init();
        State validate(Instant now);
        String token();
    }

    private class IdentityImpl implements tech.ydb.auth.AuthIdentity {
        private final AtomicReference state = new AtomicReference<>(new NullState());
        private final StaticCredentialsRpc rpc;

        IdentityImpl(GrpcAuthRpc authRpc) {
            this.rpc = new StaticCredentialsRpc(authRpc, request, clock);
        }

        private State updateState(State current, State next) {
            if (state.compareAndSet(current, next)) {
                next.init();
            }
            return state.get();
        }

        @Override
        public String getToken() {
            return state.get().validate(clock.instant()).token();
        }

        private class NullState implements State {
            @Override
            public void init() {
                // Nothing
            }

            @Override
            public String token() {
                throw new IllegalStateException("Get token for null state");
            }

            @Override
            public State validate(Instant now) {
                return updateState(this, new SyncLogin()).validate(now);
            }
        }

        private class SyncLogin implements State {
            private final CompletableFuture future = new CompletableFuture<>();

            @Override
            public void init() {
                rpc.loginAsync().whenComplete((token, th) -> {
                    if (token != null) {
                        future.complete(new LoggedInState(token));
                    } else {
                        future.complete(new ErrorState(th));
                    }
                });
            }

            @Override
            public String token() {
                throw new IllegalStateException("Get token for unfinished sync state");
            }

            @Override
            public State validate(Instant now) {
                return updateState(this, rpc.unwrap(future));
            }
        }

        private class BackgroundLogin implements State {
            private final StaticCredentialsRpc.Token token;
            private final CompletableFuture future = new CompletableFuture<>();

            BackgroundLogin(StaticCredentialsRpc.Token token) {
                this.token = token;
            }

            @Override
            public void init() {
                rpc.loginAsync().whenComplete((nextToken, th) -> {
                    if (nextToken != null) {
                        future.complete(new LoggedInState(nextToken));
                    } else {
                        future.completeExceptionally(th);
                    }
                });
            }

            @Override
            public String token() {
                return token.token();
            }

            @Override
            public State validate(Instant now) {
                if (future.isCompletedExceptionally()) {
                    if (now.isAfter(token.expiredAt())) {
                        // If token had already expired, switch to sync mode and wait for finishing
                        return updateState(this, new SyncLogin()).validate(now);
                    }
                    // else retry background login
                    return updateState(this, new BackgroundLogin(token));
                }

                if (future.isDone()) {
                    return updateState(this, future.join());
                }

                return this;
            }
        }

        private class LoggedInState implements State {
            private final StaticCredentialsRpc.Token token;

            LoggedInState(StaticCredentialsRpc.Token token) {
                this.token = token;
                logger.debug("logged in with expired at {} and updating at {}", token.expiredAt(), token.updateAt());
            }

            @Override
            public void init() { }

            @Override
            public String token() {
                return token.token();
            }

            @Override
            public State validate(Instant now) {
                if (now.isAfter(token.expiredAt())) {
                    // If token had already expired, switch to sync mode and wait for finishing
                    return updateState(this, new SyncLogin()).validate(now);
                }
                if (now.isAfter(token.updateAt())) {
                    return updateState(this, new BackgroundLogin(token));
                }
                return this;
            }
        }

        private class ErrorState implements State {
            private final RuntimeException ex;

            ErrorState(Throwable ex) {
                this.ex = ex instanceof RuntimeException ?
                        (RuntimeException) ex : new RuntimeException("can't login", ex);
            }

            @Override
            public void init() { }

            @Override
            public String token() {
                throw ex;
            }

            @Override
            public State validate(Instant instant) {
                return this;
            }
        }
    }
}