From d3d8f1a487dddf49b6d6217a9f165823f1d58701 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@vmware.com>
Date: Tue, 10 Nov 2020 19:59:36 +0000
Subject: [PATCH] LimitedDataBufferList adds or raises error

Closes gh-26060
---
 .../core/io/buffer/LimitedDataBufferList.java  |  9 +++------
 .../core/io/buffer/DataBufferUtilsTests.java   | 18 ++++++++++++++++++
 .../io/buffer/LimitedDataBufferListTests.java  | 10 +++++++---
 3 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/spring-core/src/main/java/org/springframework/core/io/buffer/LimitedDataBufferList.java b/spring-core/src/main/java/org/springframework/core/io/buffer/LimitedDataBufferList.java
index fb8c42aeeb0..d95e426d385 100644
--- a/spring-core/src/main/java/org/springframework/core/io/buffer/LimitedDataBufferList.java
+++ b/spring-core/src/main/java/org/springframework/core/io/buffer/LimitedDataBufferList.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -54,11 +54,8 @@ public class LimitedDataBufferList extends ArrayList<DataBuffer> {
 
 	@Override
 	public boolean add(DataBuffer buffer) {
-		boolean result = super.add(buffer);
-		if (result) {
-			updateCount(buffer.readableByteCount());
-		}
-		return result;
+		updateCount(buffer.readableByteCount());
+		return super.add(buffer);
 	}
 
 	@Override
diff --git a/spring-core/src/test/java/org/springframework/core/io/buffer/DataBufferUtilsTests.java b/spring-core/src/test/java/org/springframework/core/io/buffer/DataBufferUtilsTests.java
index 7c71dc8b728..8615551b319 100644
--- a/spring-core/src/test/java/org/springframework/core/io/buffer/DataBufferUtilsTests.java
+++ b/spring-core/src/test/java/org/springframework/core/io/buffer/DataBufferUtilsTests.java
@@ -35,6 +35,8 @@ import java.util.List;
 import java.util.concurrent.CountDownLatch;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.PooledByteBufAllocator;
+import org.junit.jupiter.api.Test;
 import org.mockito.stubbing.Answer;
 import org.reactivestreams.Subscription;
 import reactor.core.publisher.BaseSubscriber;
@@ -834,6 +836,22 @@ class DataBufferUtilsTests extends AbstractDataBufferAllocatingTests {
 				.verifyError(DataBufferLimitException.class);
 	}
 
+	@Test // gh-26060
+	void joinWithLimitDoesNotOverRelease() {
+		NettyDataBufferFactory bufferFactory = new NettyDataBufferFactory(PooledByteBufAllocator.DEFAULT);
+		byte[] bytes = "foo-bar-baz".getBytes(StandardCharsets.UTF_8);
+
+		NettyDataBuffer buffer = bufferFactory.allocateBuffer(bytes.length);
+		buffer.getNativeBuffer().retain(); // should be at 2 now
+		buffer.write(bytes);
+
+		Mono<DataBuffer> result = DataBufferUtils.join(Flux.just(buffer), 8);
+
+		StepVerifier.create(result).verifyError(DataBufferLimitException.class);
+		assertThat(buffer.getNativeBuffer().refCnt()).isEqualTo(1);
+		buffer.release();
+	}
+
 	@ParameterizedDataBufferAllocatingTest
 	void joinErrors(String displayName, DataBufferFactory bufferFactory) {
 		super.bufferFactory = bufferFactory;
diff --git a/spring-core/src/test/java/org/springframework/core/io/buffer/LimitedDataBufferListTests.java b/spring-core/src/test/java/org/springframework/core/io/buffer/LimitedDataBufferListTests.java
index eeb816fe274..cb5c43ebe95 100644
--- a/spring-core/src/test/java/org/springframework/core/io/buffer/LimitedDataBufferListTests.java
+++ b/spring-core/src/test/java/org/springframework/core/io/buffer/LimitedDataBufferListTests.java
@@ -17,9 +17,11 @@ package org.springframework.core.io.buffer;
 
 import java.nio.charset.StandardCharsets;
 
-import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.Test;
 
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
 /**
  * Unit tests for {@link LimitedDataBufferList}.
  * @author Rossen Stoyanchev
@@ -32,8 +34,10 @@ public class LimitedDataBufferListTests {
 
 	@Test
 	void limitEnforced() {
-		Assertions.assertThatThrownBy(() -> new LimitedDataBufferList(5).add(toDataBuffer("123456")))
-				.isInstanceOf(DataBufferLimitException.class);
+		LimitedDataBufferList list = new LimitedDataBufferList(5);
+
+		assertThatThrownBy(() -> list.add(toDataBuffer("123456"))).isInstanceOf(DataBufferLimitException.class);
+		assertThat(list).isEmpty();
 	}
 
 	@Test