package handlers import ( "bytes" "fmt" "testing" "time" "github.com/mimecast/dtail/internal" "github.com/mimecast/dtail/internal/io/dlog" "github.com/mimecast/dtail/internal/io/line" "github.com/mimecast/dtail/internal/protocol" ) // newReadTestHandler returns a baseHandler suitable for exercising Read // directly, without any output or generation scoping involved. func newReadTestHandler() baseHandler { return baseHandler{ done: internal.NewDone(), lines: make(chan *line.Line, 4), serverMessages: make(chan string, 4), maprMessages: make(chan string, 4), hostname: "testhost", } } // readExactly drains wantLen bytes from the handler using a buffer of bufSize // bytes per Read call, failing the test when Read errors or stalls. Stalls are // detected fast via consecutive empty reads and a wall-clock deadline: an // empty Read means the handler blocked ~1s in its poll select without data, // so two in a row indicate the remaining message bytes were dropped (e.g. a // regression back to resetting readBuf) and we fail with a diagnostic instead // of running into the go test timeout. func readExactly(t *testing.T, handler *baseHandler, bufSize, wantLen int) []byte { t.Helper() var got []byte p := make([]byte, bufSize) deadline := time.Now().Add(10 * time.Second) emptyReads := 0 for len(got) < wantLen { if time.Now().After(deadline) { t.Fatalf("Read stalled: got %d of %d bytes before deadline", len(got), wantLen) } n, err := handler.Read(p) if err != nil { t.Fatalf("Read() error = %v", err) } if n == 0 { emptyReads++ if emptyReads >= 2 { t.Fatalf("Read stalled: got %d of %d bytes, %d consecutive empty reads", len(got), wantLen, emptyReads) } continue } emptyReads = 0 got = append(got, p[:n]...) } return got } // expectedRemoteLine renders the protocol message Read is expected to emit // for a line delivered via the lines channel in non-plain mode. func expectedRemoteLine(content []byte, count uint64, sourceID string) []byte { var want bytes.Buffer formatRemoteLine(&want, "testhost", fmt.Sprintf("%3d", 100), count, sourceID, content) return want.Bytes() } // TestBaseHandlerReadLargeLineAcrossMultipleReads reproduces the original // bug: a line larger than the caller's buffer must arrive completely, // including the trailing message delimiter, across multiple Read calls. func TestBaseHandlerReadLargeLineAcrossMultipleReads(t *testing.T) { handler := newReadTestHandler() content := bytes.Repeat([]byte("x"), 1000) handler.lines <- line.New(bytes.NewBuffer(append([]byte{}, content...)), 1, 100, "test.log") want := expectedRemoteLine(content, 1, "test.log") got := readExactly(t, &handler, 32, len(want)) if !bytes.Equal(got, want) { t.Fatalf("large line corrupted across reads:\ngot %q\nwant %q", got, want) } if got[len(got)-1] != protocol.MessageDelimiter { t.Fatalf("message delimiter lost, last byte = %q", got[len(got)-1]) } } // TestBaseHandlerReadLargeServerMessageAcrossMultipleReads verifies the // serverMessages path keeps its remainder across Read calls as well. func TestBaseHandlerReadLargeServerMessageAcrossMultipleReads(t *testing.T) { handler := newReadTestHandler() message := "server says: " + string(bytes.Repeat([]byte("y"), 500)) handler.serverMessages <- message var want bytes.Buffer formatServerMessage(&want, "testhost", message, false) got := readExactly(t, &handler, 32, want.Len()) if !bytes.Equal(got, want.Bytes()) { t.Fatalf("large server message corrupted across reads:\ngot %q\nwant %q", got, want.Bytes()) } } // TestBaseHandlerReadLargeMaprMessageAcrossMultipleReads verifies the // maprMessages AGGREGATE path keeps its remainder across Read calls. func TestBaseHandlerReadLargeMaprMessageAcrossMultipleReads(t *testing.T) { handler := newReadTestHandler() message := "aggregated " + string(bytes.Repeat([]byte("m"), 500)) handler.maprMessages <- message var want bytes.Buffer want.WriteString("AGGREGATE") want.WriteString(protocol.FieldDelimiter) want.WriteString("testhost") want.WriteString(protocol.FieldDelimiter) want.WriteString(message) want.WriteByte(protocol.MessageDelimiter) got := readExactly(t, &handler, 32, want.Len()) if !bytes.Equal(got, want.Bytes()) { t.Fatalf("large mapr message corrupted across reads:\ngot %q\nwant %q", got, want.Bytes()) } } // TestBaseHandlerReadLargeHiddenMessageAcrossMultipleReads verifies the // hidden-message path (messages starting with '.') keeps its remainder across // Read calls: hidden messages are forwarded verbatim plus the delimiter. func TestBaseHandlerReadLargeHiddenMessageAcrossMultipleReads(t *testing.T) { handler := newReadTestHandler() message := ".hidden " + string(bytes.Repeat([]byte("h"), 500)) handler.serverMessages <- message want := append([]byte(message), protocol.MessageDelimiter) got := readExactly(t, &handler, 32, len(want)) if !bytes.Equal(got, want) { t.Fatalf("large hidden message corrupted across reads:\ngot %q\nwant %q", got, want) } } // TestBaseHandlerReadDrainsRemainderBeforeOutputData verifies that a pending // readBuf remainder is fully delivered before any output payload when output // mode is toggled on between Reads: the remainder belongs to a message that // was already accepted for delivery, so output output must not preempt it. func TestBaseHandlerReadDrainsRemainderBeforeOutputData(t *testing.T) { // The output read path logs via dlog.Server, which is nil in unit tests; // stub it out like the other output tests in this package do. originalLogger := dlog.Server dlog.Server = &dlog.DLog{} t.Cleanup(func() { dlog.Server = originalLogger }) handler := newReadTestHandler() message := "regular " + string(bytes.Repeat([]byte("r"), 200)) handler.serverMessages <- message var want bytes.Buffer formatServerMessage(&want, "testhost", message, false) // First small Read leaves the rest of the message in readBuf. p := make([]byte, 16) n, err := handler.Read(p) if err != nil { t.Fatalf("Read() error = %v", err) } got := append([]byte{}, p[:n]...) // Toggle output mode on mid-message and queue a output payload. if !handler.output.enable() { t.Fatal("expected output enable to switch output mode on") } outputPayload := []byte("output payload data") handler.output.channel() <- outputPayload total := want.Len() + len(outputPayload) got = append(got, readExactly(t, &handler, 16, total-len(got))...) if !bytes.Equal(got[:want.Len()], want.Bytes()) { t.Fatalf("output data preempted pending remainder:\ngot %q\nwant %q", got[:want.Len()], want.Bytes()) } if !bytes.Equal(got[want.Len():], outputPayload) { t.Fatalf("output payload corrupted:\ngot %q\nwant %q", got[want.Len():], outputPayload) } } // TestBaseHandlerReadDelimiterAloneInFinalRead verifies the boundary where // the caller's buffer holds everything except the trailing delimiter, so the // final Read must deliver the delimiter as its sole byte. func TestBaseHandlerReadDelimiterAloneInFinalRead(t *testing.T) { handler := newReadTestHandler() message := "delimiter boundary" var want bytes.Buffer formatServerMessage(&want, "testhost", message, false) handler.serverMessages <- message p := make([]byte, want.Len()-1) n, err := handler.Read(p) if err != nil { t.Fatalf("Read() error = %v", err) } if n != want.Len()-1 || !bytes.Equal(p[:n], want.Bytes()[:n]) { t.Fatalf("first read mismatch: n = %d, want %d", n, want.Len()-1) } n, err = handler.Read(p) if err != nil { t.Fatalf("Read() error = %v", err) } if n != 1 || p[0] != protocol.MessageDelimiter { t.Fatalf("expected lone delimiter in final read, got %q", p[:n]) } } // TestBaseHandlerReadExactFitBuffer verifies that a message exactly filling // the caller's buffer leaves no stale remainder behind: the next message must // start fresh instead of being glued to leftover bytes. func TestBaseHandlerReadExactFitBuffer(t *testing.T) { handler := newReadTestHandler() first := "exact fit" var firstWant bytes.Buffer formatServerMessage(&firstWant, "testhost", first, false) handler.serverMessages <- first p := make([]byte, firstWant.Len()) n, err := handler.Read(p) if err != nil { t.Fatalf("Read() error = %v", err) } if n != firstWant.Len() || !bytes.Equal(p[:n], firstWant.Bytes()) { t.Fatalf("exact-fit read mismatch:\ngot %q\nwant %q", p[:n], firstWant.Bytes()) } second := "next message" var secondWant bytes.Buffer formatServerMessage(&secondWant, "testhost", second, false) handler.serverMessages <- second got := readExactly(t, &handler, firstWant.Len(), secondWant.Len()) if !bytes.Equal(got, secondWant.Bytes()) { t.Fatalf("stale remainder leaked into next message:\ngot %q\nwant %q", got, secondWant.Bytes()) } } // TestBaseHandlerReadMultipleQueuedMessages verifies that several queued // messages, each larger than the read buffer, arrive back to back in order // and without corruption. func TestBaseHandlerReadMultipleQueuedMessages(t *testing.T) { handler := newReadTestHandler() var want bytes.Buffer for i := 0; i < 3; i++ { content := bytes.Repeat([]byte{byte('a' + i)}, 100) handler.lines <- line.New(bytes.NewBuffer(append([]byte{}, content...)), uint64(i+1), 100, "queued.log") want.Write(expectedRemoteLine(content, uint64(i+1), "queued.log")) } got := readExactly(t, &handler, 16, want.Len()) if !bytes.Equal(got, want.Bytes()) { t.Fatalf("queued messages corrupted:\ngot %q\nwant %q", got, want.Bytes()) } } // TestBaseHandlerReadPlainEmptyLine verifies the edge case of an empty line // in plain mode: the message consists of the delimiter only. func TestBaseHandlerReadPlainEmptyLine(t *testing.T) { handler := newReadTestHandler() handler.plain = true handler.lines <- line.New(&bytes.Buffer{}, 1, 100, "empty.log") p := make([]byte, 8) n, err := handler.Read(p) if err != nil { t.Fatalf("Read() error = %v", err) } if n != 1 || p[0] != protocol.MessageDelimiter { t.Fatalf("expected single delimiter byte, got %q", p[:n]) } }