From ac0ffefa3342e91f320652f3d2bce6463b0617b0 Mon Sep 17 00:00:00 2001
From: Kavin <20838718+FireMasterK@users.noreply.github.com>
Date: Tue, 9 May 2023 22:31:17 +0100
Subject: [PATCH] Add optional limit parameter.

---
 .../me/kavin/piped/server/ServerLauncher.java | 42 +++++++++++++++++--
 .../server/handlers/auth/FeedHandlers.java    |  8 ++--
 2 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/src/main/java/me/kavin/piped/server/ServerLauncher.java b/src/main/java/me/kavin/piped/server/ServerLauncher.java
index 64fcf16..ee0fc04 100644
--- a/src/main/java/me/kavin/piped/server/ServerLauncher.java
+++ b/src/main/java/me/kavin/piped/server/ServerLauncher.java
@@ -320,7 +320,19 @@ public class ServerLauncher extends MultithreadedHttpServerLauncher {
                     }
                 })).map(GET, "/feed", AsyncServlet.ofBlocking(executor, request -> {
                     try {
-                        return getJsonResponse(FeedHandlers.feedResponse(request.getQueryParameter("authToken"), request.getQueryParameter("start")),
+
+                        int limit;
+                        try {
+                            var limitStr = request.getQueryParameter("limit");
+                            if (!StringUtils.isEmpty(limitStr))
+                                limit = Math.min(Integer.parseInt(limitStr), 100);
+                            else
+                                limit = 100;
+                        } catch (NumberFormatException e) {
+                            limit = 100;
+                        }
+
+                        return getJsonResponse(FeedHandlers.feedResponse(request.getQueryParameter("authToken"), request.getQueryParameter("start"), limit),
                                 "private");
                     } catch (Exception e) {
                         return getErrorResponse(e, request.getPath());
@@ -334,8 +346,20 @@ public class ServerLauncher extends MultithreadedHttpServerLauncher {
                     }
                 })).map(GET, "/feed/unauthenticated", AsyncServlet.ofBlocking(executor, request -> {
                     try {
+
+                        int limit;
+                        try {
+                            var limitStr = request.getQueryParameter("limit");
+                            if (!StringUtils.isEmpty(limitStr))
+                                limit = Math.min(Integer.parseInt(limitStr), 100);
+                            else
+                                limit = 100;
+                        } catch (NumberFormatException e) {
+                            limit = 100;
+                        }
+
                         return getJsonResponse(FeedHandlers.unauthenticatedFeedResponse(
-                                getArray(request.getQueryParameter("channels")), request.getQueryParameter("start")
+                                getArray(request.getQueryParameter("channels")), request.getQueryParameter("start"), limit
                         ), "public, s-maxage=120");
                     } catch (Exception e) {
                         return getErrorResponse(e, request.getPath());
@@ -344,7 +368,19 @@ public class ServerLauncher extends MultithreadedHttpServerLauncher {
                     try {
                         String[] subscriptions = mapper.readValue(request.loadBody().getResult().asArray(),
                                 String[].class);
-                        return getJsonResponse(FeedHandlers.unauthenticatedFeedResponse(subscriptions, request.getQueryParameter("start")), "public, s-maxage=120");
+
+                        int limit;
+                        try {
+                            var limitStr = request.getQueryParameter("limit");
+                            if (!StringUtils.isEmpty(limitStr))
+                                limit = Math.min(Integer.parseInt(limitStr), 100);
+                            else
+                                limit = 100;
+                        } catch (NumberFormatException e) {
+                            limit = 100;
+                        }
+
+                        return getJsonResponse(FeedHandlers.unauthenticatedFeedResponse(subscriptions, request.getQueryParameter("start"), limit), "public, s-maxage=120");
                     } catch (Exception e) {
                         return getErrorResponse(e, request.getPath());
                     }
diff --git a/src/main/java/me/kavin/piped/server/handlers/auth/FeedHandlers.java b/src/main/java/me/kavin/piped/server/handlers/auth/FeedHandlers.java
index 1c988ee..355ae2c 100644
--- a/src/main/java/me/kavin/piped/server/handlers/auth/FeedHandlers.java
+++ b/src/main/java/me/kavin/piped/server/handlers/auth/FeedHandlers.java
@@ -93,7 +93,7 @@ public class FeedHandlers {
         }
     }
 
-    public static byte[] feedResponse(String session, String start) throws IOException {
+    public static byte[] feedResponse(String session, String start, int limit) throws IOException {
 
         if (StringUtils.isBlank(session))
             ExceptionHandler.throwErrorResponse(new InvalidRequestResponse("session is a required parameter"));
@@ -126,7 +126,7 @@ public class FeedHandlers {
                         .orderBy(cb.desc(root.get("uploaded")));
 
                 List<StreamItem> feedItems = s.createQuery(criteria)
-                        .setMaxResults(100)
+                        .setMaxResults(limit)
                         .setTimeout(20).stream()
                         .parallel().map(video -> {
                             var channel = video.getChannel();
@@ -200,7 +200,7 @@ public class FeedHandlers {
         return null;
     }
 
-    public static byte[] unauthenticatedFeedResponse(String[] channelIds, String start) throws Exception {
+    public static byte[] unauthenticatedFeedResponse(String[] channelIds, String start, int limit) throws Exception {
 
         Set<String> filtered = Arrays.stream(channelIds)
                 .filter(ChannelHelpers::isValidId)
@@ -230,7 +230,7 @@ public class FeedHandlers {
                     .orderBy(cb.desc(root.get("uploaded")));
 
             List<StreamItem> feedItems = s.createQuery(criteria)
-                    .setMaxResults(100)
+                    .setMaxResults(limit)
                     .setTimeout(20).stream()
                     .parallel().map(video -> {
                         var channel = video.getChannel();