DRILL-8558: Fix credential leak and cross-user OAuth token reuse

contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/DrillDataStore.java

@@ -41,7 +41,12 @@ public class DrillDataStore<V extends Serializable> extends AbstractMemoryDataSt
   DrillDataStore(OAuthTokenProvider tokenProvider, String pluginName, String userID, DrillDataStoreFactory dataStoreFactory) {
     super(dataStoreFactory, userID);
     this.drillDataStoreFactory = dataStoreFactory;
-    TokenRegistry tokenRegistry = tokenProvider.getOauthTokenRegistry(userID);
+    // The registry identity comes from the factory, not from userID: userID is the
+    // Google-side data store id, which for a shared plugin is a placeholder name
+    // rather than a Drill user. Using it here would look up a per-user registry
+    // for a plugin whose tokens live in the shared one.
+    TokenRegistry tokenRegistry =
+      tokenProvider.getOauthTokenRegistry(dataStoreFactory.getRegistryIdentity());
     this.tokenTable = tokenRegistry.getTokenTable(pluginName);
     if (hasValidTokens(tokenTable)) {
       keyValueMap.put(tokenTable.ACCESS_TOKEN_KEY, tokenTable.getAccessToken().getBytes(StandardCharsets.UTF_8));

contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/DrillDataStoreFactory.java

@@ -28,10 +28,24 @@
 public class DrillDataStoreFactory extends AbstractDataStoreFactory {
   private final OAuthTokenProvider tokenProvider;
   private final String pluginName;
+  private final String registryIdentity;
 
-  public DrillDataStoreFactory(OAuthTokenProvider tokenProvider, String pluginName) {
+  /**
+   * @param registryIdentity the identity whose Drill token registry backs these
+   *     data stores: the query user for a {@code USER_TRANSLATION} plugin, or
+   *     null for the shared registry. This is deliberately separate from the
+   *     Google-side data store id passed to {@link #createDataStore(String)},
+   *     which must be a non-null string and is not a Drill identity.
+   */
+  public DrillDataStoreFactory(OAuthTokenProvider tokenProvider, String pluginName,
+                               String registryIdentity) {
     this.tokenProvider = tokenProvider;
     this.pluginName = pluginName;
+    this.registryIdentity = registryIdentity;
+  }
+
+  public String getRegistryIdentity() {
+    return registryIdentity;
   }
 
   @Override
@@ -44,6 +58,7 @@ public String toString() {
     return new PlanStringBuilder(this)
       .field("tokenProvider", tokenProvider)
       .field("pluginName", pluginName)
+      .field("registryIdentity", registryIdentity)
       .toString();
   }
 }

contrib/storage-googlesheets/src/main/java/org/apache/drill/exec/store/googlesheets/GoogleSheetsStoragePlugin.java

@@ -44,6 +44,7 @@
 import org.apache.drill.exec.store.base.filter.FilterPushDownUtils;
 import org.apache.drill.exec.store.googlesheets.schema.GoogleSheetsSchemaFactory;
 import org.apache.drill.exec.store.googlesheets.utils.GoogleSheetsUtils;
+import org.apache.commons.lang3.StringUtils;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableSet;
 import org.slf4j.Logger;
@@ -52,19 +53,36 @@
 import java.io.IOException;
 import java.security.GeneralSecurityException;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
 
 public class GoogleSheetsStoragePlugin extends AbstractStoragePlugin {
   private final static Logger logger = LoggerFactory.getLogger(GoogleSheetsStoragePlugin.class);
   private final static String SHARED_USERNAME = "anonymous";
   private final GoogleSheetsStoragePluginConfig config;
   private final GoogleSheetsSchemaFactory schemaFactory;
   private final OAuthTokenProvider tokenProvider;
-  private DataStore<StoredCredential> dataStore;
-  private Sheets service;
-  private Drive driveService;
-  private TokenRegistry tokenRegistry;
-  private String username;
+  /**
+   * Google credential stores and services, cached per identity.
+   *
+   * <p>These were previously single fields on the plugin. A plugin instance is
+   * shared by every query, so the first caller's authenticated service was handed
+   * to every later caller regardless of which user asked for it, which under
+   * {@code USER_TRANSLATION} meant one user's queries reaching Google as another
+   * user. They are keyed by {@link #cacheKey(String)} so each identity gets, and
+   * keeps, only its own.
+   */
+  private final Map<String, DataStore<StoredCredential>> dataStores = new ConcurrentHashMap<>();
+  private final Map<String, Sheets> sheetsServices = new ConcurrentHashMap<>();
+  private final Map<String, Drive> driveServices = new ConcurrentHashMap<>();
+
+  /**
+   * The access token each identity's cached entries were built from; see
+   * {@link #cached(Map, String, Supplier)}.
+   */
+  private final Map<String, String> cachedAccessTokens = new ConcurrentHashMap<>();
 
 
   public GoogleSheetsStoragePlugin(GoogleSheetsStoragePluginConfig configuration, DrillbitContext context, String name) {
@@ -95,22 +113,109 @@ public void initializeOauthTokenTable(SchemaConfig schemaConfig) {
     // the plugin is loaded, we read the tokens from the persistent store into a GS dataStore.
     // This happens when the plugin is registered.
 
-    if (config.getAuthMode() == AuthMode.USER_TRANSLATION) {
-      this.username = schemaConfig.getUserName();
-      tokenRegistry = tokenProvider.getOauthTokenRegistry(this.username);
-    } else {
-      this.username = SHARED_USERNAME;
-      tokenRegistry = tokenProvider.getOauthTokenRegistry(null);
-    }
+    getTokenRegistry(schemaConfig == null ? null : schemaConfig.getUserName());
+  }
+
+  /**
+   * The token registry holding {@code queryUser}'s tokens, with this plugin's
+   * token table created in it if absent. Resolved per call; see
+   * {@link #dataStores} for why nothing here is cached on the plugin.
+   */
+  public TokenRegistry getTokenRegistry(String queryUser) {
+    TokenRegistry tokenRegistry = tokenProvider.getOauthTokenRegistry(registryIdentity(queryUser));
     tokenRegistry.createTokenTable(getName());
-    this.dataStore = new DrillDataStoreFactory(tokenProvider, getName()).createDataStore(this.username);
+    return tokenRegistry;
+  }
+
+  /**
+   * The identity whose Drill token registry holds this plugin's tokens: the query
+   * user under user translation, or null for the single shared registry. Note that
+   * {@link #SHARED_USERNAME} is only a Google-side placeholder and must not be used
+   * as a registry identity, or a shared plugin's tokens would be looked up in a
+   * per-user registry that the REST credential endpoints never write to.
+   */
+  private String registryIdentity(String queryUser) {
+    return config.getAuthMode() == AuthMode.USER_TRANSLATION ? queryUser : null;
+  }
+
+  /**
+   * The Google-side data store id, which must be a non-null string and is not a
+   * Drill identity. Under user translation it is the query user; otherwise a
+   * fixed placeholder.
+   */
+  private String googleStoreId(String queryUser) {
+    return config.getAuthMode() == AuthMode.USER_TRANSLATION ? queryUser : SHARED_USERNAME;
+  }
+
+  /**
+   * Cache key for the per-identity maps. Prefixed rather than using the user name
+   * directly, so that a user literally named {@link #SHARED_USERNAME} — which is
+   * what every web user is called when authentication is disabled — cannot share
+   * a cache entry with a shared-credential plugin.
+   */
+  private String cacheKey(String queryUser) {
+    return config.getAuthMode() == AuthMode.USER_TRANSLATION ? "user:" + queryUser : "shared";
   }
 
-  public DataStore<StoredCredential> getDataStore(String username) {
-    if (this.dataStore == null) {
-      this.dataStore = new DrillDataStoreFactory(tokenProvider, getName()).createDataStore(username);
+  private DataStore<StoredCredential> buildDataStore(String queryUser) {
+    return new DrillDataStoreFactory(tokenProvider, getName(), registryIdentity(queryUser))
+      .createDataStore(googleStoreId(queryUser));
+  }
+
+  public DataStore<StoredCredential> getDataStore(String queryUser) {
+    return cached(dataStores, queryUser, () -> buildDataStore(queryUser));
+  }
+
+  /**
+   * This identity's cached value, built on first use, but only kept for as long
+   * as the tokens it was built from are the tokens in the token table.
+   *
+   * <p>Google's {@code Credential} copies the access and refresh tokens once, and
+   * {@link DrillDataStore} reads the token table only in its constructor. So an
+   * entry built before the user authorized would serve an empty credential until
+   * the drillbit restarted, and under user translation that is the normal flow
+   * for a new user: query, get an auth error, authorize, query again. Nothing is
+   * cached while the identity has no access token, and an entry is dropped once
+   * the stored token differs from the one it was built from, so authorizing or
+   * re-authorizing takes effect on the next query.
+   *
+   * <p>A token refresh rewrites the stored token as well, which costs one extra
+   * rebuild; racing callers may likewise build one entry more than they use.
+   * Both are cheap next to serving the wrong identity's credential.
+   */
+  private <T> T cached(Map<String, T> cache, String queryUser, Supplier<T> builder) {
+    String key = cacheKey(queryUser);
+    String accessToken = getTokenTable(queryUser).getAccessToken();
+
+    if (StringUtils.isEmpty(accessToken)) {
+      invalidate(key);
+      return builder.get();
     }
-    return dataStore;
+
+    if (!accessToken.equals(cachedAccessTokens.get(key))) {
+      invalidate(key);
+      cachedAccessTokens.put(key, accessToken);
+    }
+
+    T cachedValue = cache.get(key);
+    if (cachedValue != null) {
+      return cachedValue;
+    }
+
+    // Built outside computeIfAbsent: building a service asks for this identity's
+    // data store, which comes back through here, and a mapping function must not
+    // touch the map it is updating.
+    T value = builder.get();
+    cache.put(key, value);
+    return value;
+  }
+
+  /** Forgets every service cached for one identity. */
+  private void invalidate(String key) {
+    dataStores.remove(key);
+    sheetsServices.remove(key);
+    driveServices.remove(key);
+    cachedAccessTokens.remove(key);
   }
 
 
@@ -120,8 +225,17 @@ public void registerSchemas(SchemaConfig schemaConfig, SchemaPlus parent) {
     schemaFactory.registerSchemas(schemaConfig, parent);
   }
 
+  /**
+   * This plugin's token table for the shared identity. Prefer
+   * {@link #getTokenTable(String)} wherever the query user is known.
+   */
   public PersistentTokenTable getTokenTable() {
-    return tokenRegistry.getTokenTable(getName());
+    return getTokenTable(null);
+  }
+
+  /** This plugin's token table holding {@code queryUser}'s tokens. */
+  public PersistentTokenTable getTokenTable(String queryUser) {
+    return getTokenRegistry(queryUser).getTokenTable(getName());
   }
 
   @Override
@@ -153,7 +267,7 @@ public AbstractGroupScan getPhysicalScan(String userName, JSONOptions selection)
   public AbstractGroupScan getPhysicalScan(String userName, JSONOptions selection, List<SchemaPath> columns, SessionOptionManager options,
                                            MetadataProviderManager metadataProviderManager) throws IOException {
     GoogleSheetsScanSpec scanSpec = selection.getListWith(context.getLpPersistence().getMapper(), new TypeReference<GoogleSheetsScanSpec>() {});
-    return new GoogleSheetsGroupScan(this.username, scanSpec, this, metadataProviderManager);
+    return new GoogleSheetsGroupScan(userName, scanSpec, this, metadataProviderManager);
   }
 
   @Override
@@ -182,8 +296,7 @@ public StoragePluginConfig getConfig() {
    */
   @VisibleForTesting
   public void initializeTokenTableForTesting() {
-    OAuthTokenProvider tokenProvider = context.getOauthTokenProvider();
-    tokenRegistry = tokenProvider.getOauthTokenRegistry(null);
+    getTokenRegistry(null);
   }
 
   @Override
@@ -201,26 +314,17 @@ public boolean supportsInsert() {
    * @return An authenticated {@link Sheets} Google Sheets service.
    */
   public Sheets getSheetsService(String queryUser) {
-    if (service != null && dataStore != null) {
-      return service;
-    } else {
-      // Check if datastore is null and initialize if so.
-      if (dataStore == null) {
-        this.dataStore = getDataStore(queryUser);
-      }
-
-      try {
-        if (config.getAuthMode() == AuthMode.USER_TRANSLATION) {
-          service = GoogleSheetsUtils.getSheetsService(config, dataStore, queryUser);
-        } else {
-          service = GoogleSheetsUtils.getSheetsService(config, dataStore, SHARED_USERNAME);
-        }
-        return service;
-      } catch (IOException | GeneralSecurityException e) {
-        throw UserException.connectionError(e)
-          .message("Error connecting to Googlesheets Service: " + e.getMessage())
-          .build(logger);
-      }
+    return cached(sheetsServices, queryUser, () -> buildSheetsService(queryUser));
+  }
+
+  private Sheets buildSheetsService(String queryUser) {
+    try {
+      return GoogleSheetsUtils.getSheetsService(
+        config, getDataStore(queryUser), googleStoreId(queryUser));
+    } catch (IOException | GeneralSecurityException e) {
+      throw UserException.connectionError(e)
+        .message("Error connecting to Googlesheets Service: " + e.getMessage())
+        .build(logger);
     }
   }
 
@@ -231,26 +335,17 @@ public Sheets getSheetsService(String queryUser) {
    * @return A validated and authenticated {@link Drive} instance.
    */
   public Drive getDriveService(String queryUser) {
-    if (driveService != null && dataStore != null) {
-      return driveService;
-    } else {
-      // Check if datastore is null and initialize if so.
-      if (dataStore == null) {
-        this.dataStore = getDataStore(queryUser);
-      }
-
-      try {
-        if (config.getAuthMode() == AuthMode.USER_TRANSLATION) {
-          driveService = GoogleSheetsUtils.getDriveService(config, dataStore, queryUser);
-        } else {
-          driveService = GoogleSheetsUtils.getDriveService(config, dataStore, SHARED_USERNAME);
-        }
-        return driveService;
-      } catch (IOException | GeneralSecurityException e) {
-        throw UserException.connectionError(e)
-          .message("Error connecting to Google Drive Service: " + e.getMessage())
-          .build(logger);
-      }
+    return cached(driveServices, queryUser, () -> buildDriveService(queryUser));
+  }
+
+  private Drive buildDriveService(String queryUser) {
+    try {
+      return GoogleSheetsUtils.getDriveService(
+        config, getDataStore(queryUser), googleStoreId(queryUser));
+    } catch (IOException | GeneralSecurityException e) {
+      throw UserException.connectionError(e)
+        .message("Error connecting to Google Drive Service: " + e.getMessage())
+        .build(logger);
     }
   }
 }

contrib/storage-http/src/main/java/org/apache/drill/exec/store/http/HttpAPIConnectionSchema.java

@@ -68,7 +68,7 @@ public Table getTable(String tableName) {
       // Register a new table
       return registerTable(tableName, new DynamicDrillTable(plugin, plugin.getName(),
         new HttpScanSpec(plugin.getName(), name, tableName,
-              plugin.getConfig().copyForPlan(name), plugin.getTokenTable(), queryUserName, plugin.getRegistry())));
+              plugin.getConfig().copyForPlan(name), plugin.getTokenTable(queryUserName), queryUserName, plugin.getRegistry())));
     }
   }
 

contrib/storage-http/src/main/java/org/apache/drill/exec/store/http/HttpSchemaFactory.java

@@ -106,7 +106,7 @@ public Table getTable(String name) {
         // Register a new table
         return registerTable(name, new DynamicDrillTable(plugin, plugin.getName(),
             new HttpScanSpec(plugin.getName(), name, null,
-                plugin.getConfig().copyForPlan(name), plugin.getTokenTable(), queryUserName, plugin.getRegistry())));
+                plugin.getConfig().copyForPlan(name), plugin.getTokenTable(queryUserName), queryUserName, plugin.getRegistry())));
       } else {
         return null; // Unknown table
       }