From a0558397fd0fa8afb7cc906930629fdf2ad46b7c Mon Sep 17 00:00:00 2001
From: soven <nicko54lay@mail.ru>
Date: Thu, 26 Nov 2020 01:37:04 +0300
Subject: [PATCH] fix log stdlibadapter when prefixed (#1036)

* fix log stdlibadapter when prefixed

* add some more doc for Prefix stdlib adapter option

* fix typo

* Update log/stdlib.go

Co-authored-by: Peter Bourgon <peterbourgon@users.noreply.github.com>
---
 log/stdlib.go      | 41 ++++++++++++++++++++++++++++++++++---
 log/stdlib_test.go | 50 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 88 insertions(+), 3 deletions(-)

diff --git a/log/stdlib.go b/log/stdlib.go
index ff96b5dee..2d030c456 100644
--- a/log/stdlib.go
+++ b/log/stdlib.go
@@ -1,6 +1,7 @@
 package log
 
 import (
+	"bytes"
 	"io"
 	"log"
 	"regexp"
@@ -26,9 +27,11 @@ func (w StdlibWriter) Write(p []byte) (int, error) {
 // messages, and place them under relevant keys.
 type StdlibAdapter struct {
 	Logger
-	timestampKey string
-	fileKey      string
-	messageKey   string
+	timestampKey    string
+	fileKey         string
+	messageKey      string
+	prefix          string
+	joinPrefixToMsg bool
 }
 
 // StdlibAdapterOption sets a parameter for the StdlibAdapter.
@@ -49,6 +52,16 @@ func MessageKey(key string) StdlibAdapterOption {
 	return func(a *StdlibAdapter) { a.messageKey = key }
 }
 
+// Prefix configures the adapter to parse a prefix from stdlib log events. If
+// you provide a non-empty prefix to the stdlib logger, then your should provide
+// that same prefix to the adapter via this option.
+//
+// By default, the prefix isn't included in the msg key. Set joinPrefixToMsg to
+// true if you want to include the parsed prefix in the msg.
+func Prefix(prefix string, joinPrefixToMsg bool) StdlibAdapterOption {
+	return func(a *StdlibAdapter) { a.prefix = prefix; a.joinPrefixToMsg = joinPrefixToMsg }
+}
+
 // NewStdlibAdapter returns a new StdlibAdapter wrapper around the passed
 // logger. It's designed to be passed to log.SetOutput.
 func NewStdlibAdapter(logger Logger, options ...StdlibAdapterOption) io.Writer {
@@ -65,6 +78,8 @@ func NewStdlibAdapter(logger Logger, options ...StdlibAdapterOption) io.Writer {
 }
 
 func (a StdlibAdapter) Write(p []byte) (int, error) {
+	p = a.handlePrefix(p)
+
 	result := subexps(p)
 	keyvals := []interface{}{}
 	var timestamp string
@@ -84,6 +99,7 @@ func (a StdlibAdapter) Write(p []byte) (int, error) {
 		keyvals = append(keyvals, a.fileKey, file)
 	}
 	if msg, ok := result["msg"]; ok {
+		msg = a.handleMessagePrefix(msg)
 		keyvals = append(keyvals, a.messageKey, msg)
 	}
 	if err := a.Logger.Log(keyvals...); err != nil {
@@ -92,6 +108,25 @@ func (a StdlibAdapter) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
+func (a StdlibAdapter) handlePrefix(p []byte) []byte {
+	if a.prefix != "" {
+		p = bytes.TrimPrefix(p, []byte(a.prefix))
+	}
+	return p
+}
+
+func (a StdlibAdapter) handleMessagePrefix(msg string) string {
+	if a.prefix == "" {
+		return msg
+	}
+
+	msg = strings.TrimPrefix(msg, a.prefix)
+	if a.joinPrefixToMsg {
+		msg = a.prefix + msg
+	}
+	return msg
+}
+
 const (
 	logRegexpDate = `(?P<date>[0-9]{4}/[0-9]{2}/[0-9]{2})?[ ]?`
 	logRegexpTime = `(?P<time>[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+)?)?[ ]?`
diff --git a/log/stdlib_test.go b/log/stdlib_test.go
index 90ebc5fa4..ae9d82ce4 100644
--- a/log/stdlib_test.go
+++ b/log/stdlib_test.go
@@ -72,6 +72,56 @@ func TestStdLibAdapterExtraction(t *testing.T) {
 	}
 }
 
+func TestStdLibAdapterPrefixedExtraction(t *testing.T) {
+	buf := &bytes.Buffer{}
+	logger := NewLogfmtLogger(buf)
+	writer := NewStdlibAdapter(logger, Prefix("some prefix ", false))
+	for input, want := range map[string]string{
+		"some prefix hello":                                            "msg=hello\n",
+		"some prefix 2009/01/23: hello":                                "ts=2009/01/23 msg=hello\n",
+		"some prefix 2009/01/23 01:23:23: hello":                       "ts=\"2009/01/23 01:23:23\" msg=hello\n",
+		"some prefix 01:23:23: hello":                                  "ts=01:23:23 msg=hello\n",
+		"some prefix 2009/01/23 01:23:23.123123: hello":                "ts=\"2009/01/23 01:23:23.123123\" msg=hello\n",
+		"some prefix 2009/01/23 01:23:23.123123 /a/b/c/d.go:23: hello": "ts=\"2009/01/23 01:23:23.123123\" caller=/a/b/c/d.go:23 msg=hello\n",
+		"some prefix 01:23:23.123123 /a/b/c/d.go:23: hello":            "ts=01:23:23.123123 caller=/a/b/c/d.go:23 msg=hello\n",
+		"some prefix 2009/01/23 01:23:23 /a/b/c/d.go:23: hello":        "ts=\"2009/01/23 01:23:23\" caller=/a/b/c/d.go:23 msg=hello\n",
+		"some prefix 2009/01/23 /a/b/c/d.go:23: hello":                 "ts=2009/01/23 caller=/a/b/c/d.go:23 msg=hello\n",
+		"some prefix /a/b/c/d.go:23: hello":                            "caller=/a/b/c/d.go:23 msg=hello\n",
+		"/a/b/c/d.go:23: some prefix hello":                            "caller=/a/b/c/d.go:23 msg=hello\n",
+	} {
+		buf.Reset()
+		fmt.Fprint(writer, input)
+		if have := buf.String(); want != have {
+			t.Errorf("%q: want %#v, have %#v", input, want, have)
+		}
+	}
+}
+
+func TestStdLibAdapterPrefixedExtractionWithJoinToMessage(t *testing.T) {
+	buf := &bytes.Buffer{}
+	logger := NewLogfmtLogger(buf)
+	writer := NewStdlibAdapter(logger, Prefix("some prefix ", true))
+	for input, want := range map[string]string{
+		"some prefix hello":                                            "msg=\"some prefix hello\"\n",
+		"some prefix 2009/01/23: hello":                                "ts=2009/01/23 msg=\"some prefix hello\"\n",
+		"some prefix 2009/01/23 01:23:23: hello":                       "ts=\"2009/01/23 01:23:23\" msg=\"some prefix hello\"\n",
+		"some prefix 01:23:23: hello":                                  "ts=01:23:23 msg=\"some prefix hello\"\n",
+		"some prefix 2009/01/23 01:23:23.123123: hello":                "ts=\"2009/01/23 01:23:23.123123\" msg=\"some prefix hello\"\n",
+		"some prefix 2009/01/23 01:23:23.123123 /a/b/c/d.go:23: hello": "ts=\"2009/01/23 01:23:23.123123\" caller=/a/b/c/d.go:23 msg=\"some prefix hello\"\n",
+		"some prefix 01:23:23.123123 /a/b/c/d.go:23: hello":            "ts=01:23:23.123123 caller=/a/b/c/d.go:23 msg=\"some prefix hello\"\n",
+		"some prefix 2009/01/23 01:23:23 /a/b/c/d.go:23: hello":        "ts=\"2009/01/23 01:23:23\" caller=/a/b/c/d.go:23 msg=\"some prefix hello\"\n",
+		"some prefix 2009/01/23 /a/b/c/d.go:23: hello":                 "ts=2009/01/23 caller=/a/b/c/d.go:23 msg=\"some prefix hello\"\n",
+		"some prefix /a/b/c/d.go:23: hello":                            "caller=/a/b/c/d.go:23 msg=\"some prefix hello\"\n",
+		"/a/b/c/d.go:23: some prefix hello":                            "caller=/a/b/c/d.go:23 msg=\"some prefix hello\"\n",
+	} {
+		buf.Reset()
+		fmt.Fprint(writer, input)
+		if have := buf.String(); want != have {
+			t.Errorf("%q: want %#v, have %#v", input, want, have)
+		}
+	}
+}
+
 func TestStdlibAdapterSubexps(t *testing.T) {
 	for input, wantMap := range map[string]map[string]string{
 		"hello world": {