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.
This commit is contained in:
NovaFox161
2020-05-29 23:30:47 -05:00
parent 3ca03e68ac
commit 003dc7130c
10 changed files with 135 additions and 63 deletions

View File

@@ -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();

View File

@@ -410,7 +410,7 @@ public class AnnouncementMessageFormatter {
mentions.append(s).append(" ");
}
return mentions.toString().replaceAll("@", "@\\u200B");
return mentions.toString().replaceAll("@", "@\u200B");
}));
});
}

View File

@@ -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<Calendar> discalService;
private final Map<Snowflake, Mono<GuildSettings>> allSettings = new ConcurrentHashMap<>();
private final Map<Snowflake, Mono<CalendarData>> calendars = new ConcurrentHashMap<>();
private final Map<Snowflake, Mono<Calendar>> customServices = new ConcurrentHashMap<>();
private final Map<Snowflake, Mono<List<Event>>> allEvents = new ConcurrentHashMap<>();
private final Map<Integer, Mono<Calendar>> 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<Void> run() {
@@ -232,7 +237,7 @@ public class AnnouncementThread {
return customServices.get(gs.getGuildID());
}
return discalService;
return discalServices.get(gs.getCredentialsId());
}
private Mono<List<Event>> getEvents(GuildSettings gs, CalendarData cd, Calendar service) {

View File

@@ -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<String> SCOPES = Arrays.asList(CalendarScopes.CALENDAR);
private static final Map<Integer, FileDataStoreFactory> DATA_STORE_FACTORIES;
static {
try {
HTTP_TRANSPORT = GoogleNetHttpTransport.newTrustedTransport();
DATA_STORE_FACTORY = new FileDataStoreFactory(DATA_STORE_DIR);
Map<Integer, FileDataStoreFactory> 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<Credential> authorize() {
private static Mono<Credential> 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<Calendar> getCalendarService(GuildSettings g) {
public static Mono<Calendar> 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<Calendar> 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();
}
}

View File

@@ -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));

View File

@@ -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,

View File

@@ -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");

View File

@@ -0,0 +1,5 @@
# noinspection SqlResolveForFile
ALTER TABLE `${prefix}guild_settings`
ADD COLUMN CREDENTIALS_ID INTEGER NOT NULL DEFAULT 0
AFTER PRIVATE_KEY;

View File

@@ -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"));

View File

@@ -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);
}