Enforce write interest registration state

TODO.md

@@ -210,16 +210,6 @@ These are bugs, correctness issues, or missing functionality that may affect pro
 
 ---
 
-### 17. SocketWrapperBase Write Interest Enforcement (1 item)
-
-| # | File:Line | Description | Fix Idea | Effort | Difficulty |
-|---|-----------|-------------|----------|--------|------------|
-| 17.1 | `SocketWrapperBase.java:511` | `isReadyForWrite()` restriction not enforced in `registerWriteInterest()` | Add a state guard in `registerWriteInterest()` that throws `IllegalStateException` if called when a pending write callback hasn't fired. | 1 day | Medium |
-
-**Total estimated effort: 1 day, Medium difficulty**
-
----
-
 ### 18. WebSocket POJO Handler Accessibility (1 item)
 
 | # | File:Line | Description | Fix Idea | Effort | Difficulty |

java/org/apache/tomcat/util/net/LocalStrings.properties

@@ -146,6 +146,7 @@ sniExtractor.tooEarly=It is illegal to call this method before the client hello
 
 socket.closed=The socket associated with this connection has been closed.
 socket.sslreneg=Exception re-negotiating SSL connection
+socket.writeInterest=Write interest has already been registered for this connection
 
 socketProperties.negativeUnlockTimeout=The negative value for unlockTimeout has been ignored
 

java/org/apache/tomcat/util/net/NioEndpoint.java

@@ -1159,6 +1159,7 @@ protected void processKey(SelectionKey sk, NioSocketWrapper socketWrapper) {
                                     }
                                 } else if (socketWrapper.writeBlocking) {
                                     synchronized (socketWrapper.writeLock) {
+                                        socketWrapper.clearWriteInterest();
                                         socketWrapper.writeBlocking = false;
                                         socketWrapper.writeLock.notify();
                                     }
@@ -1886,7 +1887,7 @@ public void registerReadInterest() {
 
 
         @Override
-        public void registerWriteInterest() {
+        protected void doRegisterWriteInterest() {
             if (log.isTraceEnabled()) {
                 log.trace(sm.getString("endpoint.debug.registerWrite", this));
             }
@@ -2041,6 +2042,9 @@ protected boolean hasOutboundRemaining() {
 
             @Override
             public void run() {
+                if (!read && !inline) {
+                    clearWriteInterest();
+                }
                 // Perform the IO operation
                 // Called from the poller to continue the IO operation
                 long nBytes = 0;

java/org/apache/tomcat/util/net/SocketProcessorBase.java

@@ -71,6 +71,9 @@ public final void run() {
             if (socketWrapper.isClosed()) {
                 return;
             }
+            if (event == SocketEvent.OPEN_WRITE) {
+                socketWrapper.clearWriteInterest();
+            }
             doRun();
         } finally {
             lock.unlock();

java/org/apache/tomcat/util/net/SocketWrapperBase.java

@@ -68,6 +68,8 @@ public abstract class SocketWrapperBase<E> {
     /** Indicates whether the socket has been closed. */
     protected final AtomicBoolean closed = new AtomicBoolean(false);
 
+    private final AtomicBoolean writeInterest = new AtomicBoolean(false);
+
     // Volatile because I/O and setting the timeout values occurs on a different
     // thread to the thread checking the timeout.
     /** Read timeout in milliseconds. */
@@ -507,9 +509,8 @@ public boolean hasDataToWrite() {
      * Checks to see if there are any writes pending and if there are calls {@link #registerWriteInterest()} to trigger
      * a callback once the pending writes have completed.
      * <p>
-     * Note: Once this method has returned <code>false</code> it <b>MUST NOT</b> be called again until the pending write
-     * has completed and the callback has been fired. TODO: Modify {@link #registerWriteInterest()} so the above
-     * restriction is enforced there rather than relying on the caller.
+     * Once this method has returned <code>false</code>, it must not be called again until the pending write has completed
+     * and the callback has been fired.
      *
      * @return <code>true</code> if no writes are pending and data can be written otherwise <code>false</code>
      */
@@ -1000,8 +1001,28 @@ public void processSocket(SocketEvent socketStatus, boolean dispatch) {
 
     /**
      * Registers interest in write events.
+     *
+     * @throws IllegalStateException If write interest has already been registered and the associated callback has not
+     *                                   started
+     */
+    public final void registerWriteInterest() {
+        if (!writeInterest.compareAndSet(false, true)) {
+            throw new IllegalStateException(sm.getString("socket.writeInterest"));
+        }
+        doRegisterWriteInterest();
+    }
+
+    /**
+     * Clears the write interest registration when write event processing starts.
+     */
+    final void clearWriteInterest() {
+        writeInterest.set(false);
+    }
+
+    /**
+     * Registers interest in write events with the endpoint implementation.
      */
-    public abstract void registerWriteInterest();
+    protected abstract void doRegisterWriteInterest();
 
     /**
      * Creates a sendfile data object for the specified file.

test/org/apache/tomcat/util/net/TestSocketWrapperBase.java

@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.tomcat.util.net;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import org.easymock.EasyMock;
+
+public class TestSocketWrapperBase {
+
+    @Test
+    public void testDuplicateWriteInterestRegistration() {
+        SocketWrapperBase<?> socketWrapper = EasyMock.createMockBuilder(SocketWrapperBase.class)
+                .withConstructor(Object.class, AbstractEndpoint.class)
+                .withArgs(new Object(), EasyMock.createNiceMock(AbstractEndpoint.class))
+                .addMockedMethod("doRegisterWriteInterest")
+                .createMock();
+
+        socketWrapper.doRegisterWriteInterest();
+        EasyMock.expectLastCall().once();
+        EasyMock.replay(socketWrapper);
+
+        socketWrapper.registerWriteInterest();
+        Assert.assertThrows(IllegalStateException.class, socketWrapper::registerWriteInterest);
+
+        EasyMock.verify(socketWrapper);
+    }
+
+    @Test
+    public void testWriteInterestRegistrationAfterClear() {
+        SocketWrapperBase<?> socketWrapper = EasyMock.createMockBuilder(SocketWrapperBase.class)
+                .withConstructor(Object.class, AbstractEndpoint.class)
+                .withArgs(new Object(), EasyMock.createNiceMock(AbstractEndpoint.class))
+                .addMockedMethod("doRegisterWriteInterest")
+                .createMock();
+
+        socketWrapper.doRegisterWriteInterest();
+        EasyMock.expectLastCall().times(2);
+        EasyMock.replay(socketWrapper);
+
+        socketWrapper.registerWriteInterest();
+        socketWrapper.clearWriteInterest();
+        socketWrapper.registerWriteInterest();
+
+        EasyMock.verify(socketWrapper);
+    }
+}

webapps/docs/changelog.xml

@@ -301,6 +301,10 @@
       <update>
         Remove support for HTTP 0.9. (markt)
       </update>
+      <fix>
+        Enforce that write interest is not registered more than once before
+        write event processing starts. (sainadh777)
+      </fix>
       <!-- Entries for backport and removal before 12.0.0-M1 below this line -->
     </changelog>
   </subsection>