Allow specifying custom headers for Loki log output (#3227)

log/loki.go

@@ -19,6 +19,7 @@ import (
 type lokiHook struct {
 	fallbackLogger logrus.FieldLogger
 	addr           string
+	headers        [][2]string
 	labels         [][2]string
 	ch             chan *logrus.Entry
 	limit          int
@@ -124,6 +125,11 @@ func (h *lokiHook) parseArgs(line string) error {
 				labelKey := strings.TrimPrefix(key, "label.")
 				h.labels = append(h.labels, [2]string{labelKey, value})
 
+				continue
+			} else if strings.HasPrefix(key, "header.") {
+				headerKey := strings.TrimPrefix(key, "header.")
+				h.headers = append(h.headers, [2]string{headerKey, value})
+
 				continue
 			}
 
@@ -336,6 +342,10 @@ func (h *lokiHook) push(b bytes.Buffer) error {
 
 	req.Header.Set("Content-Type", "application/json")
 
+	for _, header := range h.headers {
+		req.Header.Add(header[0], header[1])
+	}
+
 	res, err := h.client.Do(req)
 
 	if res != nil {

log/loki_test.go

@@ -36,9 +36,10 @@ func TestSyslogFromConfigLine(t *testing.T) {
 			},
 		},
 		{
-			line: "loki=somewhere:1233,label.something=else,label.foo=bar,limit=32,level=info,allowedLabels=[something],pushPeriod=5m32s,msgMaxSize=1231",
+			line: "loki=somewhere:1233,label.something=else,label.foo=bar,limit=32,level=info,allowedLabels=[something],pushPeriod=5m32s,msgMaxSize=1231,header.x-test=123,header.authorization=token foobar",
 			res: lokiHook{
 				addr:          "somewhere:1233",
+				headers:       [][2]string{{"x-test", "123"}, {"authorization", "token foobar"}},
 				limit:         32,
 				pushPeriod:    time.Minute*5 + time.Second*32,
 				levels:        logrus.AllLevels[:5],
@@ -190,3 +191,38 @@ func TestLokiFlushingOnStop(t *testing.T) {
 		t.Fatal("No logs were received from loki before hook has finished")
 	}
 }
+
+func TestLokiHeaders(t *testing.T) {
+	t.Parallel()
+	receivedHeaders := make(chan http.Header, 1)
+	srv := httptest.NewServer(
+		http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+			receivedHeaders <- req.Header
+			close(receivedHeaders) // see comment in TestLokiFlushingOnStop
+		}),
+	)
+	configLine := fmt.Sprintf("loki=%s,pushPeriod=1h,header.X-Foo=bar,header.Test=hello world,header.X-Foo=baz", srv.URL)
+	h, err := LokiFromConfigLine(nil, configLine)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	wg := new(sync.WaitGroup)
+	now := time.Now()
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		err = h.Fire(&logrus.Entry{Time: now, Level: logrus.InfoLevel, Message: "test message"})
+		time.Sleep(time.Millisecond * 10)
+		cancel()
+	}()
+	h.Listen(ctx)
+	wg.Wait()
+
+	select {
+	case headers := <-receivedHeaders:
+		require.Equal(t, []string{"bar", "baz"}, headers["X-Foo"])
+		require.Equal(t, []string{"hello world"}, headers["Test"])
+	default:
+		t.Fatal("No logs were received from loki before hook has finished")
+	}
+}