From e6bc6564e97bc791de82aed22fa68db7aa69c426 Mon Sep 17 00:00:00 2001
From: Alexandr Gorshenin <alexandr268@ydb.tech>
Date: Sun, 27 Oct 2024 10:20:10 +0000
Subject: [PATCH] Added environment auth provider with runtime class loading

---
 .../java/tech/ydb/auth/OAuth2AuthHelper.java  | 20 +++++
 .../ydb/core/auth/EnvironAuthProvider.java    | 87 +++++++++++++++++++
 .../java/tech/ydb/core/auth/JwtUtils.java     |  2 +-
 3 files changed, 108 insertions(+), 1 deletion(-)
 create mode 100644 auth-providers/oauth2-provider/src/main/java/tech/ydb/auth/OAuth2AuthHelper.java
 create mode 100644 core/src/main/java/tech/ydb/core/auth/EnvironAuthProvider.java

diff --git a/auth-providers/oauth2-provider/src/main/java/tech/ydb/auth/OAuth2AuthHelper.java b/auth-providers/oauth2-provider/src/main/java/tech/ydb/auth/OAuth2AuthHelper.java
new file mode 100644
index 00000000..90b9b44c
--- /dev/null
+++ b/auth-providers/oauth2-provider/src/main/java/tech/ydb/auth/OAuth2AuthHelper.java
@@ -0,0 +1,20 @@
+package tech.ydb.auth;
+
+import java.nio.file.Path;
+
+import tech.ydb.core.impl.auth.GrpcAuthRpc;
+
+/**
+ *
+ * @author Aleksandr Gorshenin
+ */
+public class OAuth2AuthHelper {
+    private OAuth2AuthHelper() { }
+
+    public static AuthIdentity configFileIdentity(Path file, GrpcAuthRpc rpc) {
+        return OAuth2TokenExchangeProvider
+                .fromFile(file.toFile())
+                .build()
+                .createAuthIdentity(rpc);
+    }
+}
diff --git a/core/src/main/java/tech/ydb/core/auth/EnvironAuthProvider.java b/core/src/main/java/tech/ydb/core/auth/EnvironAuthProvider.java
new file mode 100644
index 00000000..6eceb5b0
--- /dev/null
+++ b/core/src/main/java/tech/ydb/core/auth/EnvironAuthProvider.java
@@ -0,0 +1,87 @@
+package tech.ydb.core.auth;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
+import tech.ydb.auth.AuthIdentity;
+import tech.ydb.auth.AuthRpcProvider;
+import tech.ydb.core.impl.auth.GrpcAuthRpc;
+
+/**
+ *
+ * @author Aleksandr Gorshenin
+ */
+public class EnvironAuthProvider implements AuthRpcProvider<GrpcAuthRpc> {
+    private static final String IAM_CLASS_NAME = "tech.ydb.auth.iam.CloudAuthIdentity";
+    private static final String IAM_CLASS_ERROR = "Cannot find CloudAuthIdentity class, "
+            + "you have to add tech.ydb.auth:yc-auth-provider artifact to classpath";
+
+    private static final String OAUTH2_CLASS_NAME = "tech.ydb.auth.OAuth2AuthHelper";
+    private static final String OAUTH2_CLASS_ERROR = "Cannot find OAuth2AuthHelper class, "
+            + "you have to add tech.ydb.auth:ydb-oauth2-provider artifact to classpath";
+
+    @Override
+    public AuthIdentity createAuthIdentity(GrpcAuthRpc rpc) {
+        String anonCredentials = System.getenv("YDB_ANONYMOUS_CREDENTIALS");
+        if (anonCredentials != null && anonCredentials.equals("1")) {
+            return null;
+        }
+
+        String saKeyFile = System.getenv("YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS");
+        if (saKeyFile != null) {
+            return loadServiceAccountIdentity(Paths.get(saKeyFile));
+        }
+
+        String metadataCredentials = System.getenv("YDB_METADATA_CREDENTIALS");
+        if (metadataCredentials != null && metadataCredentials.equals("1")) {
+            return loadMetadataIdentity();
+        }
+
+        String accessToken = System.getenv("YDB_ACCESS_TOKEN_CREDENTIALS");
+        if (accessToken != null) {
+            return () -> accessToken;
+        }
+
+        String oauth2KeyFile = System.getenv("YDB_OAUTH2_KEY_FILE");
+        if (oauth2KeyFile != null) {
+            return loadOAuth2KeyProvider(rpc, Paths.get(oauth2KeyFile));
+        }
+
+        return loadMetadataIdentity();
+    }
+
+    private static AuthIdentity loadServiceAccountIdentity(Path saKeyFile) {
+        try {
+            Class<?> clazz = Class.forName(IAM_CLASS_NAME);
+            Method method = clazz.getMethod("serviceAccountIdentity", Path.class, String.class);
+            return (AuthIdentity) method.invoke(null, saKeyFile, null);
+        } catch (ClassNotFoundException | NoSuchMethodException | SecurityException |
+                IllegalAccessException | IllegalArgumentException | InvocationTargetException ex) {
+            throw new RuntimeException(IAM_CLASS_ERROR, ex);
+        }
+    }
+
+    private static AuthIdentity loadMetadataIdentity() {
+        try {
+            Class<?> clazz = Class.forName(IAM_CLASS_NAME);
+            Method method = clazz.getMethod("metadataIdentity", String.class);
+            return (AuthIdentity) method.invoke(null, (Object) null);
+        } catch (ClassNotFoundException | NoSuchMethodException | SecurityException |
+                IllegalAccessException | IllegalArgumentException | InvocationTargetException ex) {
+            throw new RuntimeException(IAM_CLASS_ERROR, ex);
+        }
+    }
+
+    private static AuthIdentity loadOAuth2KeyProvider(GrpcAuthRpc rpc, Path configFile) {
+        try {
+            Class<?> clazz = Class.forName(OAUTH2_CLASS_NAME);
+            Method method = clazz.getMethod("configFileIdentity", Path.class, GrpcAuthRpc.class);
+            return (AuthIdentity) method.invoke(null, configFile, rpc);
+        } catch (ClassNotFoundException | NoSuchMethodException | SecurityException |
+                IllegalAccessException | IllegalArgumentException | InvocationTargetException ex) {
+            throw new RuntimeException(OAUTH2_CLASS_ERROR, ex);
+        }
+    }
+}
diff --git a/core/src/main/java/tech/ydb/core/auth/JwtUtils.java b/core/src/main/java/tech/ydb/core/auth/JwtUtils.java
index e82e24b5..73f70edf 100644
--- a/core/src/main/java/tech/ydb/core/auth/JwtUtils.java
+++ b/core/src/main/java/tech/ydb/core/auth/JwtUtils.java
@@ -14,7 +14,7 @@
  *
  * @author Aleksandr Gorshenin
  */
-public class JwtUtils {
+class JwtUtils {
     private static final Logger logger = LoggerFactory.getLogger(JwtUtils.class);
 
     private static final char SEPARATOR_CHAR = '.';