From c26dc75c44fa27b8bf6e749d1c12fb1b3f869aac Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Fri, 10 May 2019 04:44:01 -0400
Subject: [PATCH] More robust check

---
 .../proxy/util/collect/CappedCollection.java       | 10 +++++-----
 .../proxy/util/collect/CappedCollectionTest.java   | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/util/collect/CappedCollection.java b/proxy/src/main/java/com/velocitypowered/proxy/util/collect/CappedCollection.java
index 64f5e92a5..01645e973 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/util/collect/CappedCollection.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/util/collect/CappedCollection.java
@@ -36,16 +36,16 @@ public class CappedCollection<T> extends ForwardingCollection<T> {
 
   @Override
   public boolean add(T element) {
-    Preconditions.checkState(this.delegate.size() + 1 <= upperSize, "collection is too large (%s)",
-        this.delegate.size());
+    if (this.delegate.size() >= upperSize) {
+      Preconditions.checkState(this.delegate.contains(element), "collection is too large (%s >= %s)",
+          this.delegate.size(), this.upperSize);
+      return false;
+    }
     return this.delegate.add(element);
   }
 
   @Override
   public boolean addAll(Collection<? extends T> collection) {
-    Preconditions.checkState(this.delegate.size() + collection.size() <= upperSize,
-        "collection would grow too large (%s + %s > %s)",
-        this.delegate.size(), collection.size(), upperSize);
     return this.standardAddAll(collection);
   }
 }
diff --git a/proxy/src/test/java/com/velocitypowered/proxy/util/collect/CappedCollectionTest.java b/proxy/src/test/java/com/velocitypowered/proxy/util/collect/CappedCollectionTest.java
index 634cadf2c..f6fbcd919 100644
--- a/proxy/src/test/java/com/velocitypowered/proxy/util/collect/CappedCollectionTest.java
+++ b/proxy/src/test/java/com/velocitypowered/proxy/util/collect/CappedCollectionTest.java
@@ -32,4 +32,18 @@ class CappedCollectionTest {
         "items added to collection although it is too full");
     assertEquals(3, coll.size(), "collection grew in size unexpectedly");
   }
+
+  @Test
+  void handlesSetBehaviorCorrectly() {
+    Set<String> doesFill1 = ImmutableSet.of("coffee", "tea");
+    Set<String> doesFill2 = ImmutableSet.of("coffee", "chocolate");
+    Set<String> overfill = ImmutableSet.of("coffee", "Coke", "Pepsi");
+
+    Collection<String> coll = CappedCollection.newCappedSet(3);
+    assertTrue(coll.addAll(doesFill1), "did not add items");
+    assertTrue(coll.addAll(doesFill2), "did not add items");
+    assertThrows(IllegalStateException.class, () -> coll.addAll(overfill),
+        "items added to collection although it is too full");
+    assertEquals(3, coll.size(), "collection grew in size unexpectedly");
+  }
 }
\ No newline at end of file