From 003dc7130c04fc368f2210b97d0ceff21ccb1b5c Mon Sep 17 00:00:00 2001 From: NovaFox161 Date: Fri, 29 May 2020 23:30:47 -0500 Subject: [PATCH] Add support for multiple credentials for discal This should be the way to resolve the API rate limit with calendar creation. By randomly assigning new guilds one of n amount of accounts, we spread out the calendar creation load onto many accounts. This should overall fix the API rate limit of calendar creation per account as well as allow further expansion. --- .../discal/client/DisCalClient.java | 9 +++ .../message/AnnouncementMessageFormatter.java | 2 +- .../announcement/AnnouncementThread.java | 11 +++- .../discal/core/calendar/CalendarAuth.java | 64 +++++++++++------- .../discal/core/database/DatabaseManager.java | 65 ++++++++++--------- .../discal/core/object/BotSettings.java | 4 +- .../discal/core/object/GuildSettings.java | 17 +++++ .../V7__Alter_Guild_Settings_Add_Creds_Id.sql | 5 ++ .../discal/server/DisCalServer.java | 9 +++ .../dreamexposure/discal/web/DisCalWeb.java | 12 +++- 10 files changed, 135 insertions(+), 63 deletions(-) create mode 100644 core/src/main/resources/db/migration/V7__Alter_Guild_Settings_Add_Creds_Id.sql diff --git a/client/src/main/java/org/dreamexposure/discal/client/DisCalClient.java b/client/src/main/java/org/dreamexposure/discal/client/DisCalClient.java index 23f9228a..c9f5a050 100644 --- a/client/src/main/java/org/dreamexposure/discal/client/DisCalClient.java +++ b/client/src/main/java/org/dreamexposure/discal/client/DisCalClient.java @@ -20,6 +20,7 @@ import org.dreamexposure.discal.client.module.command.RsvpCommand; import org.dreamexposure.discal.client.module.command.TimeCommand; import org.dreamexposure.discal.client.service.KeepAliveHandler; import org.dreamexposure.discal.client.service.TimeManager; +import org.dreamexposure.discal.core.calendar.CalendarAuth; import org.dreamexposure.discal.core.database.DatabaseManager; import org.dreamexposure.discal.core.logger.LogFeed; import org.dreamexposure.discal.core.logger.object.LogObject; @@ -65,6 +66,14 @@ public class DisCalClient { p.load(new FileReader(new File("settings.properties"))); BotSettings.init(p); + if (args.length > 1 && args[0].equalsIgnoreCase("-forceNewAuth")) { + //Forcefully start a browser for google account authorization. + CalendarAuth.getCalendarService(Integer.parseInt(args[1])).block(); //Block until auth completes... + + //Kill the running instance as this is only meant for generating new credentials... Illegal State basically. + System.exit(100); + } + //Start Google authorization daemon Authorization.getAuth().init(); diff --git a/client/src/main/java/org/dreamexposure/discal/client/message/AnnouncementMessageFormatter.java b/client/src/main/java/org/dreamexposure/discal/client/message/AnnouncementMessageFormatter.java index 403f0168..18bd5701 100644 --- a/client/src/main/java/org/dreamexposure/discal/client/message/AnnouncementMessageFormatter.java +++ b/client/src/main/java/org/dreamexposure/discal/client/message/AnnouncementMessageFormatter.java @@ -410,7 +410,7 @@ public class AnnouncementMessageFormatter { mentions.append(s).append(" "); } - return mentions.toString().replaceAll("@", "@\\u200B"); + return mentions.toString().replaceAll("@", "@\u200B"); })); }); } diff --git a/client/src/main/java/org/dreamexposure/discal/client/module/announcement/AnnouncementThread.java b/client/src/main/java/org/dreamexposure/discal/client/module/announcement/AnnouncementThread.java index 671978d6..e8ce58d1 100644 --- a/client/src/main/java/org/dreamexposure/discal/client/module/announcement/AnnouncementThread.java +++ b/client/src/main/java/org/dreamexposure/discal/client/module/announcement/AnnouncementThread.java @@ -15,6 +15,7 @@ import org.dreamexposure.discal.core.object.calendar.CalendarData; import org.dreamexposure.discal.core.utils.GlobalConst; import org.dreamexposure.discal.core.wrapper.google.EventWrapper; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -30,16 +31,20 @@ import static reactor.function.TupleUtils.function; public class AnnouncementThread { private final GatewayDiscordClient client; - private final Mono discalService; private final Map> allSettings = new ConcurrentHashMap<>(); private final Map> calendars = new ConcurrentHashMap<>(); private final Map> customServices = new ConcurrentHashMap<>(); private final Map>> allEvents = new ConcurrentHashMap<>(); + private final Map> discalServices = new HashMap<>(); + public AnnouncementThread(GatewayDiscordClient client) { this.client = client; - this.discalService = CalendarAuth.getCalendarService(null).cache(); + + for (int i = 0; i < CalendarAuth.credentialsCount(); i++) { + this.discalServices.put(i, CalendarAuth.getCalendarService(i).cache()); + } } public Mono run() { @@ -232,7 +237,7 @@ public class AnnouncementThread { return customServices.get(gs.getGuildID()); } - return discalService; + return discalServices.get(gs.getCredentialsId()); } private Mono> getEvents(GuildSettings gs, CalendarData cd, Calendar service) { diff --git a/core/src/main/java/org/dreamexposure/discal/core/calendar/CalendarAuth.java b/core/src/main/java/org/dreamexposure/discal/core/calendar/CalendarAuth.java index 2e1de863..19711929 100644 --- a/core/src/main/java/org/dreamexposure/discal/core/calendar/CalendarAuth.java +++ b/core/src/main/java/org/dreamexposure/discal/core/calendar/CalendarAuth.java @@ -17,14 +17,19 @@ import com.google.api.services.calendar.CalendarScopes; import org.dreamexposure.discal.core.crypto.AESEncryption; import org.dreamexposure.discal.core.network.google.Authorization; +import org.dreamexposure.discal.core.object.BotSettings; import org.dreamexposure.discal.core.object.GuildSettings; +import org.jetbrains.annotations.NotNull; import java.io.File; import java.io.FileInputStream; import java.io.InputStream; import java.io.InputStreamReader; import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.function.Function; import reactor.core.publisher.Mono; @@ -42,16 +47,6 @@ public class CalendarAuth { */ private static final String APPLICATION_NAME = "DisCal"; - /** - * Directory to store user credentials for this application. - */ - private static final java.io.File DATA_STORE_DIR = new java.io.File(System.getProperty("user.home"), ".credentials/DisCal"); - - /** - * Global instance of the {@link FileDataStoreFactory}. - */ - private static FileDataStoreFactory DATA_STORE_FACTORY; - /** * Global instance of the JSON factory. */ @@ -60,7 +55,7 @@ public class CalendarAuth { /** * Global instance of the HTTP transport. */ - private static HttpTransport HTTP_TRANSPORT; + private final static HttpTransport HTTP_TRANSPORT; /** * Global instance of the scopes required by this quickstart. @@ -70,13 +65,23 @@ public class CalendarAuth { */ private static final List SCOPES = Arrays.asList(CalendarScopes.CALENDAR); + private static final Map DATA_STORE_FACTORIES; + static { try { HTTP_TRANSPORT = GoogleNetHttpTransport.newTrustedTransport(); - DATA_STORE_FACTORY = new FileDataStoreFactory(DATA_STORE_DIR); + + Map dataStoreFactories = new HashMap<>(); + int credCount = Integer.parseInt(BotSettings.CREDENTIALS_COUNT.get()); + for (int i = 0; i < credCount; i++) { + dataStoreFactories.put(i, new FileDataStoreFactory(getCredentialsFolder(i))); + } + + DATA_STORE_FACTORIES = Collections.unmodifiableMap(dataStoreFactories); } catch (Throwable t) { t.printStackTrace(); System.exit(1); + throw new RuntimeException(t); //Never reached, makes compiler happy :) } } @@ -85,15 +90,19 @@ public class CalendarAuth { * * @return an authorized Credential object. */ - private static Mono authorize() { + private static Mono authorize(int credentialId) { return Mono.fromCallable(() -> { // Load client secrets. - //InputStream in = CalendarAuth.class.getResourceAsStream("/client_secret.json"); <- in case it breaks, this is still here InputStream in = new FileInputStream(new File("client_secret.json")); GoogleClientSecrets clientSecrets = GoogleClientSecrets.load(JSON_FACTORY, new InputStreamReader(in)); // Build flow and trigger user authorization request. - GoogleAuthorizationCodeFlow flow = new GoogleAuthorizationCodeFlow.Builder(HTTP_TRANSPORT, JSON_FACTORY, clientSecrets, SCOPES).setDataStoreFactory(DATA_STORE_FACTORY).setAccessType("offline").build(); + GoogleAuthorizationCodeFlow flow = new GoogleAuthorizationCodeFlow + .Builder(HTTP_TRANSPORT, JSON_FACTORY, clientSecrets, SCOPES) + .setDataStoreFactory(DATA_STORE_FACTORIES.get(credentialId)) + .setAccessType("offline") + .build(); + Credential credential = new AuthorizationCodeInstalledApp(flow, new LocalServerReceiver()).authorize("user"); //Try to close input stream since I don't think it was ever closed? @@ -118,15 +127,8 @@ public class CalendarAuth { }).subscribeOn(Schedulers.boundedElastic()); } - public static Mono getCalendarService(GuildSettings g) { + public static Mono getCalendarService(@NotNull GuildSettings g) { return Mono.fromCallable(() -> { - if (g == null) { - return authorize().map(cred -> - new Calendar.Builder(HTTP_TRANSPORT, JSON_FACTORY, cred) - .setApplicationName(APPLICATION_NAME) - .build()); - } - if (g.useExternalCalendar()) { return authorize(g).map(cred -> new Calendar. @@ -134,11 +136,25 @@ public class CalendarAuth { .setApplicationName(APPLICATION_NAME) .build()); } else { - return authorize().map(cred -> + return authorize(g.getCredentialsId()).map(cred -> new Calendar.Builder(HTTP_TRANSPORT, JSON_FACTORY, cred) .setApplicationName(APPLICATION_NAME) .build()); } }).flatMap(Function.identity()); } + + public static Mono getCalendarService(int credentialId) { + return authorize(credentialId).map(cred -> new Calendar.Builder(HTTP_TRANSPORT, JSON_FACTORY, cred) + .setApplicationName(APPLICATION_NAME) + .build()); + } + + private static File getCredentialsFolder(int credentialId) { + return new File(BotSettings.CREDENTIAL_FOLDER.get() + "/" + credentialId); + } + + public static int credentialsCount() { + return DATA_STORE_FACTORIES.size(); + } } \ No newline at end of file diff --git a/core/src/main/java/org/dreamexposure/discal/core/database/DatabaseManager.java b/core/src/main/java/org/dreamexposure/discal/core/database/DatabaseManager.java index 4a1d5d77..807ced83 100644 --- a/core/src/main/java/org/dreamexposure/discal/core/database/DatabaseManager.java +++ b/core/src/main/java/org/dreamexposure/discal/core/database/DatabaseManager.java @@ -175,7 +175,7 @@ public class DatabaseManager { .flatMap(exists -> { if (exists) { String update = "UPDATE " + table - + " SET EXTERNAL_CALENDAR = ?, PRIVATE_KEY = ?," + + " SET EXTERNAL_CALENDAR = ?, PRIVATE_KEY = ?, CREDENTIALS_ID = ?," + " ACCESS_TOKEN = ?, REFRESH_TOKEN = ?," + " CONTROL_ROLE = ?, DISCAL_CHANNEL = ?, SIMPLE_ANNOUNCEMENT = ?," + " LANG = ?, PREFIX = ?, PATRON_GUILD = ?, DEV_GUILD = ?," @@ -185,36 +185,7 @@ public class DatabaseManager { return connect(master, c -> Mono.from(c.createStatement(update) .bind(0, set.useExternalCalendar()) .bind(1, set.getPrivateKey()) - .bind(2, set.getEncryptedAccessToken()) - .bind(3, set.getEncryptedRefreshToken()) - .bind(4, set.getControlRole()) - .bind(5, set.getDiscalChannel()) - .bind(6, set.usingSimpleAnnouncements()) - .bind(7, set.getLang()) - .bind(8, set.getPrefix()) - .bind(9, set.isPatronGuild()) - .bind(10, set.isDevGuild()) - .bind(11, set.getMaxCalendars()) - .bind(12, set.getDmAnnouncementsString()) - .bind(13, set.useTwelveHour()) - .bind(14, set.isBranded()) - .bind(15, set.getGuildID().asString()) - .execute()) - ).flatMap(res -> Mono.from(res.getRowsUpdated())) - .hasElement() - .thenReturn(true); - } else { - String insertCommand = "INSERT INTO " + table + "(GUILD_ID, " + - "EXTERNAL_CALENDAR, PRIVATE_KEY, ACCESS_TOKEN, REFRESH_TOKEN, " + - "CONTROL_ROLE, DISCAL_CHANNEL, SIMPLE_ANNOUNCEMENT, LANG, " + - "PREFIX, PATRON_GUILD, DEV_GUILD, MAX_CALENDARS, " + - "DM_ANNOUNCEMENTS, 12_HOUR, BRANDED) " + - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"; - - return connect(master, c -> Mono.from(c.createStatement(insertCommand) - .bind(0, set.getGuildID().asString()) - .bind(1, set.useExternalCalendar()) - .bind(2, set.getPrivateKey()) + .bind(2, set.getCredentialsId()) .bind(3, set.getEncryptedAccessToken()) .bind(4, set.getEncryptedRefreshToken()) .bind(5, set.getControlRole()) @@ -228,6 +199,37 @@ public class DatabaseManager { .bind(13, set.getDmAnnouncementsString()) .bind(14, set.useTwelveHour()) .bind(15, set.isBranded()) + .bind(16, set.getGuildID().asString()) + .execute()) + ).flatMap(res -> Mono.from(res.getRowsUpdated())) + .hasElement() + .thenReturn(true); + } else { + String insertCommand = "INSERT INTO " + table + "(GUILD_ID, " + + "EXTERNAL_CALENDAR, PRIVATE_KEY, CREDENTIALS_ID, ACCESS_TOKEN, REFRESH_TOKEN, " + + "CONTROL_ROLE, DISCAL_CHANNEL, SIMPLE_ANNOUNCEMENT, LANG, " + + "PREFIX, PATRON_GUILD, DEV_GUILD, MAX_CALENDARS, " + + "DM_ANNOUNCEMENTS, 12_HOUR, BRANDED) " + + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"; + + return connect(master, c -> Mono.from(c.createStatement(insertCommand) + .bind(0, set.getGuildID().asString()) + .bind(1, set.useExternalCalendar()) + .bind(2, set.getPrivateKey()) + .bind(3, set.getCredentialsId()) + .bind(4, set.getEncryptedAccessToken()) + .bind(5, set.getEncryptedRefreshToken()) + .bind(6, set.getControlRole()) + .bind(7, set.getDiscalChannel()) + .bind(8, set.usingSimpleAnnouncements()) + .bind(9, set.getLang()) + .bind(10, set.getPrefix()) + .bind(11, set.isPatronGuild()) + .bind(12, set.isDevGuild()) + .bind(13, set.getMaxCalendars()) + .bind(14, set.getDmAnnouncementsString()) + .bind(15, set.useTwelveHour()) + .bind(16, set.isBranded()) .execute()) ).flatMap(res -> Mono.from(res.getRowsUpdated())) .hasElement() @@ -502,6 +504,7 @@ public class DatabaseManager { set.setUseExternalCalendar(row.get("EXTERNAL_CALENDAR", Boolean.class)); set.setPrivateKey(row.get("PRIVATE_KEY", String.class)); + set.setCredentialsId(row.get("CREDENTIALS_ID", Integer.class)); set.setEncryptedAccessToken(row.get("ACCESS_TOKEN", String.class)); set.setEncryptedRefreshToken(row.get("REFRESH_TOKEN", String.class)); set.setControlRole(row.get("CONTROL_ROLE", String.class)); diff --git a/core/src/main/java/org/dreamexposure/discal/core/object/BotSettings.java b/core/src/main/java/org/dreamexposure/discal/core/object/BotSettings.java index edc3ce7d..f66b23ac 100644 --- a/core/src/main/java/org/dreamexposure/discal/core/object/BotSettings.java +++ b/core/src/main/java/org/dreamexposure/discal/core/object/BotSettings.java @@ -23,11 +23,11 @@ public enum BotSettings { TOKEN, SECRET, ID, - GOOGLE_CLIENT_ID, GOOGLE_CLIENT_SECRET, + GOOGLE_CLIENT_ID, GOOGLE_CLIENT_SECRET, CREDENTIALS_COUNT, SHARD_COUNT, SHARD_INDEX, - LANG_PATH, LOG_FOLDER, + LANG_PATH, LOG_FOLDER, CREDENTIAL_FOLDER, PW_TOKEN, DBO_TOKEN, diff --git a/core/src/main/java/org/dreamexposure/discal/core/object/GuildSettings.java b/core/src/main/java/org/dreamexposure/discal/core/object/GuildSettings.java index d0c8ddb3..614c55e2 100644 --- a/core/src/main/java/org/dreamexposure/discal/core/object/GuildSettings.java +++ b/core/src/main/java/org/dreamexposure/discal/core/object/GuildSettings.java @@ -1,9 +1,11 @@ package org.dreamexposure.discal.core.object; +import org.dreamexposure.discal.core.calendar.CalendarAuth; import org.json.JSONObject; import java.util.ArrayList; import java.util.Collections; +import java.util.Random; import discord4j.common.util.Snowflake; @@ -19,6 +21,8 @@ public class GuildSettings { private boolean externalCalendar; private String privateKey; + private int credentialsId; + private String encryptedAccessToken; private String encryptedRefreshToken; @@ -44,6 +48,8 @@ public class GuildSettings { externalCalendar = false; privateKey = "N/a"; + credentialsId = new Random().nextInt(CalendarAuth.credentialsCount()); + encryptedAccessToken = "N/a"; encryptedRefreshToken = "N/a"; @@ -74,6 +80,10 @@ public class GuildSettings { return privateKey; } + public int getCredentialsId() { + return credentialsId; + } + public String getEncryptedAccessToken() { return encryptedAccessToken; } @@ -150,6 +160,10 @@ public class GuildSettings { privateKey = _privateKey; } + public void setCredentialsId(int credentialsId) { + this.credentialsId = credentialsId; + } + public void setEncryptedAccessToken(String _access) { encryptedAccessToken = _access; } @@ -209,6 +223,7 @@ public class GuildSettings { data.put("guild_id", guildID.asString()); data.put("external_calendar", externalCalendar); data.put("private_key", privateKey); + data.put("credentials_id", credentialsId); data.put("access_token", encryptedAccessToken); data.put("refresh_token", encryptedRefreshToken); data.put("control_role", controlRole); @@ -248,6 +263,7 @@ public class GuildSettings { guildID = Snowflake.of(data.getString("guild_id")); externalCalendar = data.getBoolean("external_calendar"); privateKey = data.getString("private_key"); + credentialsId = data.getInt("credentials_id"); encryptedAccessToken = data.getString("access_token"); encryptedRefreshToken = data.getString("refresh_token"); controlRole = data.getString("control_role"); @@ -267,6 +283,7 @@ public class GuildSettings { public GuildSettings fromJsonSecure(JSONObject data) { guildID = Snowflake.of(data.getString("guild_id")); externalCalendar = data.getBoolean("external_calendar"); + //credentialsId = data.getInt("credentials_id"); //privateKey = data.getString("PrivateKey"); //encryptedAccessToken = data.getString("AccessToken"); //encryptedRefreshToken = data.getString("RefreshToken"); diff --git a/core/src/main/resources/db/migration/V7__Alter_Guild_Settings_Add_Creds_Id.sql b/core/src/main/resources/db/migration/V7__Alter_Guild_Settings_Add_Creds_Id.sql new file mode 100644 index 00000000..a6d6c100 --- /dev/null +++ b/core/src/main/resources/db/migration/V7__Alter_Guild_Settings_Add_Creds_Id.sql @@ -0,0 +1,5 @@ +# noinspection SqlResolveForFile + +ALTER TABLE `${prefix}guild_settings` + ADD COLUMN CREDENTIALS_ID INTEGER NOT NULL DEFAULT 0 + AFTER PRIVATE_KEY; \ No newline at end of file diff --git a/server/src/main/java/org/dreamexposure/discal/server/DisCalServer.java b/server/src/main/java/org/dreamexposure/discal/server/DisCalServer.java index 80da668d..199260e2 100644 --- a/server/src/main/java/org/dreamexposure/discal/server/DisCalServer.java +++ b/server/src/main/java/org/dreamexposure/discal/server/DisCalServer.java @@ -1,5 +1,6 @@ package org.dreamexposure.discal.server; +import org.dreamexposure.discal.core.calendar.CalendarAuth; import org.dreamexposure.discal.core.database.DatabaseManager; import org.dreamexposure.discal.core.logger.LogFeed; import org.dreamexposure.discal.core.logger.object.LogObject; @@ -39,6 +40,14 @@ public class DisCalServer { p.load(new FileReader(new File("settings.properties"))); BotSettings.init(p); + if (args.length > 1 && args[0].equalsIgnoreCase("-forceNewAuth")) { + //Forcefully start a browser for google account authorization. + CalendarAuth.getCalendarService(Integer.parseInt(args[1])).block(); //Block until auth completes... + + //Kill the running instance as this is only meant for generating new credentials... Illegal State basically. + System.exit(100); + } + //Handle database migrations handleMigrations(args.length > 0 && args[0].equalsIgnoreCase("--repair")); diff --git a/web/src/main/java/org/dreamexposure/discal/web/DisCalWeb.java b/web/src/main/java/org/dreamexposure/discal/web/DisCalWeb.java index 1297a3d0..3769c600 100644 --- a/web/src/main/java/org/dreamexposure/discal/web/DisCalWeb.java +++ b/web/src/main/java/org/dreamexposure/discal/web/DisCalWeb.java @@ -1,5 +1,6 @@ package org.dreamexposure.discal.web; +import org.dreamexposure.discal.core.calendar.CalendarAuth; import org.dreamexposure.discal.core.logger.LogFeed; import org.dreamexposure.discal.core.logger.object.LogObject; import org.dreamexposure.discal.core.network.google.Authorization; @@ -22,6 +23,14 @@ public class DisCalWeb { p.load(new FileReader(new File("settings.properties"))); BotSettings.init(p); + if (args.length > 1 && args[0].equalsIgnoreCase("-forceNewAuth")) { + //Forcefully start a browser for google account authorization. + CalendarAuth.getCalendarService(Integer.parseInt(args[1])).block(); //Block until auth completes... + + //Kill the running instance as this is only meant for generating new credentials... Illegal State basically. + System.exit(100); + } + //Start Google authorization daemon Authorization.getAuth().init(); @@ -33,8 +42,7 @@ public class DisCalWeb { app.run(args); } catch (Exception e) { e.printStackTrace(); - LogFeed.log(LogObject - .forException("'Spring error", "by 'PANIC! AT THE WEBSITE'", e, DisCalWeb.class)); + LogFeed.log(LogObject.forException("'Spring error", "by 'PANIC! AT THE WEBSITE'", e, DisCalWeb.class)); System.exit(4); }