reject invalid token type

ant-media · Nov 4, 2024 · f274f91 · f274f91
1 parent 1571e8d
commit f274f91
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 31 deletions.
diff --git a/src/main/java/io/antmedia/rest/RestServiceBase.java b/src/main/java/io/antmedia/rest/RestServiceBase.java
@@ -1736,37 +1736,28 @@ protected ITokenService getTokenService()
 		return null;
 	}
 
-	protected Object getToken (String streamId, long expireDate, String type, String roomId)
-	{
-		Token token = null;
-		String message = "Define Stream ID, Token Type and Expire Date (unix time)";
-		if(streamId != null && type != null && expireDate > 0) {
+	protected Object getToken(String streamId, long expireDate, String type, String roomId) {
+		if (streamId == null || type == null || expireDate <= 0) {
+			return new Result(false, "Define Stream ID, Token Type and Expire Date (unix time)");
+		}
+		if (!type.equals(Token.PLAY_TOKEN) && !type.equals(Token.PUBLISH_TOKEN)) {
+			return new Result(false, "Invalid token type. Supported types are 'play' and 'publish'");
+		}
 
-			ITokenService tokenService = getTokenService();
+		ITokenService tokenService = getTokenService();
+		if (tokenService == null) {
+			return new Result(false, "No token service in this app");
+		}
 
-			if(tokenService != null)
-			{
-				token = tokenService.createToken(streamId, expireDate, type, roomId);
-				if(token != null)
-				{
-					if (getDataStore().saveToken(token)) {
-						//returns token only everything is OK
-						return token;
-					}
-					else {
-						message = "Cannot save token to the datastore";
-					}
-				}
-				else {
-					message = "Cannot create token. It can be a mock token service";
-				}
-			}
-			else {
-				message = "No token service in this app";
-			}
+		Token token = tokenService.createToken(streamId, expireDate, type, roomId);
+		if (token == null) {
+			return new Result(false, "Cannot create token. It may be a mock token service");
+		}
+		if (!getDataStore().saveToken(token)) {
+			return new Result(false, "Cannot save token to the datastore");
 		}
 
-		return new Result(false, message);
+		return token;
 	}
 
 	protected Object getJwtToken (String streamId, long expireDate, String type, String roomId)

diff --git a/src/test/java/io/antmedia/integration/ConsoleAppRestServiceTest.java b/src/test/java/io/antmedia/integration/ConsoleAppRestServiceTest.java
@@ -1373,10 +1373,31 @@ public void testTokenControl() {
 			long expireDate = Instant.now().getEpochSecond() + 1000;
 
 			Broadcast broadcast = RestServiceV2Test.callCreateRegularBroadcast();
+
+			//only token types 'play' and 'publish' are allowed.
+			Token nullAccessToken1 = null;
+			try {
+				nullAccessToken1 = callGetToken("http://localhost:5080/" + appName + "/rest/v2/broadcasts/" + broadcast.getStreamId() + "/token", "Play", expireDate);
+			} catch (Exception e) {
+				log.info(e.getMessage());
+			}
+			assertNull(nullAccessToken1);
+
+			Token nullAccessToken2 = null;
+			try {
+				nullAccessToken2 = callGetToken("http://localhost:5080/" + appName + "/rest/v2/broadcasts/" + broadcast.getStreamId() + "/token", "Publish", expireDate);
+			} catch (Exception e) {
+				log.info(e.getMessage());
+			}
+			assertNull(nullAccessToken2);
+
+
 			Token accessToken = callGetToken( "http://localhost:5080/"+appName+"/rest/v2/broadcasts/"+broadcast.getStreamId()+"/token", Token.PLAY_TOKEN, expireDate);
 			assertNotNull(accessToken);
 
 
+
+
 			Process rtmpSendingProcess = execute(ffmpegPath
 					+ " -re -i src/test/resources/test.flv  -codec copy -f flv rtmp://127.0.0.1/"+ appName + "/"
 					+ broadcast.getStreamId());
@@ -2368,10 +2389,6 @@ public void testRestartPeriod() {
 
 	}
 
-	//public static Token callGetToken(String streamId, String type, long expireDate) throws Exception {
-	//	return callGetToken(SERVICE_URL + "/broadcast/getToken", streamId, type, expireDate);
-	//}
-
 	public static Token callGetToken(String url, String type, long expireDate) throws Exception {
 
 		CloseableHttpClient client = HttpClients.custom().setRedirectStrategy(new LaxRedirectStrategy()).build();
@@ -2387,6 +2404,7 @@ public static Token callGetToken(String url, String type, long expireDate) throw
 		if (response.getStatusLine().getStatusCode() != 200) {
 			throw new Exception(result.toString());
 		}
+
 		System.out.println("result string: " + result.toString());
 
 		return gson.fromJson(result.toString(), Token.class);

diff --git a/src/test/java/io/antmedia/test/rest/BroadcastRestServiceV2UnitTest.java b/src/test/java/io/antmedia/test/rest/BroadcastRestServiceV2UnitTest.java
@@ -576,6 +576,24 @@ public void testGetToken() {
 			assertFalse(result.isSuccess());
 		}
 
+		{
+			//set token type to Publish (must be publish)
+			tokenReturn = restServiceReal.getTokenV2(streamId, 123432, "Publish", "testRoom").getEntity();
+			assertTrue(tokenReturn instanceof Result);
+			result = (Result) tokenReturn;
+			assertFalse(result.isSuccess());
+
+		}
+
+		{
+			//set token type to Play (must be play)
+			tokenReturn = restServiceReal.getTokenV2(streamId, 123432, "Play", "testRoom").getEntity();
+			assertTrue(tokenReturn instanceof Result);
+			result = (Result) tokenReturn;
+			assertFalse(result.isSuccess());
+
+		}
+
 		Mockito.when(datastore.saveToken(Mockito.any())).thenReturn(true);
 		tokenReturn = (Object) restServiceReal.getTokenV2(streamId, 123432, Token.PLAY_TOKEN, "testRoom").getEntity();
 		assertTrue(tokenReturn instanceof Token);