fix: strip headers from logs using log stream specification (#3226)

* modify log consumer test to show how stripping logs sometimes fails

bug introduced by #454

* fix: implement the docker log stream specification to strip headers from logs

* chore: lint

* chore: use require

---------

Co-authored-by: Manuel de la Peña <[email protected]>
Co-authored-by: Manuel de la Peña <[email protected]>

Author: LaurentGoderre

docker.go

@@ -5,6 +5,7 @@ import (
 	"bufio"
 	"context"
 	"encoding/base64"
+	"encoding/binary"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -367,8 +368,6 @@ func (c *DockerContainer) inspectRawContainer(ctx context.Context) (*container.I
 // Logs will fetch both STDOUT and STDERR from the current container. Returns a
 // ReadCloser and leaves it up to the caller to extract what it wants.
 func (c *DockerContainer) Logs(ctx context.Context) (io.ReadCloser, error) {
-	const streamHeaderSize = 8
-
 	options := container.LogsOptions{
 		ShowStdout: true,
 		ShowStderr: true,
@@ -380,42 +379,43 @@ func (c *DockerContainer) Logs(ctx context.Context) (io.ReadCloser, error) {
 	}
 	defer c.provider.Close()
 
+	resp, err := c.Inspect(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	if resp.Config.Tty {
+		return rc, nil
+	}
+
+	return c.parseMultiplexedLogs(rc), nil
+}
+
+// parseMultiplexedLogs handles the multiplexed log format used when TTY is disabled
+func (c *DockerContainer) parseMultiplexedLogs(rc io.ReadCloser) io.ReadCloser {
+	const streamHeaderSize = 8
+
 	pr, pw := io.Pipe()
 	r := bufio.NewReader(rc)
 
 	go func() {
-		lineStarted := true
-		for err == nil {
-			line, isPrefix, err := r.ReadLine()
-
-			if lineStarted && len(line) >= streamHeaderSize {
-				line = line[streamHeaderSize:] // trim stream header
-				lineStarted = false
-			}
-			if !isPrefix {
-				lineStarted = true
-			}
-
-			_, errW := pw.Write(line)
-			if errW != nil {
+		header := make([]byte, streamHeaderSize)
+		for {
+			_, errH := io.ReadFull(r, header)
+			if errH != nil {
+				_ = pw.CloseWithError(errH)
 				return
 			}
 
-			if !isPrefix {
-				_, errW := pw.Write([]byte("\n"))
-				if errW != nil {
-					return
-				}
-			}
-
-			if err != nil {
-				_ = pw.CloseWithError(err)
+			frameSize := binary.BigEndian.Uint32(header[4:])
+			if _, err := io.CopyN(pw, r, int64(frameSize)); err != nil {
+				pw.CloseWithError(err)
 				return
 			}
 		}
 	}()
 
-	return pr, nil
+	return pr
 }
 
 // Deprecated: use the ContainerRequest.LogConsumerConfig field instead.

from_dockerfile_test.go

@@ -159,7 +159,7 @@ func TestBuildImageFromDockerfile_Target(t *testing.T) {
 
 		logs, err := io.ReadAll(r)
 		require.NoError(t, err)
-		require.Equal(t, fmt.Sprintf("target%d\n\n", i), string(logs))
+		require.Equal(t, fmt.Sprintf("target%d\n", i), string(logs))
 	}
 }
 

logconsumer_test.go

@@ -363,7 +363,7 @@ func TestContainerLogsShouldBeWithoutStreamHeader(t *testing.T) {
 	ctx := context.Background()
 	req := ContainerRequest{
 		Image:      "alpine:latest",
-		Cmd:        []string{"sh", "-c", "id -u"},
+		Cmd:        []string{"sh", "-c", "echo 'abcdefghi' && echo 'foo'"},
 		WaitingFor: wait.ForExit(),
 	}
 	ctr, err := GenericContainer(ctx, GenericContainerRequest{
@@ -378,7 +378,32 @@ func TestContainerLogsShouldBeWithoutStreamHeader(t *testing.T) {
 	defer r.Close()
 	b, err := io.ReadAll(r)
 	require.NoError(t, err)
-	assert.Equal(t, "0", strings.TrimSpace(string(b)))
+	require.Equal(t, "abcdefghi\nfoo", strings.TrimSpace(string(b)))
+}
+
+func TestContainerLogsTty(t *testing.T) {
+	ctx := context.Background()
+	req := ContainerRequest{
+		Image: "alpine:latest",
+		Cmd:   []string{"sh", "-c", "echo 'abcdefghi' && echo 'foo'"},
+		ConfigModifier: func(ctr *container.Config) {
+			ctr.Tty = true
+		},
+		WaitingFor: wait.ForExit(),
+	}
+	ctr, err := GenericContainer(ctx, GenericContainerRequest{
+		ContainerRequest: req,
+		Started:          true,
+	})
+	CleanupContainer(t, ctr)
+	require.NoError(t, err)
+
+	r, err := ctr.Logs(ctx)
+	require.NoError(t, err)
+	defer r.Close()
+	b, err := io.ReadAll(r)
+	require.NoError(t, err)
+	require.Equal(t, "abcdefghi\r\nfoo", strings.TrimSpace(string(b)))
 }
 
 func TestContainerLogsEnableAtStart(t *testing.T) {