Bugfix: Nested collection replacment

impl/src/main/java/io/jsonwebtoken/impl/lang/DefaultCollectionMutator.java

@@ -23,6 +23,7 @@
 
 import java.util.Collection;
 import java.util.LinkedHashSet;
+import java.util.NoSuchElementException;
 
 public class DefaultCollectionMutator<E, M extends CollectionMutator<E, M>> implements CollectionMutator<E, M> {
 
@@ -46,6 +47,33 @@ private boolean doAdd(E e) {
         return this.collection.add(e);
     }
 
+    public M replace(E e) {
+        if (Objects.isEmpty(e)) {
+            if (e == null)  throw new NullPointerException("Cannot be null.");
+            else            throw new IllegalArgumentException("Cannot be empty.");
+        }
+
+        // Keep a copy of the original, in case remove/add fails, and need to rollback
+        E old = null;
+        for (E element : this.collection) {
+            if (element.equals(e)) {
+                old = element;
+                break;
+            }
+        }
+
+        if (old == null || !this.collection.contains(e)) {
+            String msg = this.getClass() + " does not contain " + e + ".";
+            throw new NoSuchElementException(msg);
+        }
+
+        if (this.collection.remove(e))
+            if (doAdd(e)) changed(); // removed and add successfully, notify changed()
+            else this.collection.add(old); // remove successfully but add failed, add back old (not considered a change)
+
+        return self();
+    }
+
     @Override
     public M add(E e) {
         if (doAdd(e)) changed();

impl/src/test/groovy/io/jsonwebtoken/impl/lang/DefaultCollectionMutatorTest.groovy

@@ -16,10 +16,14 @@
 package io.jsonwebtoken.impl.lang
 
 import io.jsonwebtoken.Identifiable
+import io.jsonwebtoken.Jwts
 import io.jsonwebtoken.lang.Strings
+import io.jsonwebtoken.security.MacAlgorithm
 import org.junit.Before
 import org.junit.Test
 
+import java.lang.reflect.Constructor
+
 import static org.junit.Assert.*
 
 /**
@@ -128,6 +132,35 @@ class DefaultCollectionMutatorTest {
         assertEquals 0, changeCount
     }
 
+    @Test
+    void replace() {
+        Class<?> c = Class.forName("io.jsonwebtoken.impl.security.DefaultMacAlgorithm")
+        Constructor<?> ctor = c.getDeclaredConstructor(String.class, String.class, int.class)
+        ctor.setAccessible(true)
+        MacAlgorithm custom = (MacAlgorithm) ctor.newInstance('HS512', 'HmacSHA512', 80)
+
+        m.add(Jwts.SIG.HS512)
+        m.replace(custom)
+        assertEquals 2, changeCount // replace is count as one change
+        assertEquals 1, m.getCollection().size() // existing is removed as part of replacement
+        assertEquals 80, ((MacAlgorithm) m.getCollection().toArray()[0]).getKeyBitLength()
+    }
+
+    @Test(expected = NoSuchElementException)
+    void replaceMissing() {
+        m.replace('foo')
+    }
+
+    @Test(expected = NullPointerException)
+    void replaceNull() {
+        m.replace(null)
+    }
+
+    @Test(expected = IllegalArgumentException)
+    void replaceEmpty() {
+        m.replace(Strings.EMPTY)
+    }
+
     @Test
     void clear() {
         m.add('one').add('two').add(['three', 'four'])