Merge pull request #1557 from fl4via/UNDERTOW-2336

[UNDERTOW-2336] CVE-2024-1635 At WriteTimeoutStreamSinkConduit, add a…

Author: fl4via

core/src/main/java/io/undertow/conduits/WriteTimeoutStreamSinkConduit.java

@@ -22,8 +22,8 @@
 import io.undertow.UndertowOptions;
 import io.undertow.server.OpenListener;
 import io.undertow.util.WorkerUtils;
-
 import org.xnio.Buffers;
+import org.xnio.ChannelListener;
 import org.xnio.ChannelListeners;
 import org.xnio.IoUtils;
 import org.xnio.Options;
@@ -47,7 +47,7 @@
  */
 public final class WriteTimeoutStreamSinkConduit extends AbstractStreamSinkConduit<StreamSinkConduit> {
 
-    private XnioExecutor.Key handle;
+    private volatile XnioExecutor.Key handle;
     private final StreamConnection connection;
     private volatile long expireTime = -1;
     private final OpenListener openListener;
@@ -82,6 +82,16 @@ public WriteTimeoutStreamSinkConduit(final StreamSinkConduit delegate, StreamCon
         super(delegate);
         this.connection = connection;
         this.openListener = openListener;
+        this.connection.getCloseSetter().set((ChannelListener<StreamConnection>) channel -> {
+            if (handle != null) {
+                synchronized (WriteTimeoutStreamSinkConduit.this) {
+                    if (handle != null) {
+                        handle.remove();
+                        handle = null;
+                    }
+                }
+            }
+        });
     }
 
     private void handleWriteTimeout(final long ret) throws IOException {
@@ -124,10 +134,14 @@ public long write(final ByteBuffer[] srcs, final int offset, final int length) t
     public int writeFinal(ByteBuffer src) throws IOException {
         int ret = super.writeFinal(src);
         handleWriteTimeout(ret);
-        if(!src.hasRemaining()) {
-            if(handle != null) {
-                handle.remove();
-                handle = null;
+        if (!src.hasRemaining()) {
+            if (handle != null) {
+                synchronized (this) {
+                    if (handle != null) {
+                        handle.remove();
+                        handle = null;
+                    }
+                }
             }
         }
         return ret;
@@ -137,10 +151,14 @@ public int writeFinal(ByteBuffer src) throws IOException {
     public long writeFinal(ByteBuffer[] srcs, int offset, int length) throws IOException {
         long ret = super.writeFinal(srcs, offset, length);
         handleWriteTimeout(ret);
-        if(!Buffers.hasRemaining(srcs)) {
-            if(handle != null) {
-                handle.remove();
-                handle = null;
+        if (!Buffers.hasRemaining(srcs)) {
+            if (handle != null) {
+                synchronized (this) {
+                    if (handle != null) {
+                        handle.remove();
+                        handle = null;
+                    }
+                }
             }
         }
         return ret;
@@ -200,19 +218,33 @@ private Integer getTimeout() {
 
     @Override
     public void terminateWrites() throws IOException {
-        super.terminateWrites();
-        if(handle != null) {
-            handle.remove();
-            handle = null;
+        try {
+            super.terminateWrites();
+        } finally {
+            if(handle != null) {
+                synchronized (this) {
+                    if (this.handle != null) {
+                        handle.remove();
+                        handle = null;
+                    }
+                }
+            }
         }
     }
 
     @Override
     public void truncateWrites() throws IOException {
-        super.truncateWrites();
-        if(handle != null) {
-            handle.remove();
-            handle = null;
+        try {
+            super.truncateWrites();
+        } finally {
+            if (handle != null) {
+                synchronized (this) {
+                    if (this.handle != null) {
+                        handle.remove();
+                        handle = null;
+                    }
+                }
+            }
         }
     }
 
@@ -233,8 +265,12 @@ public void suspendWrites() {
 
         XnioExecutor.Key handle = this.handle;
         if(handle != null) {
-            handle.remove();
-            this.handle = null;
+            synchronized (this) {
+                if (this.handle != null) {
+                    handle.remove();
+                    this.handle = null;
+                }
+            }
         }
     }
 
@@ -253,7 +289,11 @@ private void handleResumeTimeout() {
         expireTime = currentTime + timeout;
         XnioExecutor.Key key = handle;
         if (key == null) {
-            handle = connection.getIoThread().executeAfter(timeoutCommand, timeout, TimeUnit.MILLISECONDS);
+            synchronized (this) {
+                if (handle == null) {
+                    handle = connection.getIoThread().executeAfter(timeoutCommand, timeout, TimeUnit.MILLISECONDS);
+                }
+            }
         }
     }
 }