From 55a2fe890e3ef2090179812f015f83d380ffd968 Mon Sep 17 00:00:00 2001
From: vzikratyi <vzikratyi@thingsboard.io>
Date: Wed, 24 Jun 2020 17:36:26 +0300
Subject: [PATCH] Added method to load oauth2 config with tenantId

---
 .../Oauth2AuthenticationSuccessHandler.java   |  2 +-
 .../server/dao/oauth2/OAuth2Service.java      |  3 ++
 .../server/dao/oauth2/OAuth2ServiceImpl.java  | 35 ++++++++++++++-----
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/application/src/main/java/org/thingsboard/server/service/security/auth/oauth2/Oauth2AuthenticationSuccessHandler.java b/application/src/main/java/org/thingsboard/server/service/security/auth/oauth2/Oauth2AuthenticationSuccessHandler.java
index 4dda36a0fb..03c9db757f 100644
--- a/application/src/main/java/org/thingsboard/server/service/security/auth/oauth2/Oauth2AuthenticationSuccessHandler.java
+++ b/application/src/main/java/org/thingsboard/server/service/security/auth/oauth2/Oauth2AuthenticationSuccessHandler.java
@@ -64,7 +64,7 @@ public class Oauth2AuthenticationSuccessHandler extends SimpleUrlAuthenticationS
         try {
             OAuth2AuthenticationToken token = (OAuth2AuthenticationToken) authentication;
 
-            OAuth2ClientRegistration clientRegistration = oAuth2Service.getClientRegistration(token.getAuthorizedClientRegistrationId());
+            OAuth2ClientRegistration clientRegistration = oAuth2Service.getClientRegistrationWithTenant(token.getAuthorizedClientRegistrationId()).getRight();
             OAuth2ClientMapper mapper = oauth2ClientMapperProvider.getOAuth2ClientMapperByType(clientRegistration.getMapperConfig().getType());
             SecurityUser securityUser = mapper.getOrCreateUserByClientPrincipal(token, clientRegistration.getMapperConfig());
 
diff --git a/common/dao-api/src/main/java/org/thingsboard/server/dao/oauth2/OAuth2Service.java b/common/dao-api/src/main/java/org/thingsboard/server/dao/oauth2/OAuth2Service.java
index b02a48650f..17dc6d5f61 100644
--- a/common/dao-api/src/main/java/org/thingsboard/server/dao/oauth2/OAuth2Service.java
+++ b/common/dao-api/src/main/java/org/thingsboard/server/dao/oauth2/OAuth2Service.java
@@ -15,6 +15,7 @@
  */
 package org.thingsboard.server.dao.oauth2;
 
+import org.apache.commons.lang3.tuple.Pair;
 import org.thingsboard.server.common.data.id.CustomerId;
 import org.thingsboard.server.common.data.id.EntityId;
 import org.thingsboard.server.common.data.id.TenantId;
@@ -25,6 +26,8 @@ import org.thingsboard.server.common.data.oauth2.OAuth2ClientsParams;
 import java.util.List;
 
 public interface OAuth2Service {
+    Pair<TenantId, OAuth2ClientRegistration> getClientRegistrationWithTenant(String registrationId);
+
     OAuth2ClientRegistration getClientRegistration(String registrationId);
 
     List<OAuth2ClientInfo> getOAuth2Clients(String domainName);
diff --git a/dao/src/main/java/org/thingsboard/server/dao/oauth2/OAuth2ServiceImpl.java b/dao/src/main/java/org/thingsboard/server/dao/oauth2/OAuth2ServiceImpl.java
index b2824fdf4e..e23dbf4a20 100644
--- a/dao/src/main/java/org/thingsboard/server/dao/oauth2/OAuth2ServiceImpl.java
+++ b/dao/src/main/java/org/thingsboard/server/dao/oauth2/OAuth2ServiceImpl.java
@@ -1,12 +1,12 @@
 /**
  * Copyright © 2016-2020 The Thingsboard Authors
- *
+ * <p>
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -24,6 +24,8 @@ import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.collections.ListUtils;
+import org.apache.commons.lang3.tuple.ImmutablePair;
+import org.apache.commons.lang3.tuple.Pair;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.core.env.Environment;
 import org.springframework.stereotype.Service;
@@ -89,12 +91,27 @@ public class OAuth2ServiceImpl implements OAuth2Service {
     }
 
     @Override
-    public OAuth2ClientRegistration getClientRegistration(String registrationId) {
-        return clientsParams.values().stream()
-                .flatMap(oAuth2ClientsParams -> oAuth2ClientsParams.getClientRegistrations().stream())
-                .filter(clientRegistration -> registrationId.equals(clientRegistration.getRegistrationId()))
+    public Pair<TenantId, OAuth2ClientRegistration> getClientRegistrationWithTenant(String registrationId) {
+        return clientsParams.entrySet().stream()
+                .map(entry -> {
+                    TenantId tenantId = entry.getKey();
+                    OAuth2ClientRegistration clientRegistration = entry.getValue().getClientRegistrations().stream()
+                            .filter(registration -> registrationId.equals(registration.getRegistrationId()))
+                            .findFirst()
+                            .orElse(null);
+                    return clientRegistration != null ?
+                            ImmutablePair.of(tenantId, clientRegistration) : null;
+                })
+                .filter(Objects::nonNull)
                 .findFirst()
-                .orElse(null);
+                .orElse(null)
+                ;
+    }
+
+    @Override
+    public OAuth2ClientRegistration getClientRegistration(String registrationId) {
+        Pair<TenantId, OAuth2ClientRegistration> clientRegistrationPair = getClientRegistrationWithTenant(registrationId);
+        return clientRegistrationPair != null ? clientRegistrationPair.getRight() : null;
     }
 
     @Override