From cc8a3e57712c749f44420b62f3672c3e981080f5 Mon Sep 17 00:00:00 2001
From: Pavlo Strokov <pstrokov@gitlab.com>
Date: Thu, 29 Jun 2023 15:08:17 +0300
Subject: [PATCH] gitaly: Buffer command output on context done

This is an attempt to prevent a failure on reading from the
command output and context cancellation. When the context is done
the termination signal is sent to the command to stop it. The
handling of the signal by the command may delay, so it may
proceed its execution, produce an output and finish with 0 code.
That is why before discarding the full command output we now
store its 1024 bytes and discard anything else because of safety
reasons. This data can now be read from the command by another
goroutine.

Closes: https://gitlab.com/gitlab-org/gitaly/-/issues/5021
---
 internal/command/command.go | 35 +++++++++++++++++++++++++++--------
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/internal/command/command.go b/internal/command/command.go
index bbea588c8a..2be1bdafc2 100644
--- a/internal/command/command.go
+++ b/internal/command/command.go
@@ -1,6 +1,7 @@
 package command
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -10,6 +11,7 @@ import (
 	"path"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -131,7 +133,7 @@ const (
 // terminated and reaped automatically when the context.Context that
 // created it is canceled.
 type Command struct {
-	reader       io.Reader
+	reader       atomic.Pointer[io.ReadCloser]
 	writer       io.WriteCloser
 	stderrBuffer *stderrBuffer
 	cmd          *exec.Cmd
@@ -275,7 +277,7 @@ func New(ctx context.Context, nameAndArgs []string, opts ...Option) (*Command, e
 			return nil, fmt.Errorf("creating stdout pipe: %w", err)
 		}
 
-		command.reader = pipe
+		command.reader.Store(&pipe)
 	}
 
 	if cfg.stderr != nil {
@@ -339,13 +341,16 @@ func New(ctx context.Context, nameAndArgs []string, opts ...Option) (*Command, e
 	return command, nil
 }
 
-// Read calls Read() on the stdout pipe of the command.
+// Read calls Read() on the stdout pipe of the command or temporary buffer
+// that contains max 1024 bytes of the output in case the command execution
+// context is cancelled.
 func (c *Command) Read(p []byte) (int, error) {
-	if c.reader == nil {
+	readerPtr := c.reader.Load()
+	if readerPtr == nil || *readerPtr == nil {
 		panic("command has no reader")
 	}
 
-	return c.reader.Read(p)
+	return (*readerPtr).Read(p)
 }
 
 // Write calls Write() on the stdin pipe of the command.
@@ -379,16 +384,30 @@ func (c *Command) wait() {
 		c.writer.Close()
 	}
 
-	if c.reader != nil {
+	contextIsDone := c.context.Err() != nil
+
+	if readerPtr := c.reader.Load(); readerPtr != nil && *readerPtr != nil {
+		if contextIsDone {
+			// When context is done it is not yet mean the running command should fail.
+			// If the termination error is not yet processed by the command it may produce
+			// useful output that can be processed by the caller. That is why we consume
+			// max up to 1024 bytes from it and discard all the other info because of
+			// safety reasons. The consumed data can be read from the buffer by another
+			// goroutine. All remaining output of the command will be discarded.
+			buffer := &bytes.Buffer{}
+			_, _ = io.Copy(buffer, io.LimitReader(*readerPtr, 1024))
+			tmpBuffer := io.NopCloser(buffer)
+			c.reader.Store(&tmpBuffer)
+		}
 		// Prevent the command from blocking on writing to its stdout.
-		_, _ = io.Copy(io.Discard, c.reader)
+		_, _ = io.Copy(io.Discard, *readerPtr)
 	}
 
 	c.waitError = c.cmd.Wait()
 
 	// If the context is done, the process was likely terminated due to it. If so,
 	// we return the context error to correctly report the reason.
-	if c.context.Err() != nil {
+	if contextIsDone {
 		// The standard library sets exit status -1 if the process was terminated by a signal,
 		// such as the SIGTERM sent when context is done.
 		if exitCode, ok := ExitStatus(c.waitError); ok && exitCode == -1 {
-- 
GitLab