// Copyright 2020 Google LLC
// 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
// 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.
// See the License for the specific language governing permissions and
// limitations under the License.

package com.google.cloud.broker.hadoop.fs;

import java.io.IOException;
import java.security.PrivilegedAction;
import java.util.Collections;
import java.util.Set;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.security.UserGroupInformation;

import com.google.cloud.hadoop.util.AccessTokenProvider;

// Classes dynamically generated by protobuf-maven-plugin:
import com.google.cloud.broker.apps.brokerserver.protobuf.GetAccessTokenRequest;
import com.google.cloud.broker.apps.brokerserver.protobuf.GetAccessTokenResponse;


public final class BrokerAccessTokenProvider implements AccessTokenProvider {

    private Configuration config;
    private AccessToken accessToken;
    private BrokerTokenIdentifier tokenIdentifier;
    private Text service;

    private final static AccessToken EXPIRED_TOKEN = new AccessToken("", -1L);

    public BrokerAccessTokenProvider(Text service) {
        this.service = service;
        this.accessToken = EXPIRED_TOKEN;
    }

    public BrokerAccessTokenProvider(Text service, BrokerTokenIdentifier bti) {
        this.service = service;
        this.tokenIdentifier = bti;
        this.accessToken = EXPIRED_TOKEN;
    }

    @Override
    public AccessToken getAccessToken() {
        return this.accessToken;
    }

    @Override
    public void refresh() {
        UserGroupInformation currentUser;
        UserGroupInformation loginUser;
        try {
            currentUser = UserGroupInformation.getCurrentUser();
            loginUser = UserGroupInformation.getLoginUser();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        String sessionToken;
        String target;
        String owner;
        Set<String> scopes;

        if (tokenIdentifier != null) {
            sessionToken = tokenIdentifier.getSessionToken();
            owner = null;
            target = null;
            scopes = null;
        }
        else {
            sessionToken = null;
            owner = currentUser.getUserName();
            target = service.toString();
            scopes = Collections.singleton(BrokerTokenIdentifier.GCS_SCOPE);
        }

        GetAccessTokenResponse response = loginUser.doAs((PrivilegedAction<GetAccessTokenResponse>) () -> {
            BrokerGateway gateway = new BrokerGateway(config, sessionToken);
            GetAccessTokenRequest request = GetAccessTokenRequest.newBuilder()
                .addAllScopes(scopes)
                .setOwner(owner)
                .setTarget(target)
                .build();
            GetAccessTokenResponse r = gateway.getStub().getAccessToken(request);
            gateway.getManagedChannel().shutdown();
            return r;
        });

        String tokenString = response.getAccessToken();
        long expiresAt = response.getExpiresAt();
        accessToken = new AccessToken(tokenString, expiresAt);
    }

    @Override
    public void setConf(Configuration config) {
        this.config = config;
    }

    @Override
    public Configuration getConf() {
        return this.config;
    }

}
@@ -2249,82 +2249,155 @@
     
-String sessionToken;%0A        String target;%0A        String owner;%0A
+GetAccessTokenResponse response = loginUser.doAs((PrivilegedAction%3CGetAccessTokenResponse%3E) () -%3E %7B%0A            BrokerGateway gateway;%0A    
         
 Set%3C
@@ -2396,29 +2396,43 @@
     
-Set%3CString%3E scop
+GetAccessTokenRequest requ
 es
+t
 ;%0A
-%0A
+    
     
@@ -2459,9 +2459,9 @@
 ier 
-!
+=
 = nu
@@ -2469,122 +2469,92 @@
 l) %7B
-%0A
   
-          sessionToken = tokenIdentifier.getSessionToken();%0A            owner = null;%0A            target =
+// Direct authentication%0A                gateway = new BrokerGateway(config,
  null
+)
 ;%0A  
@@ -2567,22 +2567,56 @@
     
-scopes = null;
+    request = GetAccessTokenRequest.newBuilder()
 %0A   
@@ -2624,18 +2624,16 @@
     
-%7D%0A
         
 else
@@ -2632,68 +2632,112 @@
     
-else %7B%0A            sessionToken = null;%0A            o
+    .addAllScopes(Collections.singleton(BrokerTokenIdentifier.GCS_SCOPE))%0A                    .setO
 wner
- = 
+(
 curr
@@ -2757,18 +2757,26 @@
 erName()
-;%0A
+)%0A        
         
@@ -2783,17 +2783,19 @@
     
-t
+.setT
 arget
- = 
+(
 serv
@@ -2808,18 +2808,20 @@
 String()
-;%0A
+)%0A  
         
@@ -2828,73 +2828,29 @@
     
-scopes = Collections.singleton(BrokerTokenIdentifier.GCS_SCOPE);%0A
+      .build();%0A%0A    
     
@@ -2851,33 +2851,32 @@
           %7D%0A
-%0A
         
 GetAccessTok
@@ -2867,133 +2867,63 @@
     
-GetAccessTokenResponse response = loginUser.doAs((PrivilegedAction%3CGetAccessTokenResponse%3E) () -%3E %7B%0A            BrokerGateway
+    else %7B  // Delegated authentication%0A               
  gat
@@ -2955,17 +2955,36 @@
 config, 
-s
+tokenIdentifier.getS
 essionTo
@@ -2986,16 +2986,18 @@
 ionToken
+()
 );%0A     
@@ -3003,37 +3003,19 @@
         
-GetAccessTokenRequest
+   
  request
@@ -3055,140 +3055,31 @@
 er()
-%0A                .addAllScopes(scopes)%0A                .setOwner(owner)%0A                .setTarget(target)%0A                .build();
+.build();%0A            %7D
 %0A   

Fix connector