feat: allow set lifetime for service account creds (#516)

* feat: allow set lifetime for service account creds

* update

* update name

* update

* update

* change lifetime 0 to default

* update

* update

Co-authored-by: Jeff Ching <chingor@google.com>

oauth2_http/java/com/google/auth/oauth2/ServiceAccountCredentials.java

@@ -93,6 +93,8 @@ public class ServiceAccountCredentials extends GoogleCredentials
   private static final long serialVersionUID = 7807543542681217978L;
   private static final String GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer";
   private static final String PARSE_ERROR_PREFIX = "Error parsing token refresh response. ";
+  private static final int TWELVE_HOURS_IN_SECONDS = 43200;
+  private static final int DEFAULT_LIFETIME_IN_SECONDS = 3600;
 
   private final String clientId;
   private final String clientEmail;
@@ -104,6 +106,7 @@ public class ServiceAccountCredentials extends GoogleCredentials
   private final URI tokenServerUri;
   private final Collection<String> scopes;
   private final String quotaProjectId;
+  private final int lifetime;
 
   private transient HttpTransportFactory transportFactory;
 
@@ -123,6 +126,10 @@ public class ServiceAccountCredentials extends GoogleCredentials
    *     authority to the service account.
    * @param projectId the project used for billing
    * @param quotaProjectId The project used for quota and billing purposes. May be null.
+   * @param lifetime number of seconds the access token should be valid for. The value should be at
+   *     most 43200 (12 hours). If the token is used for calling a Google API, then the value should
+   *     be at most 3600 (1 hour). If the given value is 0, then the default value 3600 will be used
+   *     when creating the credentials.
    */
   ServiceAccountCredentials(
       String clientId,
@@ -134,7 +141,8 @@ public class ServiceAccountCredentials extends GoogleCredentials
       URI tokenServerUri,
       String serviceAccountUser,
       String projectId,
-      String quotaProjectId) {
+      String quotaProjectId,
+      int lifetime) {
     this.clientId = clientId;
     this.clientEmail = Preconditions.checkNotNull(clientEmail);
     this.privateKey = Preconditions.checkNotNull(privateKey);
@@ -149,6 +157,10 @@ public class ServiceAccountCredentials extends GoogleCredentials
     this.serviceAccountUser = serviceAccountUser;
     this.projectId = projectId;
     this.quotaProjectId = quotaProjectId;
+    if (lifetime > TWELVE_HOURS_IN_SECONDS) {
+      throw new IllegalStateException("lifetime must be less than or equal to 43200");
+    }
+    this.lifetime = lifetime;
   }
 
   /**
@@ -325,7 +337,8 @@ static ServiceAccountCredentials fromPkcs8(
         tokenServerUri,
         serviceAccountUser,
         projectId,
-        quotaProject);
+        quotaProject,
+        DEFAULT_LIFETIME_IN_SECONDS);
   }
 
   /** Helper to convert from a PKCS#8 String to an RSA private key */
@@ -513,7 +526,21 @@ public GoogleCredentials createScoped(Collection<String> newScopes) {
         tokenServerUri,
         serviceAccountUser,
         projectId,
-        quotaProjectId);
+        quotaProjectId,
+        lifetime);
+  }
+
+  /**
+   * Clones the service account with a new lifetime value.
+   *
+   * @param lifetime life time value in seconds. The value should be at most 43200 (12 hours). If
+   *     the token is used for calling a Google API, then the value should be at most 3600 (1 hour).
+   *     If the given value is 0, then the default value 3600 will be used when creating the
+   *     credentials.
+   * @return the cloned service account credentials with the given custom life time
+   */
+  public ServiceAccountCredentials createWithCustomLifetime(int lifetime) {
+    return this.toBuilder().setLifetime(lifetime).build();
   }
 
   @Override
@@ -528,7 +555,8 @@ public GoogleCredentials createDelegated(String user) {
         tokenServerUri,
         user,
         projectId,
-        quotaProjectId);
+        quotaProjectId,
+        lifetime);
   }
 
   public final String getClientId() {
@@ -563,6 +591,11 @@ public final URI getTokenServerUri() {
     return tokenServerUri;
   }
 
+  @VisibleForTesting
+  int getLifetime() {
+    return lifetime;
+  }
+
   @Override
   public String getAccount() {
     return getClientEmail();
@@ -618,7 +651,8 @@ public int hashCode() {
         transportFactoryClassName,
         tokenServerUri,
         scopes,
-        quotaProjectId);
+        quotaProjectId,
+        lifetime);
   }
 
   @Override
@@ -632,6 +666,7 @@ public String toString() {
         .add("scopes", scopes)
         .add("serviceAccountUser", serviceAccountUser)
         .add("quotaProjectId", quotaProjectId)
+        .add("lifetime", lifetime)
         .toString();
   }
 
@@ -648,7 +683,8 @@ public boolean equals(Object obj) {
         && Objects.equals(this.transportFactoryClassName, other.transportFactoryClassName)
         && Objects.equals(this.tokenServerUri, other.tokenServerUri)
         && Objects.equals(this.scopes, other.scopes)
-        && Objects.equals(this.quotaProjectId, other.quotaProjectId);
+        && Objects.equals(this.quotaProjectId, other.quotaProjectId)
+        && Objects.equals(this.lifetime, other.lifetime);
   }
 
   String createAssertion(JsonFactory jsonFactory, long currentTime, String audience)
@@ -661,7 +697,7 @@ String createAssertion(JsonFactory jsonFactory, long currentTime, String audienc
     JsonWebToken.Payload payload = new JsonWebToken.Payload();
     payload.setIssuer(clientEmail);
     payload.setIssuedAtTimeSeconds(currentTime / 1000);
-    payload.setExpirationTimeSeconds(currentTime / 1000 + 3600);
+    payload.setExpirationTimeSeconds(currentTime / 1000 + this.lifetime);
     payload.setSubject(serviceAccountUser);
     payload.put("scope", Joiner.on(' ').join(scopes));
 
@@ -693,7 +729,7 @@ String createAssertionForIdToken(
     JsonWebToken.Payload payload = new JsonWebToken.Payload();
     payload.setIssuer(clientEmail);
     payload.setIssuedAtTimeSeconds(currentTime / 1000);
-    payload.setExpirationTimeSeconds(currentTime / 1000 + 3600);
+    payload.setExpirationTimeSeconds(currentTime / 1000 + this.lifetime);
     payload.setSubject(serviceAccountUser);
 
     if (audience == null) {
@@ -746,6 +782,7 @@ public static class Builder extends GoogleCredentials.Builder {
     private Collection<String> scopes;
     private HttpTransportFactory transportFactory;
     private String quotaProjectId;
+    private int lifetime = DEFAULT_LIFETIME_IN_SECONDS;
 
     protected Builder() {}
 
@@ -760,6 +797,7 @@ protected Builder(ServiceAccountCredentials credentials) {
       this.serviceAccountUser = credentials.serviceAccountUser;
       this.projectId = credentials.projectId;
       this.quotaProjectId = credentials.quotaProjectId;
+      this.lifetime = credentials.lifetime;
     }
 
     public Builder setClientId(String clientId) {
@@ -812,6 +850,11 @@ public Builder setQuotaProjectId(String quotaProjectId) {
       return this;
     }
 
+    public Builder setLifetime(int lifetime) {
+      this.lifetime = lifetime == 0 ? DEFAULT_LIFETIME_IN_SECONDS : lifetime;
+      return this;
+    }
+
     public String getClientId() {
       return clientId;
     }
@@ -852,6 +895,10 @@ public String getQuotaProjectId() {
       return quotaProjectId;
     }
 
+    public int getLifetime() {
+      return lifetime;
+    }
+
     public ServiceAccountCredentials build() {
       return new ServiceAccountCredentials(
           clientId,
@@ -863,7 +910,8 @@ public ServiceAccountCredentials build() {
           tokenServerUri,
           serviceAccountUser,
           projectId,
-          quotaProjectId);
+          quotaProjectId,
+          lifetime);
     }
   }
 }

oauth2_http/javatests/com/google/auth/oauth2/ServiceAccountCredentialsTest.java

@@ -111,6 +111,54 @@ public class ServiceAccountCredentialsTest extends BaseSerializationTest {
           + "aXNzIjoiaHR0cHM6Ly9hY2NvdW50cy5nb29nbGUuY29tIiwic3ViIjoiMTAyMTAxNTUwODM0MjAwNzA4NTY4In0"
           + ".redacted";
   private static final String QUOTA_PROJECT = "sample-quota-project-id";
+  private static final int DEFAULT_LIFETIME_IN_SECONDS = 3600;
+  private static final int INVALID_LIFETIME = 43210;
+
+  private ServiceAccountCredentials.Builder createDefaultBuilder() throws IOException {
+    PrivateKey privateKey = ServiceAccountCredentials.privateKeyFromPkcs8(PRIVATE_KEY_PKCS8);
+    return ServiceAccountCredentials.newBuilder()
+        .setClientId(CLIENT_ID)
+        .setClientEmail(CLIENT_EMAIL)
+        .setPrivateKey(privateKey)
+        .setPrivateKeyId(PRIVATE_KEY_ID)
+        .setScopes(SCOPES)
+        .setServiceAccountUser(USER)
+        .setProjectId(PROJECT_ID);
+  }
+
+  @Test
+  public void setLifetime() throws IOException {
+    ServiceAccountCredentials.Builder builder = createDefaultBuilder();
+    assertEquals(DEFAULT_LIFETIME_IN_SECONDS, builder.getLifetime());
+    assertEquals(DEFAULT_LIFETIME_IN_SECONDS, builder.build().getLifetime());
+
+    builder.setLifetime(4000);
+    assertEquals(4000, builder.getLifetime());
+    assertEquals(4000, builder.build().getLifetime());
+
+    builder.setLifetime(0);
+    assertEquals(DEFAULT_LIFETIME_IN_SECONDS, builder.build().getLifetime());
+  }
+
+  @Test
+  public void setLifetime_invalid_lifetime() throws IOException, IllegalStateException {
+    try {
+      createDefaultBuilder().setLifetime(INVALID_LIFETIME).build();
+      fail(
+          String.format(
+              "Should throw exception with message containing '%s'",
+              "lifetime must be less than or equal to 43200"));
+    } catch (IllegalStateException expected) {
+      assertTrue(expected.getMessage().contains("lifetime must be less than or equal to 43200"));
+    }
+  }
+
+  @Test
+  public void createWithCustomLifetime() throws IOException {
+    ServiceAccountCredentials credentials = createDefaultBuilder().build();
+    credentials = credentials.createWithCustomLifetime(4000);
+    assertEquals(4000, credentials.getLifetime());
+  }
 
   @Test
   public void createdScoped_clones() throws IOException {
@@ -202,6 +250,19 @@ public void createAssertion_correct() throws IOException {
     assertEquals(Joiner.on(' ').join(scopes), payload.get("scope"));
   }
 
+  @Test
+  public void createAssertion_custom_lifetime() throws IOException {
+    ServiceAccountCredentials credentials = createDefaultBuilder().setLifetime(4000).build();
+
+    JsonFactory jsonFactory = OAuth2Utils.JSON_FACTORY;
+    long currentTimeMillis = Clock.SYSTEM.currentTimeMillis();
+    String assertion = credentials.createAssertion(jsonFactory, currentTimeMillis, null);
+
+    JsonWebSignature signature = JsonWebSignature.parse(jsonFactory, assertion);
+    JsonWebToken.Payload payload = signature.getPayload();
+    assertEquals(currentTimeMillis / 1000 + 4000, (long) payload.getExpirationTimeSeconds());
+  }
+
   @Test
   public void createAssertionForIdToken_correct() throws IOException {
 
@@ -231,6 +292,22 @@ public void createAssertionForIdToken_correct() throws IOException {
     assertEquals(USER, payload.getSubject());
   }
 
+  @Test
+  public void createAssertionForIdToken_custom_lifetime() throws IOException {
+
+    ServiceAccountCredentials credentials = createDefaultBuilder().setLifetime(4000).build();
+
+    JsonFactory jsonFactory = OAuth2Utils.JSON_FACTORY;
+    long currentTimeMillis = Clock.SYSTEM.currentTimeMillis();
+    String assertion =
+        credentials.createAssertionForIdToken(
+            jsonFactory, currentTimeMillis, null, "https://foo.com/bar");
+
+    JsonWebSignature signature = JsonWebSignature.parse(jsonFactory, assertion);
+    JsonWebToken.Payload payload = signature.getPayload();
+    assertEquals(currentTimeMillis / 1000 + 4000, (long) payload.getExpirationTimeSeconds());
+  }
+
   @Test
   public void createAssertionForIdToken_incorrect() throws IOException {
 
@@ -904,7 +981,7 @@ public void toString_containsFields() throws IOException {
         String.format(
             "ServiceAccountCredentials{clientId=%s, clientEmail=%s, privateKeyId=%s, "
                 + "transportFactoryClassName=%s, tokenServerUri=%s, scopes=%s, serviceAccountUser=%s, "
-                + "quotaProjectId=%s}",
+                + "quotaProjectId=%s, lifetime=3600}",
             CLIENT_ID,
             CLIENT_EMAIL,
             PRIVATE_KEY_ID,