github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/clients/hadoopfs/src/test/java/io/lakefs/auth/AWSLakeFSTokenProviderTest.java (about)

     1  package io.lakefs.auth;
     2  
     3  import com.google.gson.FieldNamingPolicy;
     4  import com.google.gson.Gson;
     5  import com.google.gson.GsonBuilder;
     6  import io.lakefs.Constants;
     7  import io.lakefs.FSConfiguration;
     8  import io.lakefs.clients.sdk.model.AuthenticationToken;
     9  import org.apache.commons.codec.binary.Base64;
    10  import org.apache.hadoop.conf.Configuration;
    11  import org.junit.Assert;
    12  
    13  import org.junit.Rule;
    14  import org.junit.Test;
    15  import org.mockserver.client.MockServerClient;
    16  import org.mockserver.junit.MockServerRule;
    17  import org.mockserver.matchers.Times;
    18  import org.mockserver.model.Cookie;
    19  import org.mockserver.model.HttpRequest;
    20  
    21  import static org.mockserver.model.HttpResponse.response;
    22  
    23  
    24  public class AWSLakeFSTokenProviderTest {
    25      @Rule
    26      public MockServerRule mockServerRule = new MockServerRule(this);
    27      protected MockServerClient mockServerClient;
    28      protected final Gson gson = new GsonBuilder().setFieldNamingPolicy(FieldNamingPolicy.LOWER_CASE_WITH_UNDERSCORES).create();
    29  
    30      @Test
    31      public void testProviderIdentityTokenSerde() throws Exception {
    32          Configuration conf = new Configuration(false);
    33          conf.set("fs.lakefs." + Constants.LAKEFS_AUTH_PROVIDER_KEY_SUFFIX, TemporaryAWSCredentialsLakeFSTokenProvider.NAME);
    34          conf.set("fs.lakefs." + Constants.TOKEN_AWS_CREDENTIALS_PROVIDER_ACCESS_KEY_SUFFIX, "accessKeyId");
    35          conf.set("fs.lakefs." + Constants.TOKEN_AWS_CREDENTIALS_PROVIDER_SECRET_KEY_SUFFIX, "secretAccessKey");
    36          conf.set("fs.lakefs." + Constants.TOKEN_AWS_CREDENTIALS_PROVIDER_SESSION_TOKEN_KEY_SUFFIX, "sessionToken");
    37          conf.set("fs.lakefs." + Constants.TOKEN_AWS_STS_ENDPOINT, "https://sts.amazonaws.com");
    38  
    39          AWSLakeFSTokenProvider provider = (AWSLakeFSTokenProvider) LakeFSTokenProviderFactory.newLakeFSTokenProvider(Constants.DEFAULT_SCHEME, conf);
    40          String identityToken = provider.newPresignedGetCallerIdentityToken();
    41          String decodedToken = new String(Base64.decodeBase64(identityToken.getBytes()));
    42          LakeFSExternalPrincipalIdentityRequest request = LakeFSExternalPrincipalIdentityRequest.fromJSON(decodedToken);
    43          Assert.assertEquals("POST", request.getMethod());
    44          Assert.assertEquals("sts.amazonaws.com", request.getHost());
    45          Assert.assertEquals("us-east-1", request.getRegion());
    46          Assert.assertEquals("GetCallerIdentity", request.getAction());
    47          Assert.assertTrue(request.getDate().matches("\\d{8}T\\d{6}Z"));
    48          Assert.assertEquals("60", request.getExpirationDuration());
    49          Assert.assertEquals(FSConfiguration.get(conf, "lakefs", Constants.TOKEN_AWS_CREDENTIALS_PROVIDER_ACCESS_KEY_SUFFIX), request.getAccessKeyId());
    50          Assert.assertTrue(request.getSignature().matches("[0-9a-fA-F]{64}"));
    51          Assert.assertEquals("host", request.getSignedHeaders().get(0));
    52          Assert.assertEquals("x-lakefs-server-id", request.getSignedHeaders().get(1));
    53          Assert.assertEquals("2011-06-15", request.getVersion());
    54          Assert.assertEquals("AWS4-HMAC-SHA256", request.getAlgorithm());
    55          Assert.assertEquals(FSConfiguration.get(conf, "lakefs", Constants.TOKEN_AWS_CREDENTIALS_PROVIDER_SESSION_TOKEN_KEY_SUFFIX), request.getSecurityToken());
    56      }
    57  
    58      protected void mockExternalPrincipalLogin(Long tokenExpiration, String token, String sessionID) {
    59          // lakeFSFS initialization requires a blockstore.
    60          HttpRequest request = HttpRequest.request().withCookie(new Cookie("sessionId", sessionID));
    61  
    62          mockServerClient
    63                  .when(
    64                          request.withMethod("POST").withPath("/auth/external/principal/login"),
    65                          Times.once())
    66                  .respond(
    67                          response().withStatusCode(200).withBody(new AuthenticationToken().token(token).tokenExpiration(tokenExpiration).toJson())
    68                  );
    69      }
    70  
    71      @Test
    72      public void testProviderToken() throws Exception {
    73          String sessionID = "testProviderToken";
    74          String expectedToken = "lakefs-jwt-token";
    75          Configuration conf = new Configuration(false);
    76          conf.set("fs.lakefs." + Constants.LAKEFS_AUTH_PROVIDER_KEY_SUFFIX, TemporaryAWSCredentialsLakeFSTokenProvider.NAME);
    77          conf.set("fs.lakefs." + Constants.TOKEN_AWS_CREDENTIALS_PROVIDER_ACCESS_KEY_SUFFIX, "accessKeyId");
    78          conf.set("fs.lakefs." + Constants.TOKEN_AWS_CREDENTIALS_PROVIDER_SECRET_KEY_SUFFIX, "secretAccessKey");
    79          conf.set("fs.lakefs." + Constants.TOKEN_AWS_CREDENTIALS_PROVIDER_SESSION_TOKEN_KEY_SUFFIX, "sessionToken");
    80          conf.setInt("fs.lakefs." + Constants.LAKEFS_AUTH_TOKEN_TTL_KEY_SUFFIX, 120);
    81          conf.set("fs.lakefs." + Constants.TOKEN_AWS_STS_ENDPOINT, "https://sts.amazonaws.com");
    82          conf.set("fs.lakefs.endpoint", String.format("http://localhost:%d/", mockServerClient.getPort()));
    83          conf.set("fs.lakefs.session_id", sessionID);
    84  
    85          LakeFSTokenProvider provider = LakeFSTokenProviderFactory.newLakeFSTokenProvider(Constants.DEFAULT_SCHEME, conf);
    86          mockExternalPrincipalLogin(1000L, expectedToken, sessionID);
    87  
    88          String lakeFSJWT = provider.getToken();
    89          Assert.assertEquals(expectedToken, lakeFSJWT);
    90      }
    91  }