[UNDERTOW-1497] add request body ready bytes attribute

Author: baranowb

core/src/main/java/io/undertow/attribute/BytesReadAttribute.java

@@ -0,0 +1,84 @@
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2024 Red Hat, Inc., and individual contributors
+ * as indicated by the @author tags.
+ *
+ * Licensed 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 io.undertow.attribute;
+
+import io.undertow.server.HttpServerExchange;
+
+/**
+ * The bytes sent
+ *
+ * @author baranowb
+ */
+public class BytesReadAttribute implements ExchangeAttribute {
+
+    public static final String BYTES_READ_SHORT_UPPER = "%X";
+    public static final String BYTES_READ_SHORT_LOWER = "%x";
+    public static final String BYTES_READ = "%{BYTES_READ}";
+
+    private final boolean dashIfZero;
+
+    public BytesReadAttribute(boolean dashIfZero) {
+        this.dashIfZero = dashIfZero;
+    }
+
+
+    @Override
+    public String readAttribute(final HttpServerExchange exchange) {
+        if (dashIfZero )  {
+            long bytesSent = exchange.getRequestBytesRead();
+            return bytesSent == 0 ? "-" : Long.toString(bytesSent);
+        } else {
+            return Long.toString(exchange.getRequestBytesRead());
+        }
+    }
+
+    @Override
+    public void writeAttribute(final HttpServerExchange exchange, final String newValue) throws ReadOnlyAttributeException {
+        throw new ReadOnlyAttributeException("Bytes read", newValue);
+    }
+
+    @Override
+    public String toString() {
+        return BYTES_READ;
+    }
+
+    public static final class Builder implements ExchangeAttributeBuilder {
+
+        @Override
+        public String name() {
+            return "Bytes Read";
+        }
+
+        @Override
+        public ExchangeAttribute build(final String token) {
+            if(token.equals(BYTES_READ_SHORT_LOWER)) {
+                return new BytesReadAttribute(true);
+            }
+            if (token.equals(BYTES_READ) || token.equals(BYTES_READ_SHORT_UPPER)) {
+                return new BytesReadAttribute(false);
+            }
+            return null;
+        }
+
+        @Override
+        public int priority() {
+            return 0;
+        }
+    }
+}

core/src/main/java/io/undertow/io/UndertowInputStream.java

@@ -19,6 +19,7 @@
 package io.undertow.io;
 
 import io.undertow.UndertowMessages;
+import io.undertow.server.Connectors;
 import io.undertow.server.HttpServerExchange;
 import io.undertow.connector.ByteBufferPool;
 import io.undertow.connector.PooledByteBuffer;
@@ -49,6 +50,7 @@ public class UndertowInputStream extends InputStream {
     private final StreamSourceChannel channel;
     private final ByteBufferPool bufferPool;
     private final int readTimeout;
+    private final HttpServerExchange exchange;
 
     /**
      * If this stream is ready for a read
@@ -60,6 +62,7 @@ public class UndertowInputStream extends InputStream {
     private PooledByteBuffer pooled;
 
     public UndertowInputStream(final HttpServerExchange exchange) {
+        this.exchange = exchange;
         if (exchange.isRequestChannelAvailable()) {
             this.channel = exchange.getRequestChannel();
         } else {
@@ -123,6 +126,7 @@ public int read(final byte[] b, final int off, final int len) throws IOException
             pooled.close();
             pooled = null;
         }
+        Connectors.updateRequestBytesRead(exchange, copied);
         return copied;
     }
 

core/src/main/java/io/undertow/server/Connectors.java

@@ -610,6 +610,10 @@ public static void updateResponseBytesSent(HttpServerExchange exchange, long byt
         exchange.updateBytesSent(bytes);
     }
 
+    public static void updateRequestBytesRead(HttpServerExchange exchange, long bytes) {
+        exchange.updateBytesRead(bytes);
+    }
+
     public static ConduitStreamSinkChannel getConduitSinkChannel(HttpServerExchange exchange) {
         return exchange.getConnection().getSinkChannel();
     }

core/src/main/java/io/undertow/server/HttpServerExchange.java

@@ -262,6 +262,11 @@ public final class HttpServerExchange extends AbstractAttachable {
      */
     private long responseBytesSent = 0;
 
+    /**
+     * The number of bytes that have been read from remote client. This does not include headers,
+     * only the entity body, and does not take any transfer or content encoding into account.
+     */
+    private long requestBytesRead = 0;
 
     private static final int MASK_RESPONSE_CODE = intBitMask(0, 9);
 
@@ -761,6 +766,14 @@ public long getResponseBytesSent() {
         }
     }
 
+    /**
+     *
+     * @return numbers of bytes read from request body
+     */
+    public long getRequestBytesRead() {
+        return this.requestBytesRead;
+    }
+
     /**
      * Updates the number of response bytes sent. Used when compression is in use
      * @param bytes The number of bytes to increase the response size by. May be negative
@@ -771,6 +784,10 @@ void updateBytesSent(long bytes) {
         }
     }
 
+    void updateBytesRead(long bytes) {
+        requestBytesRead += bytes;
+    }
+
     public HttpServerExchange setPersistent(final boolean persistent) {
         if (persistent) {
             setFlags(FLAG_PERSISTENT);

core/src/main/resources/META-INF/services/io.undertow.attribute.ExchangeAttributeBuilder

@@ -8,6 +8,7 @@ io.undertow.attribute.RequestMethodAttribute$Builder
 io.undertow.attribute.QueryStringAttribute$Builder
 io.undertow.attribute.RequestLineAttribute$Builder
 io.undertow.attribute.BytesSentAttribute$Builder
+io.undertow.attribute.BytesReadAttribute$Builder
 io.undertow.attribute.DateTimeAttribute$Builder
 io.undertow.attribute.RemoteUserAttribute$Builder
 io.undertow.attribute.RequestURLAttribute$Builder

servlet/src/main/java/io/undertow/servlet/spec/ServletInputStreamImpl.java

@@ -35,6 +35,7 @@
 import org.xnio.channels.StreamSourceChannel;
 import io.undertow.connector.ByteBufferPool;
 import io.undertow.connector.PooledByteBuffer;
+import io.undertow.server.Connectors;
 import io.undertow.servlet.UndertowServletMessages;
 
 /**
@@ -191,6 +192,7 @@ public int read(final byte[] b, final int off, final int len) throws IOException
                 readIntoBufferNonBlocking();
             }
         }
+        Connectors.updateRequestBytesRead(request.getExchange(), copied);
         return copied;
     }