From efab7f37b72e69e1112ed5f79384698df7030da8 Mon Sep 17 00:00:00 2001
From: Lyra Naeseth <lyra@medium.com>
Date: Tue, 13 Feb 2018 17:44:51 -0800
Subject: [PATCH] Have prefixFieldClashes respect the JSON FieldMap

Currently, it will fix clashes that don't exist and miss clashes that
do exist due to field remapping.

Fixes #708
---
 formatter.go           | 17 +++++++------
 json_formatter.go      |  2 +-
 json_formatter_test.go | 55 ++++++++++++++++++++++++++++++++++++++++++
 text_formatter.go      |  3 ++-
 4 files changed, 68 insertions(+), 9 deletions(-)

diff --git a/formatter.go b/formatter.go
index b183ff5..849dc8b 100644
--- a/formatter.go
+++ b/formatter.go
@@ -30,16 +30,19 @@ type Formatter interface {
 //
 // It's not exported because it's still using Data in an opinionated way. It's to
 // avoid code duplication between the two default formatters.
-func prefixFieldClashes(data Fields) {
-	if t, ok := data["time"]; ok {
-		data["fields.time"] = t
+func prefixFieldClashes(data Fields, fieldMap FieldMap) {
+	timeKey := fieldMap.resolve(FieldKeyTime)
+	if t, ok := data[timeKey]; ok {
+		data["fields."+timeKey] = t
 	}
 
-	if m, ok := data["msg"]; ok {
-		data["fields.msg"] = m
+	msgKey := fieldMap.resolve(FieldKeyMsg)
+	if m, ok := data[msgKey]; ok {
+		data["fields."+msgKey] = m
 	}
 
-	if l, ok := data["level"]; ok {
-		data["fields.level"] = l
+	levelKey := fieldMap.resolve(FieldKeyLevel)
+	if l, ok := data[levelKey]; ok {
+		data["fields."+levelKey] = l
 	}
 }
diff --git a/json_formatter.go b/json_formatter.go
index fb01c1b..7064947 100644
--- a/json_formatter.go
+++ b/json_formatter.go
@@ -58,7 +58,7 @@ func (f *JSONFormatter) Format(entry *Entry) ([]byte, error) {
 			data[k] = v
 		}
 	}
-	prefixFieldClashes(data)
+	prefixFieldClashes(data, f.FieldMap)
 
 	timestampFormat := f.TimestampFormat
 	if timestampFormat == "" {
diff --git a/json_formatter_test.go b/json_formatter_test.go
index 51093a7..1c140d0 100644
--- a/json_formatter_test.go
+++ b/json_formatter_test.go
@@ -3,6 +3,7 @@ package logrus
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
 	"strings"
 	"testing"
 )
@@ -106,6 +107,60 @@ func TestFieldClashWithLevel(t *testing.T) {
 	}
 }
 
+func TestFieldClashWithRemappedFields(t *testing.T) {
+	formatter := &JSONFormatter{
+		FieldMap: FieldMap{
+			FieldKeyTime:  "@timestamp",
+			FieldKeyLevel: "@level",
+			FieldKeyMsg:   "@message",
+		},
+	}
+
+	b, err := formatter.Format(WithFields(Fields{
+		"@timestamp": "@timestamp",
+		"@level":     "@level",
+		"@message":   "@message",
+		"timestamp":  "timestamp",
+		"level":      "level",
+		"msg":        "msg",
+	}))
+	if err != nil {
+		t.Fatal("Unable to format entry: ", err)
+	}
+
+	entry := make(map[string]interface{})
+	err = json.Unmarshal(b, &entry)
+	if err != nil {
+		t.Fatal("Unable to unmarshal formatted entry: ", err)
+	}
+
+	for _, field := range []string{"timestamp", "level", "msg"} {
+		if entry[field] != field {
+			t.Errorf("Expected field %v to be untouched; got %v", field, entry[field])
+		}
+
+		remappedKey := fmt.Sprintf("fields.%s", field)
+		if remapped, ok := entry[remappedKey]; ok {
+			t.Errorf("Expected %s to be empty; got %v", remappedKey, remapped)
+		}
+	}
+
+	for _, field := range []string{"@timestamp", "@level", "@message"} {
+		if entry[field] == field {
+			t.Errorf("Expected field %v to be mapped to an Entry value", field)
+		}
+
+		remappedKey := fmt.Sprintf("fields.%s", field)
+		if remapped, ok := entry[remappedKey]; ok {
+			if remapped != field {
+				t.Errorf("Expected field %v to be copied to %s; got %v", field, remappedKey, remapped)
+			}
+		} else {
+			t.Errorf("Expected field %v to be copied to %s; was absent", field, remappedKey)
+		}
+	}
+}
+
 func TestJSONEntryEndsWithNewline(t *testing.T) {
 	formatter := &JSONFormatter{}
 
diff --git a/text_formatter.go b/text_formatter.go
index 61b21ca..62763eb 100644
--- a/text_formatter.go
+++ b/text_formatter.go
@@ -20,6 +20,7 @@ const (
 
 var (
 	baseTimestamp time.Time
+	emptyFieldMap FieldMap
 )
 
 func init() {
@@ -82,7 +83,7 @@ func (f *TextFormatter) Format(entry *Entry) ([]byte, error) {
 		b = &bytes.Buffer{}
 	}
 
-	prefixFieldClashes(entry.Data)
+	prefixFieldClashes(entry.Data, emptyFieldMap)
 
 	f.Do(func() { f.init(entry) })