From f6e695445043c02321386bcf094049d91449b804 Mon Sep 17 00:00:00 2001
From: enricovianello <maverick17586@gmail.com>
Date: Mon, 26 Feb 2018 10:35:43 +0100
Subject: [PATCH] Include additional claims in ID token

Read https://github.com/indigo-iam/iam/issues/202
---
 .../service/IDTokenClaimsEnhancer.java        | 12 +++++++++
 .../impl/DefaultIdTokenClaimsEnhancer.java    | 27 +++++++++++++++++++
 .../service/impl/DefaultOIDCTokenService.java |  6 +++++
 3 files changed, 45 insertions(+)
 create mode 100644 openid-connect-common/src/main/java/org/mitre/openid/connect/service/IDTokenClaimsEnhancer.java
 create mode 100644 openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/DefaultIdTokenClaimsEnhancer.java

diff --git a/openid-connect-common/src/main/java/org/mitre/openid/connect/service/IDTokenClaimsEnhancer.java b/openid-connect-common/src/main/java/org/mitre/openid/connect/service/IDTokenClaimsEnhancer.java
new file mode 100644
index 000000000..2a628f0b9
--- /dev/null
+++ b/openid-connect-common/src/main/java/org/mitre/openid/connect/service/IDTokenClaimsEnhancer.java
@@ -0,0 +1,12 @@
+package org.mitre.openid.connect.service;
+
+import java.util.Date;
+import org.mitre.oauth2.model.OAuth2AccessTokenEntity;
+import org.springframework.security.oauth2.provider.OAuth2Request;
+import com.nimbusds.jwt.JWTClaimsSet;
+
+public interface IDTokenClaimsEnhancer {
+
+  void enhanceIdTokenClaims(JWTClaimsSet.Builder claimsBuilder, OAuth2Request request, Date issueTime,
+      String sub, OAuth2AccessTokenEntity accessToken);
+}
diff --git a/openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/DefaultIdTokenClaimsEnhancer.java b/openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/DefaultIdTokenClaimsEnhancer.java
new file mode 100644
index 000000000..2d802e2be
--- /dev/null
+++ b/openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/DefaultIdTokenClaimsEnhancer.java
@@ -0,0 +1,27 @@
+package org.mitre.openid.connect.service.impl;
+
+import java.util.Date;
+import org.mitre.oauth2.model.OAuth2AccessTokenEntity;
+import org.mitre.openid.connect.service.IDTokenClaimsEnhancer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.security.oauth2.provider.OAuth2Request;
+import org.springframework.stereotype.Service;
+import com.nimbusds.jwt.JWTClaimsSet;
+
+@Service("defaultIdTokenClaimsEnhancer")
+public class DefaultIdTokenClaimsEnhancer implements IDTokenClaimsEnhancer {
+
+  /**
+   * Logger for this class
+   */
+  private static final Logger logger = LoggerFactory.getLogger(DefaultOIDCTokenService.class);
+
+  @Override
+  public void enhanceIdTokenClaims(JWTClaimsSet.Builder claimsBuilder, OAuth2Request request, Date issueTime,
+      String sub, OAuth2AccessTokenEntity accessToken) {
+
+    logger.debug("Enhancing Id-Token claims: no claims added.");
+  }
+
+}
diff --git a/openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/DefaultOIDCTokenService.java b/openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/DefaultOIDCTokenService.java
index 49d7fab59..caf6421a0 100644
--- a/openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/DefaultOIDCTokenService.java
+++ b/openid-connect-server/src/main/java/org/mitre/openid/connect/service/impl/DefaultOIDCTokenService.java
@@ -36,6 +36,7 @@ import org.mitre.oauth2.service.AuthenticationHolderEntityService;
 import org.mitre.oauth2.service.OAuth2TokenEntityService;
 import org.mitre.oauth2.service.SystemScopeService;
 import org.mitre.openid.connect.config.ConfigurationPropertiesBean;
+import org.mitre.openid.connect.service.IDTokenClaimsEnhancer;
 import org.mitre.openid.connect.service.OIDCTokenService;
 import org.mitre.openid.connect.util.IdTokenHashUtils;
 import org.mitre.openid.connect.web.AuthenticationTimeStamper;
@@ -94,6 +95,9 @@ public class DefaultOIDCTokenService implements OIDCTokenService {
 	@Autowired
 	private OAuth2TokenEntityService tokenService;
 
+	@Autowired
+	private IDTokenClaimsEnhancer idTokenClaimsEnhancer;
+
 	@Override
 	public JWT createIdToken(ClientDetailsEntity client, OAuth2Request request, Date issueTime, String sub, OAuth2AccessTokenEntity accessToken) {
 
@@ -142,6 +146,8 @@ public class DefaultOIDCTokenService implements OIDCTokenService {
 			idClaims.claim("nonce", nonce);
 		}
 
+		idTokenClaimsEnhancer.enhanceIdTokenClaims(idClaims, request, issueTime, sub, accessToken);
+
 		Set<String> responseTypes = request.getResponseTypes();
 
 		if (responseTypes.contains("token")) {