Address code review feedback

- Use static imports for constants in test file
- Consolidate SUCCESS/NEED_MORE_BYTES/COMPLETED case handling in policy
- Improve segment size validation against remaining message bytes
- Simplify condition using getAvailableBytes() helper

Co-authored-by: gunjansingh-msft <[email protected]>

Author: Copilot

sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageDecoder.java

@@ -415,8 +415,10 @@ private boolean tryReadSegmentHeader(ByteBuffer buffer) {
                 "Unexpected segment number. Expected: " + (currentSegmentNumber + 1) + ", got: " + segmentNum));
         }
 
-        // Validate segment size
-        if (segmentSize < 0 || segmentSize > Integer.MAX_VALUE) {
+        // Validate segment size - must be non-negative and reasonable
+        // We can't have segments larger than the remaining message length
+        long remainingMessageBytes = messageLength - messageOffset - V1_SEGMENT_HEADER_LENGTH;
+        if (segmentSize < 0 || segmentSize > remainingMessageBytes) {
             LOGGER.error("Invalid segment length read: segmentLength={}, decoderOffset={}, lastCompleteSegment={}",
                 segmentSize, messageOffset, lastCompleteSegmentStart);
             throw LOGGER.logExceptionAsError(new IllegalArgumentException(
@@ -626,8 +628,8 @@ public DecodeResult decodeChunk(ByteBuffer buffer) {
                         "Decode completed");
                 }
 
-                // If we couldn't read any bytes and buffer is empty, need more
-                if (payloadRead == 0 && buffer.remaining() == 0 && pendingBytes.size() == 0) {
+                // If we couldn't read any bytes and no data available, need more
+                if (payloadRead == 0 && getAvailableBytes(buffer) == 0) {
                     break;
                 }
             }

sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicy.java

@@ -107,18 +107,10 @@ private Flux<ByteBuffer> decodeStream(Flux<ByteBuffer> encodedFlux, DecoderState
                 switch (result.getStatus()) {
                     case SUCCESS:
                     case NEED_MORE_BYTES:
-                        // Update state counters
-                        state.totalEncodedBytesProcessed.set(state.decoder.getMessageOffset());
-                        state.totalBytesDecoded.set(state.decoder.getTotalDecodedPayloadBytes());
-                        state.decodedBytesAtLastCompleteSegment = state.decoder.getTotalDecodedPayloadBytes();
-
-                        if (result.getDecodedPayload() != null && result.getDecodedPayload().hasRemaining()) {
-                            return Flux.just(result.getDecodedPayload());
-                        }
-                        return Flux.empty();
-
                     case COMPLETED:
-                        // Update state counters
+                        // All three cases update counters and return any decoded payload
+                        // SUCCESS and NEED_MORE_BYTES: partial decode, more data expected
+                        // COMPLETED: decode finished successfully
                         state.totalEncodedBytesProcessed.set(state.decoder.getMessageOffset());
                         state.totalBytesDecoded.set(state.decoder.getTotalDecodedPayloadBytes());
                         state.decodedBytesAtLastCompleteSegment = state.decoder.getTotalDecodedPayloadBytes();

sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/implementation/structuredmessage/StructuredMessageDecoderTests.java

@@ -10,6 +10,9 @@
 import java.nio.ByteOrder;
 import java.util.concurrent.ThreadLocalRandom;
 
+import static com.azure.storage.common.implementation.structuredmessage.StructuredMessageConstants.CRC64_LENGTH;
+import static com.azure.storage.common.implementation.structuredmessage.StructuredMessageConstants.V1_HEADER_LENGTH;
+import static com.azure.storage.common.implementation.structuredmessage.StructuredMessageConstants.V1_SEGMENT_HEADER_LENGTH;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -175,7 +178,7 @@ public void resetToLastCompleteSegmentWorks() throws IOException {
 
         // Parse first segment completely, then simulate interruption
         // First segment ends after: header(13) + segment_header(10) + content(256) + crc(8) = 287
-        int firstSegmentEnd = 13 + 10 + 256 + 8;
+        int firstSegmentEnd = V1_HEADER_LENGTH + V1_SEGMENT_HEADER_LENGTH + 256 + CRC64_LENGTH;
         ByteBuffer chunk1 = ByteBuffer.wrap(encodedBytes, 0, firstSegmentEnd + 5); // 5 bytes into second segment header
         chunk1.order(ByteOrder.LITTLE_ENDIAN);