mirror of https://github.com/grafana/loki
fix(promtail): Handle docker logs when a log is split in multiple frames (#12374)
Jonas L. B
1 year ago
committed by
GitHub
parent
76ba24e3d8
commit
c0113db4e8
File diff suppressed because one or more lines are too long
Binary file not shown.
@ -0,0 +1,173 @@ |
||||
package framedstdcopy |
||||
|
||||
import ( |
||||
"bytes" |
||||
"encoding/binary" |
||||
"fmt" |
||||
"io" |
||||
|
||||
"github.com/docker/docker/pkg/stdcopy" |
||||
) |
||||
|
||||
const ( |
||||
// From stdcopy
|
||||
stdWriterPrefixLen = 8 |
||||
stdWriterFdIndex = 0 |
||||
stdWriterSizeIndex = 4 |
||||
startingBufLen = 32*1024 + stdWriterPrefixLen + 1 |
||||
maxFrameLen = 16384 + 31 // In practice (undocumented) frame payload can be timestamp + 16k
|
||||
) |
||||
|
||||
// FramedStdCopy is a modified version of stdcopy.StdCopy.
|
||||
// FramedStdCopy will demultiplex `src` in the same manner as StdCopy, but instead of
|
||||
// using io.Writer for outputs, channels are used, since each frame payload may contain
|
||||
// its own inner header (notably, timestamps). Frame payloads are not further parsed here,
|
||||
// but are passed raw as individual slices through the output channel.
|
||||
//
|
||||
// FramedStdCopy will read until it hits EOF on `src`. It will then return a nil error.
|
||||
// In other words: if `err` is non nil, it indicates a real underlying error.
|
||||
//
|
||||
// `written` will hold the total number of bytes written to `dstout` and `dsterr`.
|
||||
func FramedStdCopy(dstout, dsterr chan []byte, src io.Reader) (written int64, err error) { |
||||
var ( |
||||
buf = make([]byte, startingBufLen) |
||||
bufLen = len(buf) |
||||
nr int |
||||
er error |
||||
out chan []byte |
||||
frameSize int |
||||
) |
||||
|
||||
for { |
||||
// Make sure we have at least a full header
|
||||
for nr < stdWriterPrefixLen { |
||||
var nr2 int |
||||
nr2, er = src.Read(buf[nr:]) |
||||
nr += nr2 |
||||
if er == io.EOF { |
||||
if nr < stdWriterPrefixLen { |
||||
return written, nil |
||||
} |
||||
break |
||||
} |
||||
if er != nil { |
||||
return 0, er |
||||
} |
||||
} |
||||
|
||||
stream := stdcopy.StdType(buf[stdWriterFdIndex]) |
||||
// Check the first byte to know where to write
|
||||
switch stream { |
||||
case stdcopy.Stdin: |
||||
fallthrough |
||||
case stdcopy.Stdout: |
||||
// Write on stdout
|
||||
out = dstout |
||||
case stdcopy.Stderr: |
||||
// Write on stderr
|
||||
out = dsterr |
||||
case stdcopy.Systemerr: |
||||
// If we're on Systemerr, we won't write anywhere.
|
||||
// NB: if this code changes later, make sure you don't try to write
|
||||
// to outstream if Systemerr is the stream
|
||||
out = nil |
||||
default: |
||||
return 0, fmt.Errorf("Unrecognized input header: %d", buf[stdWriterFdIndex]) |
||||
} |
||||
|
||||
// Retrieve the size of the frame
|
||||
frameSize = int(binary.BigEndian.Uint32(buf[stdWriterSizeIndex : stdWriterSizeIndex+4])) |
||||
|
||||
// Check if the buffer is big enough to read the frame.
|
||||
// Extend it if necessary.
|
||||
if frameSize+stdWriterPrefixLen > bufLen { |
||||
buf = append(buf, make([]byte, frameSize+stdWriterPrefixLen-bufLen+1)...) |
||||
bufLen = len(buf) |
||||
} |
||||
|
||||
// While the amount of bytes read is less than the size of the frame + header, we keep reading
|
||||
for nr < frameSize+stdWriterPrefixLen { |
||||
var nr2 int |
||||
nr2, er = src.Read(buf[nr:]) |
||||
nr += nr2 |
||||
if er == io.EOF { |
||||
if nr < frameSize+stdWriterPrefixLen { |
||||
return written, nil |
||||
} |
||||
break |
||||
} |
||||
if er != nil { |
||||
return 0, er |
||||
} |
||||
} |
||||
|
||||
// we might have an error from the source mixed up in our multiplexed
|
||||
// stream. if we do, return it.
|
||||
if stream == stdcopy.Systemerr { |
||||
return written, fmt.Errorf("error from daemon in stream: %s", string(buf[stdWriterPrefixLen:frameSize+stdWriterPrefixLen])) |
||||
} |
||||
|
||||
// Write the retrieved frame (without header)
|
||||
var newBuf = make([]byte, frameSize) |
||||
copy(newBuf, buf[stdWriterPrefixLen:]) |
||||
out <- newBuf |
||||
written += int64(frameSize) |
||||
|
||||
// Move the rest of the buffer to the beginning
|
||||
copy(buf, buf[frameSize+stdWriterPrefixLen:nr]) |
||||
// Move the index
|
||||
nr -= frameSize + stdWriterPrefixLen |
||||
} |
||||
} |
||||
|
||||
// Specialized version of FramedStdCopy for when frames have no headers.
|
||||
// This will happen for output from a container that has TTY set.
|
||||
// In theory this makes it impossible to find the frame boundaries, which also does not matter if timestamps were not requested,
|
||||
// but if they were requested, they will still be there at the start of every frame, which might be mid-line.
|
||||
// In practice we can find most boundaries by looking for newlines, since these result in a new frame.
|
||||
// Otherwise we rely on using the same max frame size as used in practice by docker.
|
||||
func NoHeaderFramedStdCopy(dstout chan []byte, src io.Reader) (written int64, err error) { |
||||
var ( |
||||
buf = make([]byte, 32768) |
||||
nrLine int |
||||
nr int |
||||
nr2 int |
||||
er error |
||||
) |
||||
for { |
||||
nr2, er = src.Read(buf[nr:]) |
||||
if er == io.EOF && nr2 == 0 { |
||||
return written, nil |
||||
} else if er != nil { |
||||
return written, er |
||||
} |
||||
nr += nr2 |
||||
|
||||
// We might have read multiple frames, output all those we find in the buffer
|
||||
for nr > 0 { |
||||
nrLine = bytes.Index(buf[:nr], []byte("\n")) + 1 |
||||
if nrLine > maxFrameLen { |
||||
// we found a newline but it's in the next frame (most likely)
|
||||
nrLine = maxFrameLen |
||||
} else if nrLine < 1 { |
||||
if nr >= maxFrameLen { |
||||
nrLine = maxFrameLen |
||||
} else { |
||||
// no end of frame found and we don't have enough bytes
|
||||
break |
||||
} |
||||
} |
||||
|
||||
// Write the frame
|
||||
var newBuf = make([]byte, nrLine) |
||||
copy(newBuf, buf) |
||||
dstout <- newBuf |
||||
written += int64(nrLine) |
||||
|
||||
// Move the rest of the buffer to the beginning
|
||||
copy(buf, buf[nrLine:nr]) |
||||
// Move the index
|
||||
nr -= nrLine |
||||
} |
||||
} |
||||
} |
||||
@ -0,0 +1,269 @@ |
||||
package framedstdcopy |
||||
|
||||
import ( |
||||
"bytes" |
||||
"errors" |
||||
"io" |
||||
"strings" |
||||
"sync" |
||||
"testing" |
||||
|
||||
"github.com/docker/docker/pkg/stdcopy" |
||||
) |
||||
|
||||
const ( |
||||
tsPrefix string = "2024-03-14T15:32:05.358979323Z " |
||||
unprefixedFramePayloadSize int = 16384 |
||||
) |
||||
|
||||
func timestamped(bytes []byte) []byte { |
||||
var ts = []byte(tsPrefix) |
||||
return append(ts, bytes...) |
||||
} |
||||
|
||||
func getSrcBuffer(stdOutFrames, stdErrFrames [][]byte) (buffer *bytes.Buffer, err error) { |
||||
buffer = new(bytes.Buffer) |
||||
dstOut := stdcopy.NewStdWriter(buffer, stdcopy.Stdout) |
||||
for _, stdOutBytes := range stdOutFrames { |
||||
_, err = dstOut.Write(timestamped(stdOutBytes)) |
||||
if err != nil { |
||||
return |
||||
} |
||||
} |
||||
dstErr := stdcopy.NewStdWriter(buffer, stdcopy.Stderr) |
||||
for _, stdErrBytes := range stdErrFrames { |
||||
_, err = dstErr.Write(timestamped(stdErrBytes)) |
||||
if err != nil { |
||||
return |
||||
} |
||||
} |
||||
return |
||||
} |
||||
|
||||
type streamChans struct { |
||||
out chan []byte |
||||
err chan []byte |
||||
outCollected [][]byte |
||||
errCollected [][]byte |
||||
wg sync.WaitGroup |
||||
} |
||||
|
||||
func newChans() streamChans { |
||||
return streamChans{ |
||||
out: make(chan []byte), |
||||
err: make(chan []byte), |
||||
outCollected: make([][]byte, 0), |
||||
errCollected: make([][]byte, 0), |
||||
} |
||||
} |
||||
|
||||
func (crx *streamChans) collectFrames() { |
||||
crx.wg.Add(1) |
||||
outClosed := false |
||||
errClosed := false |
||||
for { |
||||
if outClosed && errClosed { |
||||
crx.wg.Done() |
||||
return |
||||
} |
||||
select { |
||||
case bytes, ok := <-crx.out: |
||||
outClosed = !ok |
||||
if bytes != nil { |
||||
crx.outCollected = append(crx.outCollected, bytes) |
||||
} |
||||
case bytes, ok := <-crx.err: |
||||
errClosed = !ok |
||||
if bytes != nil { |
||||
crx.errCollected = append(crx.errCollected, bytes) |
||||
} |
||||
} |
||||
} |
||||
} |
||||
func (crx *streamChans) close() { |
||||
close(crx.out) |
||||
close(crx.err) |
||||
} |
||||
|
||||
func TestStdCopyWriteAndRead(t *testing.T) { |
||||
ostr := strings.Repeat("o", unprefixedFramePayloadSize) |
||||
estr := strings.Repeat("e", unprefixedFramePayloadSize) |
||||
buffer, err := getSrcBuffer( |
||||
[][]byte{ |
||||
[]byte(ostr), |
||||
[]byte(ostr[:3] + "\n"), |
||||
}, |
||||
[][]byte{ |
||||
[]byte(estr), |
||||
[]byte(estr[:3] + "\n"), |
||||
}, |
||||
) |
||||
if err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
|
||||
rx := newChans() |
||||
go rx.collectFrames() |
||||
written, err := FramedStdCopy(rx.out, rx.err, buffer) |
||||
rx.close() |
||||
rx.wg.Wait() |
||||
if err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
tslen := len(tsPrefix) |
||||
expectedTotalWritten := 2*maxFrameLen + 2*(4+tslen) |
||||
if written != int64(expectedTotalWritten) { |
||||
t.Fatalf("Expected to have total of %d bytes written, got %d", expectedTotalWritten, written) |
||||
} |
||||
if !bytes.Equal(rx.outCollected[0][tslen:maxFrameLen], []byte(ostr)) { |
||||
t.Fatal("Expected the first out frame to be all 'o'") |
||||
} |
||||
if !bytes.Equal(rx.outCollected[1][tslen:tslen+4], []byte("ooo\n")) { |
||||
t.Fatal("Expected the second out frame to be 'ooo\\n'") |
||||
} |
||||
if !bytes.Equal(rx.errCollected[0][tslen:maxFrameLen], []byte(estr)) { |
||||
t.Fatal("Expected the first err frame to be all 'e'") |
||||
} |
||||
if !bytes.Equal(rx.errCollected[1][tslen:tslen+4], []byte("eee\n")) { |
||||
t.Fatal("Expected the second err frame to be 'eee\\n'") |
||||
} |
||||
} |
||||
|
||||
type customReader struct { |
||||
n int |
||||
err error |
||||
totalCalls int |
||||
correctCalls int |
||||
src *bytes.Buffer |
||||
} |
||||
|
||||
func (f *customReader) Read(buf []byte) (int, error) { |
||||
f.totalCalls++ |
||||
if f.totalCalls <= f.correctCalls { |
||||
return f.src.Read(buf) |
||||
} |
||||
return f.n, f.err |
||||
} |
||||
|
||||
func TestStdCopyReturnsErrorReadingHeader(t *testing.T) { |
||||
expectedError := errors.New("error") |
||||
reader := &customReader{ |
||||
err: expectedError, |
||||
} |
||||
discard := newChans() |
||||
go discard.collectFrames() |
||||
written, err := FramedStdCopy(discard.out, discard.err, reader) |
||||
discard.close() |
||||
if written != 0 { |
||||
t.Fatalf("Expected 0 bytes read, got %d", written) |
||||
} |
||||
if err != expectedError { |
||||
t.Fatalf("Didn't get expected error") |
||||
} |
||||
} |
||||
|
||||
func TestStdCopyReturnsErrorReadingFrame(t *testing.T) { |
||||
expectedError := errors.New("error") |
||||
stdOutBytes := []byte(strings.Repeat("o", unprefixedFramePayloadSize)) |
||||
stdErrBytes := []byte(strings.Repeat("e", unprefixedFramePayloadSize)) |
||||
buffer, err := getSrcBuffer([][]byte{stdOutBytes}, [][]byte{stdErrBytes}) |
||||
if err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
reader := &customReader{ |
||||
correctCalls: 1, |
||||
n: stdWriterPrefixLen + 1, |
||||
err: expectedError, |
||||
src: buffer, |
||||
} |
||||
discard := newChans() |
||||
go discard.collectFrames() |
||||
written, err := FramedStdCopy(discard.out, discard.err, reader) |
||||
discard.close() |
||||
if written != 0 { |
||||
t.Fatalf("Expected 0 bytes read, got %d", written) |
||||
} |
||||
if err != expectedError { |
||||
t.Fatalf("Didn't get expected error") |
||||
} |
||||
} |
||||
|
||||
func TestStdCopyDetectsCorruptedFrame(t *testing.T) { |
||||
stdOutBytes := []byte(strings.Repeat("o", unprefixedFramePayloadSize)) |
||||
stdErrBytes := []byte(strings.Repeat("e", unprefixedFramePayloadSize)) |
||||
buffer, err := getSrcBuffer([][]byte{stdOutBytes}, [][]byte{stdErrBytes}) |
||||
if err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
reader := &customReader{ |
||||
correctCalls: 1, |
||||
n: stdWriterPrefixLen + 1, |
||||
err: io.EOF, |
||||
src: buffer, |
||||
} |
||||
discard := newChans() |
||||
go discard.collectFrames() |
||||
written, err := FramedStdCopy(discard.out, discard.err, reader) |
||||
discard.close() |
||||
if written != maxFrameLen { |
||||
t.Fatalf("Expected %d bytes read, got %d", 0, written) |
||||
} |
||||
if err != nil { |
||||
t.Fatal("Didn't get nil error") |
||||
} |
||||
} |
||||
|
||||
func TestStdCopyWithInvalidInputHeader(t *testing.T) { |
||||
dst := newChans() |
||||
go dst.collectFrames() |
||||
src := strings.NewReader("Invalid input") |
||||
_, err := FramedStdCopy(dst.out, dst.err, src) |
||||
dst.close() |
||||
if err == nil { |
||||
t.Fatal("FramedStdCopy with invalid input header should fail.") |
||||
} |
||||
} |
||||
|
||||
func TestStdCopyWithCorruptedPrefix(t *testing.T) { |
||||
data := []byte{0x01, 0x02, 0x03} |
||||
src := bytes.NewReader(data) |
||||
written, err := FramedStdCopy(nil, nil, src) |
||||
if err != nil { |
||||
t.Fatalf("FramedStdCopy should not return an error with corrupted prefix.") |
||||
} |
||||
if written != 0 { |
||||
t.Fatalf("FramedStdCopy should have written 0, but has written %d", written) |
||||
} |
||||
} |
||||
|
||||
// TestStdCopyReturnsErrorFromSystem tests that FramedStdCopy correctly returns an
|
||||
// error, when that error is muxed into the Systemerr stream.
|
||||
func TestStdCopyReturnsErrorFromSystem(t *testing.T) { |
||||
// write in the basic messages, just so there's some fluff in there
|
||||
stdOutBytes := []byte(strings.Repeat("o", unprefixedFramePayloadSize)) |
||||
stdErrBytes := []byte(strings.Repeat("e", unprefixedFramePayloadSize)) |
||||
buffer, err := getSrcBuffer([][]byte{stdOutBytes}, [][]byte{stdErrBytes}) |
||||
if err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
// add in an error message on the Systemerr stream
|
||||
systemErrBytes := []byte(strings.Repeat("S", unprefixedFramePayloadSize)) |
||||
systemWriter := stdcopy.NewStdWriter(buffer, stdcopy.Systemerr) |
||||
_, err = systemWriter.Write(systemErrBytes) |
||||
if err != nil { |
||||
t.Fatal(err) |
||||
} |
||||
|
||||
// now copy and demux. we should expect an error containing the string we
|
||||
// wrote out
|
||||
discard := newChans() |
||||
go discard.collectFrames() |
||||
_, err = FramedStdCopy(discard.out, discard.err, buffer) |
||||
discard.close() |
||||
if err == nil { |
||||
t.Fatal("expected error, got none") |
||||
} |
||||
if !strings.Contains(err.Error(), string(systemErrBytes)) { |
||||
t.Fatal("expected error to contain message") |
||||
} |
||||
} |
||||
Loading…
Reference in new issue