jwt blacklist for inmemory db

Author: lastpeony

src/main/java/io/antmedia/datastore/db/InMemoryDataStore.java

@@ -4,8 +4,11 @@
 import java.time.Instant;
 import java.util.*;
 import java.util.Map.Entry;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.regex.Pattern;
 
+import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
 import io.antmedia.rest.model.Result;
 import io.antmedia.security.ITokenService;
 import org.apache.commons.io.FilenameUtils;
@@ -38,8 +41,11 @@ public class InMemoryDataStore extends DataStore {
 	private Map<String, ConferenceRoom> roomMap = new LinkedHashMap<>();
 	private Map<String, WebRTCViewerInfo> webRTCViewerMap = new LinkedHashMap<>();
 
+	private Gson gson;
 
 	public InMemoryDataStore(String dbName) {
+		GsonBuilder builder = new GsonBuilder();
+		gson = builder.create();
 		available = true;
 	}
 
@@ -899,22 +905,65 @@ public boolean deleteToken(String tokenId) {
 
 	@Override
 	public boolean whiteListToken(String tokenId) {
+			Token token = getToken(tokenId);
+			if(token != null && token.isBlackListed()){
+				token.setBlackListed(false);
+				return saveToken(token);
+			}
+
+
 		return false;
 	}
 
 	@Override
 	public List<String> getBlackListedTokens() {
-		return Collections.emptyList();
+
+		ArrayList<String> tokenBlacklist = new ArrayList<>();
+			tokenMap.forEach((tokenId, token) -> {
+				if(token.isBlackListed()){
+					tokenBlacklist.add(tokenId);
+				}
+			});
+			return tokenBlacklist;
+
 	}
 
 	@Override
 	public Result deleteAllBlacklistedExpiredTokens(ITokenService tokenService) {
-		return null;
+		logger.info("Deleting all expired JWTs from token storage.");
+		AtomicInteger deletedTokenCount = new AtomicInteger();
+
+			tokenMap.forEach((tokenId, token) -> {
+				if(token.isBlackListed() && !tokenService.verifyJwt(tokenId,token.getStreamId(),token.getType())){
+					if(deleteToken(tokenId)){
+						deletedTokenCount.getAndIncrement();
+					}else{
+						logger.warn("Couldn't delete JWT:{}", tokenId);
+					}
+				}
+			});
+
+
+		if(deletedTokenCount.get() > 0){
+			final String successMsg = deletedTokenCount+" JWT deleted successfully from storage.";
+			logger.info(successMsg);
+			return new Result(true, successMsg);
+		}else{
+			final String failMsg = "No JWT deleted from storage.";
+			logger.warn(failMsg);
+			return new Result(false, failMsg);
+		}
 	}
 
 	@Override
 	public boolean whiteListAllTokens() {
-		throw new UnsupportedOperationException("");
+			tokenMap.forEach((tokenId, token) -> {
+				if(token.isBlackListed()){
+					whiteListToken(tokenId);
+				}
+			});
+
+		return true;
 	}
 
 	@Override
@@ -1060,11 +1109,22 @@ public boolean updateStreamMetaData(String streamId, String metaData) {
 
 	@Override
 	public boolean blackListToken(Token token) {
-        return false;
+		boolean result = false;
+
+			if (token.getStreamId() != null && token.getTokenId() != null) {
+				token.setBlackListed(true);
+				return saveToken(token);
+			}
+
+		return result;
     }
 
 	@Override
 	public Token getBlackListedToken(String tokenId) {
+		Token token = getToken(tokenId);
+		if(token != null && token.isBlackListed()){
+			return token;
+		}
 		return null;
 	}
 }

src/main/java/io/antmedia/rest/BroadcastRestService.java

@@ -613,7 +613,7 @@ public Result validateTokenV2(@ApiParam(value = "Token to be validated", require
 	@Consumes(MediaType.APPLICATION_JSON)
 	@Path("/jwt-black-list")
 	@Produces(MediaType.APPLICATION_JSON)
-	public Result blackListJwt(@ApiParam(value = "jwt to be added to blacklist.", required = true) @RequestBody Jwt jwt)
+	public Result blackListJwt(@ApiParam(value = "jwt to be added to blacklist.", required = true) Jwt jwt)
 	{
 		if(getAppSettings().isJwtBlacklistEnabled()){
 
@@ -662,14 +662,14 @@ public Result blackListJwt(@ApiParam(value = "jwt to be added to blacklist.", re
 	@Consumes(MediaType.APPLICATION_JSON)
 	@Path("/jwt-black-list")
 	@Produces(MediaType.APPLICATION_JSON)
-	public Result whiteListJwt(@ApiParam(value = "Jwt to be removed from blacklist.", required = true) @QueryParam("jwt") String jwt)
+	public Result whiteListJwt(@ApiParam(value = "Jwt to be removed from blacklist.", required = true) Jwt jwt)
 	{
 		if(getAppSettings().isJwtBlacklistEnabled()){
 
-			if(getDataStore().getBlackListedToken(jwt) == null){
+			if(getDataStore().getBlackListedToken(jwt.getJwt()) == null){
 				return new Result(false, "JWT does not exist in blacklist.");
 
-			}else if(getDataStore().whiteListToken(jwt)){
+			}else if(getDataStore().whiteListToken(jwt.getJwt())){
 				return new Result(true, "JWT successfully removed from blacklist.");
 
 			}else{
@@ -719,7 +719,7 @@ public Result clearJwtBlacklist()
 		if(getAppSettings().isJwtBlacklistEnabled()) {
 			getDataStore().whiteListAllTokens();
 			if(getDataStore().getBlackListedTokens().isEmpty()){
-				return new Result(true, "All blacklisted tokens are whitelisted successfully.");
+				return new Result(true, "All blacklisted tokens are removed successfully.");
 			}else{
 				return new Result(false, "JWT blacklist clear failed.");
 			}

src/main/java/io/antmedia/rest/model/Jwt.java

@@ -1,6 +1,13 @@
 package io.antmedia.rest.model;
 
+import dev.morphia.annotations.Entity;
+import io.swagger.annotations.ApiModel;
+import io.swagger.annotations.ApiModelProperty;
+
+@ApiModel(value="jwt", description="The basic jwt class for jwt blacklist")
+@Entity(value = "jwt")
 public class Jwt {
+    @ApiModelProperty(value = "the jwt")
     private String jwt;
 
     // Default constructor

src/test/java/io/antmedia/test/rest/BroadcastRestServiceV2UnitTest.java

@@ -3192,27 +3192,29 @@ public void testWhitelistJwt() {
 
 		restServiceSpy.setAppSettings(appSettings);
 		String jwtStr = "test-jwt";
+		Jwt jwt = new Jwt();
+		jwt.setJwt(jwtStr);
 
 		DataStore store = mock(MapDBStore.class);
 
 		restServiceSpy.setDataStore(store);
 
-		Result result1 = restServiceSpy.whiteListJwt(jwtStr);
+		Result result1 = restServiceSpy.whiteListJwt(jwt);
 		assertFalse(result1.isSuccess());
 
 		when(appSettings.isJwtBlacklistEnabled()).thenReturn(true);
 
 		when(store.getBlackListedToken(jwtStr)).thenReturn(null);
-		Result result2 = restServiceSpy.whiteListJwt(jwtStr);
+		Result result2 = restServiceSpy.whiteListJwt(jwt);
 		assertFalse(result2.isSuccess());
 
 		Token token = mock(Token.class);
 		when(store.getBlackListedToken(jwtStr)).thenReturn(token);
-		Result result3 = restServiceSpy.whiteListJwt(jwtStr);
+		Result result3 = restServiceSpy.whiteListJwt(jwt);
 		assertFalse(result3.isSuccess());
 
 		when(store.whiteListToken(jwtStr)).thenReturn(true);
-		Result result4 = restServiceSpy.whiteListJwt(jwtStr);
+		Result result4 = restServiceSpy.whiteListJwt(jwt);
 		assertTrue(result4.isSuccess());
 
 	}