Fix possible CCE when PostgreSQL protocol encoding binary format int8 (…

…#23860) * Fix PostgreSQL protocol encode for binary format int8 * Refactor PostgreSQLInt8BinaryProtocolValueTest
apache · Jan 31, 2023 · ce309ad · ce309ad
1 parent 0e23858
commit ce309ad
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 15 deletions.
diff --git a/...gresql/packet/command/query/extended/bind/protocol/PostgreSQLInt8BinaryProtocolValue.java b/...gresql/packet/command/query/extended/bind/protocol/PostgreSQLInt8BinaryProtocolValue.java
@@ -19,8 +19,6 @@
 
 import org.apache.shardingsphere.db.protocol.postgresql.payload.PostgreSQLPacketPayload;
 
-import java.math.BigDecimal;
-
 /**
  * Binary protocol value for int8 for PostgreSQL.
  */
@@ -38,6 +36,6 @@ public Object read(final PostgreSQLPacketPayload payload, final int parameterVal
 
     @Override
     public void write(final PostgreSQLPacketPayload payload, final Object value) {
-        payload.writeInt8(value instanceof BigDecimal ? ((BigDecimal) value).longValue() : (Long) value);
+        payload.writeInt8(((Number) value).longValue());
     }
 }
diff --git a/...ql/packet/command/query/extended/bind/protocol/PostgreSQLInt8BinaryProtocolValueTest.java b/...ql/packet/command/query/extended/bind/protocol/PostgreSQLInt8BinaryProtocolValueTest.java
@@ -17,30 +17,51 @@
 
 package org.apache.shardingsphere.db.protocol.postgresql.packet.command.query.extended.bind.protocol;
 
+import io.netty.buffer.Unpooled;
 import org.apache.shardingsphere.db.protocol.postgresql.payload.PostgreSQLPacketPayload;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnitRunner;
 
+import java.math.BigDecimal;
+import java.nio.charset.StandardCharsets;
+
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
 @RunWith(MockitoJUnitRunner.class)
 public final class PostgreSQLInt8BinaryProtocolValueTest {
 
-    @Mock
-    private PostgreSQLPacketPayload payload;
+    @Test
+    public void assertGetColumnLength() {
+        assertThat(new PostgreSQLInt8BinaryProtocolValue().getColumnLength(1L), is(8));
+    }
+
+    @Test
+    public void assertRead() {
+        byte[] input = new byte[]{
+                (byte) 0x80, 0, 0, 0, 0, 0, 0, 0,
+                (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF,
+                0, 0, 0, 0, 0, 0, 0, 0,
+                (byte) 0x7F, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF};
+        PostgreSQLPacketPayload payload = new PostgreSQLPacketPayload(Unpooled.wrappedBuffer(input), StandardCharsets.UTF_8);
+        assertThat(new PostgreSQLInt8BinaryProtocolValue().read(payload, 8), is(Long.MIN_VALUE));
+        assertThat(new PostgreSQLInt8BinaryProtocolValue().read(payload, 8), is(-1L));
+        assertThat(new PostgreSQLInt8BinaryProtocolValue().read(payload, 8), is(0L));
+        assertThat(new PostgreSQLInt8BinaryProtocolValue().read(payload, 8), is(Long.MAX_VALUE));
+    }
 
     @Test
-    public void assertNewInstance() {
-        PostgreSQLInt8BinaryProtocolValue actual = new PostgreSQLInt8BinaryProtocolValue();
-        assertThat(actual.getColumnLength(null), is(8));
-        when(payload.readInt8()).thenReturn(1L);
-        assertThat(actual.read(payload, 8), is(1L));
-        actual.write(payload, 1L);
-        verify(payload).writeInt8(1L);
+    public void assertWrite() {
+        byte[] actual = new byte[24];
+        PostgreSQLPacketPayload payload = new PostgreSQLPacketPayload(Unpooled.wrappedBuffer(actual).writerIndex(0), StandardCharsets.UTF_8);
+        new PostgreSQLInt8BinaryProtocolValue().write(payload, -1);
+        new PostgreSQLInt8BinaryProtocolValue().write(payload, Long.MAX_VALUE);
+        new PostgreSQLInt8BinaryProtocolValue().write(payload, BigDecimal.valueOf(Long.MIN_VALUE));
+        byte[] expected = new byte[]{
+                (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF,
+                (byte) 0x7F, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF, (byte) 0xFF,
+                (byte) 0x80, 0, 0, 0, 0, 0, 0, 0};
+        assertThat(actual, is(expected));
     }
 }