Repository: GoogleCloudPlatform/pgadapter
Pull request #241 (2022-06-29 14:47:46)

fix: skip bytes in an invalid stream defensively
PGAdapter should defensively skip bytes if it receives a message that it
does not recognize (i.e. the stream has become corrupted). Otherwise,
the stream might block if it tries to skip more bytes than are available
in the stream, which again could cause the connection to be blocked.

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/241?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> :exclamation: No coverage uploaded for pull request base (`postgresql-dialect@d39cec8`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#section-missing-base-commit).
> The diff coverage is `n/a`.

```diff
@@                  Coverage Diff                  @@
##             postgresql-dialect     #241   +/-   ##
=====================================================
  Coverage                      ?   80.99%           
  Complexity                    ?     1111           
=====================================================
  Files                         ?       94           
  Lines                         ?     3726           
  Branches                      ?      445           
=====================================================
  Hits                          ?     3018           
  Misses                        ?      546           
  Partials                      ?      162           
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `80.99% <0.00%> (?)` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#carryforward-flags-in-the-pull-request-comment) to find out more.


------

[Continue to review full report at Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/241?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/241?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform). Last update [d39cec8...39f3517](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/241?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform).


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
index 99251230..cd06719f 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/ControlMessage.java
@@ -93,11 +93,11 @@ public abstract class ControlMessage extends WireMessage {
             // Skip sync/flush in COPY_IN. This is consistent with real PostgreSQL which also does
             // this to accommodate clients that do not check what type of statement they sent in an
             // ExecuteMessage, and instead always blindly send a flush/sync after each execute.
-            return new SkipMessage(connection);
+            return SkipMessage.createForValidStream(connection);
           default:
             // Skip other unexpected messages and throw an exception to fail the copy operation.
             validMessage = false;
-            new SkipMessage(connection);
+            SkipMessage.createForInvalidStream(connection);
             throw new IllegalStateException(
                 String.format(
                     "Expected CopyData ('d'), CopyDone ('c') or CopyFail ('f') messages, got: '%c'",
@@ -133,7 +133,8 @@ public abstract class ControlMessage extends WireMessage {
             // terminate the connection to prevent the server from being flooded with invalid
             // messages.
             validMessage = false;
-            return new SkipMessage(connection);
+            // Note: The stream itself is still valid as we received a message that we recognized.
+            return SkipMessage.createForValidStream(connection);
           default:
             throw new IllegalStateException(String.format("Unknown message: %c", nextMsg));
         }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SkipMessage.java b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SkipMessage.java
index f4b3d8ae..7449bccc 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SkipMessage.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/wireprotocol/SkipMessage.java
@@ -16,6 +16,7 @@ package com.google.cloud.spanner.pgadapter.wireprotocol;
 
 import com.google.api.core.InternalApi;
 import com.google.cloud.spanner.pgadapter.ConnectionHandler;
+import java.io.DataInputStream;
 import java.io.IOException;
 import java.text.MessageFormat;
 
@@ -26,13 +27,40 @@ import java.text.MessageFormat;
  */
 @InternalApi
 public class SkipMessage extends ControlMessage {
-  public SkipMessage(ConnectionHandler connectionHandler) throws IOException {
+
+  public static SkipMessage createForValidStream(ConnectionHandler connectionHandler)
+      throws IOException {
+    return new SkipMessage(connectionHandler, false);
+  }
+
+  public static SkipMessage createForInvalidStream(ConnectionHandler connectionHandler)
+      throws IOException {
+    return new SkipMessage(connectionHandler, true);
+  }
+
+  private SkipMessage(ConnectionHandler connectionHandler, boolean streamIsPotentiallyInvalid)
+      throws IOException {
     super(connectionHandler);
     int skipLength = this.length - 4;
     int skipped = 0;
     // Read and skip bytes until we have reached the total message length.
-    while (skipped < skipLength) {
-      skipped += connection.getConnectionMetadata().getInputStream().skip(skipLength - skipped);
+    DataInputStream input = connection.getConnectionMetadata().getInputStream();
+    // If the stream is potentially invalid we'll use a more pessimistic skipping algorithm than if
+    // we just received a message that we want to skip. We do this because the Unix domain socket
+    // implementation might block on skip(n) if the stream has no bytes more to read.
+    if (streamIsPotentiallyInvalid) {
+      skipInvalid(input, skipLength);
+    } else {
+      while (skipped < skipLength) {
+        skipped += input.skip(skipLength - skipped);
+      }
+    }
+  }
+
+  private void skipInvalid(DataInputStream input, int numBytes) throws IOException {
+    int skipped = 0;
+    while (skipped < numBytes && input.read() != -1) {
+      skipped++;
     }
   }
 


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
index 8ef439ba..842e7b85 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/wireprotocol/ProtocolTest.java
@@ -1711,7 +1711,7 @@ public class ProtocolTest {
     when(connectionMetadata.getOutputStream()).thenReturn(outputStream);
     when(connectionHandler.getConnectionMetadata()).thenReturn(connectionMetadata);
 
-    SkipMessage message = new SkipMessage(connectionHandler);
+    SkipMessage message = SkipMessage.createForValidStream(connectionHandler);
     message.send();
 
     // Verify that nothing was written to the output.