From 4a326dd9d77e6f5faeb2889fd4e1df171d6b2041 Mon Sep 17 00:00:00 2001
From: Fae Charlton <fae.charlton@elastic.co>
Date: Wed, 22 Jul 2020 09:30:43 -0400
Subject: [PATCH] [libbeat] Fix write error in ensureWriter.Write (#20112)

(cherry picked from commit 787144fa45229b88131a453162f9a5c622730c85)
---
 libbeat/statestore/backend/memlog/util.go     |  2 +-
 .../statestore/backend/memlog/util_test.go    | 81 +++++++++++++++++++
 2 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 libbeat/statestore/backend/memlog/util_test.go

diff --git a/libbeat/statestore/backend/memlog/util.go b/libbeat/statestore/backend/memlog/util.go
index 2027c87adca..e2c5d4e6f68 100644
--- a/libbeat/statestore/backend/memlog/util.go
+++ b/libbeat/statestore/backend/memlog/util.go
@@ -55,7 +55,7 @@ func (e *ensureWriter) Write(p []byte) (int, error) {
 	for len(p) > 0 {
 		n, err := e.w.Write(p)
 		N, p = N+n, p[n:]
-		if isRetryErr(err) {
+		if err != nil && !isRetryErr(err) {
 			return N, err
 		}
 	}
diff --git a/libbeat/statestore/backend/memlog/util_test.go b/libbeat/statestore/backend/memlog/util_test.go
new file mode 100644
index 00000000000..fca2a2bbaf6
--- /dev/null
+++ b/libbeat/statestore/backend/memlog/util_test.go
@@ -0,0 +1,81 @@
+// Licensed to Elasticsearch B.V. under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. Elasticsearch B.V. licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package memlog
+
+import (
+	"syscall"
+	"testing"
+)
+
+// A mock Writer implementation that always returns a configurable
+// error on the first write call, to test error handling in ensureWriter.
+type mockErrorWriter struct {
+	errorType     error
+	reportedError bool
+}
+
+func (mew *mockErrorWriter) Write(data []byte) (n int, err error) {
+	if !mew.reportedError {
+		mew.reportedError = true
+		return 0, mew.errorType
+	}
+	return len(data), nil
+}
+
+func TestEnsureWriter_RetriableError(t *testing.T) {
+	// EAGAIN is retriable, ensureWriter.Write should succeed.
+	errorWriter := &mockErrorWriter{errorType: syscall.EAGAIN}
+	bytes := []byte{1, 2, 3}
+	writer := &ensureWriter{errorWriter}
+	written, err := writer.Write(bytes)
+	if err != nil {
+		t.Fatalf("ensureWriter shouldn't propagate retriable errors")
+	}
+	if written != len(bytes) {
+		t.Fatalf("Expected %d bytes written, got %d", len(bytes), written)
+	}
+}
+
+func TestEnsureWriter_NonRetriableError(t *testing.T) {
+	// EINVAL is not retriable, ensureWriter.Write should return an error.
+	errorWriter := &mockErrorWriter{errorType: syscall.EINVAL}
+	bytes := []byte{1, 2, 3}
+	writer := &ensureWriter{errorWriter}
+	written, err := writer.Write(bytes)
+	if err != syscall.EINVAL {
+		t.Fatalf("ensureWriter should propagate nonretriable errors")
+	}
+	if written != 0 {
+		t.Fatalf("Expected 0 bytes written, got %d", written)
+	}
+}
+
+func TestEnsureWriter_NoError(t *testing.T) {
+	// This tests the case where the underlying writer returns with no error,
+	// but without writing the full buffer.
+	var bytes []byte = []byte{1, 2, 3}
+	errorWriter := &mockErrorWriter{errorType: nil}
+	writer := &ensureWriter{errorWriter}
+	written, err := writer.Write(bytes)
+	if err != nil {
+		t.Fatalf("ensureWriter should only error if the underlying writer does")
+	}
+	if written != len(bytes) {
+		t.Fatalf("Expected %d bytes written, got %d", len(bytes), written)
+	}
+}