Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(lambda-promtail): Add possibility to set additional HTTP headers for promtail-client requests #11883

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
Open
52 changes: 42 additions & 10 deletions tools/lambda-promtail/lambda-promtail/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"fmt"
"net/url"
"os"
"regexp"
"strconv"
"strings"

Expand All @@ -25,19 +26,22 @@ const (

maxErrMsgLen = 1024

invalidExtraLabelsError = "invalid value for environment variable EXTRA_LABELS. Expected a comma separated list with an even number of entries. "
invalidExtraLabelsError = "invalid value for environment variable EXTRA_LABELS. Expected a comma separated list with an even number of entries. "
invalidExtraHeadersError = "invalid value for environment variable EXTRA_HTTP_HEADERS. Expected a comma separated list with an even number of entries."
trc-ikeskin marked this conversation as resolved.
Show resolved Hide resolved
)

var (
writeAddress *url.URL
username, password, extraLabelsRaw, dropLabelsRaw, tenantID, bearerToken string
keepStream bool
batchSize int
s3Clients map[string]*s3.Client
extraLabels model.LabelSet
dropLabels []model.LabelName
skipTlsVerify bool
printLogLine bool
writeAddress *url.URL
username, password, extraLabelsRaw, dropLabelsRaw, tenantID, bearerToken, extraHeadersRaw string
keepStream bool
batchSize int
s3Clients map[string]*s3.Client
extraLabels model.LabelSet
dropLabels []model.LabelName
skipTlsVerify bool
printLogLine bool
extraHeaders map[string]string
httpHeaderKeyRegex = regexp.MustCompile("^[-A-Za-z0-9]+$")
)

func setupArguments() {
Expand Down Expand Up @@ -106,6 +110,12 @@ func setupArguments() {
printLogLine = false
}
s3Clients = make(map[string]*s3.Client)

extraHeadersRaw = os.Getenv("EXTRA_HTTP_HEADERS")
extraHeaders, err = parseExtraHeaders(extraHeadersRaw)
if err != nil {
panic(err)
}
}

func parseExtraLabels(extraLabelsRaw string, omitPrefix bool) (model.LabelSet, error) {
Expand Down Expand Up @@ -190,6 +200,28 @@ func checkEventType(ev map[string]interface{}) (interface{}, error) {
return nil, fmt.Errorf("unknown event type!")
}

func parseExtraHeaders(extraHeadersRaw string) (map[string]string, error) {
extractedHeaders := make(map[string]string)
extraHeadersSplit := strings.Split(extraHeadersRaw, ",")

if len(extraHeadersRaw) < 1 {
return extractedHeaders, nil
}

if len(extraHeadersSplit)%2 != 0 {
return nil, fmt.Errorf(invalidExtraHeadersError)
}

for i := 0; i < len(extraHeadersSplit); i += 2 {
if !httpHeaderKeyRegex.MatchString(extraHeadersSplit[i]) {
return nil, fmt.Errorf("HTTP header key is invalid: %s", extraHeadersSplit[i])
}
extractedHeaders[extraHeadersSplit[i]] = extraHeadersSplit[i+1]
}

return extractedHeaders, nil
}

func handler(ctx context.Context, ev map[string]interface{}) error {
lvl, ok := os.LookupEnv("LOG_LEVEL")
if !ok {
Expand Down
30 changes: 30 additions & 0 deletions tools/lambda-promtail/lambda-promtail/main_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -64,3 +64,33 @@ func TestLambdaPromtail_TestDropLabels(t *testing.T) {
require.NotContains(t, modifiedLabels, model.LabelName("A1"))
require.Contains(t, modifiedLabels, model.LabelName("B2"))
}

func TestLambdaPromtail_ExtraHeadersValid(t *testing.T) {
extraHeaders, err := parseExtraHeaders("X-Custom-Header,This!sATota\\yCu$t0mHe4der,My-Server_WantsThis,What_ever could go here?,Expected4Entry,yLKc+QSB5VF/Gp3VPN7oOxa98yxWMxeHOAo+CW6trow=")
require.Nil(t, err)
require.Len(t, extraHeaders, 3)
require.Equal(t, extraHeaders["X-Custom-Header"], "This!sATota\\yCu$t0mHe4der")
require.Equal(t, extraHeaders["My-Server_WantsThis"], "What_ever could go here?")
require.Equal(t, extraHeaders["Expected4Entry"], "yLKc+QSB5VF/Gp3VPN7oOxa98yxWMxeHOAo+CW6trow=")
}

func TestLambdaPromtail_ExtraHeadersInvalidHeaderKey(t *testing.T) {
extraHeaders, err := parseExtraHeaders("Th.s_Shou|d-Fa!l,a")
require.Nil(t, extraHeaders)
require.ErrorContains(t, err, "HTTP header key is invalid:")
extraHeaders, err = parseExtraHeaders("Also Not Valid ,b")
require.Nil(t, extraHeaders)
require.ErrorContains(t, err, "HTTP header key is invalid:")
}

func TestLambdaPromtail_ExtraHeadersMissingValue(t *testing.T) {
extraHeaders, err := parseExtraHeaders("A,a,B,b,C,c,D")
require.Nil(t, extraHeaders)
require.Errorf(t, err, invalidExtraHeadersError)
}

func TestLambdaPromtail_TestParseHeadersNoneProvided(t *testing.T) {
extraLabels, err := parseExtraHeaders("")
require.Len(t, extraLabels, 0)
require.Nil(t, err)
}
Comment on lines +68 to +96
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lets simplify a bit and make this a table driven test

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not really sure how a table driven test simplifies my test, but I am also not too experience with this kind of testing. Maybe you could give me a hint and what you had in mind here?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rather than declaring 4 new tests, one for each of these success or failure scenarios, we would have one test with an array of input/expected output, see here https://go.dev/wiki/TableDrivenTests

6 changes: 6 additions & 0 deletions tools/lambda-promtail/lambda-promtail/promtail.go
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,12 @@ func (c *promtailClient) send(ctx context.Context, buf []byte) (int, error) {
req.Header.Set("Authorization", "Bearer "+bearerToken)
}

if len(extraHeaders) > 0 {
for key, value := range extraHeaders {
req.Header.Set(key, value)
}
}

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm am wondering if, to avoid any potential unexpected issues, we should include some validation here. For example, not allowing overwriting of the headers we set earlier in this function, and logging if we overwrite an existing header value (such as someone not noticing that in their configured extra headers they're setting SOME-EXTRA-HEADER=a, ... , SOME-EXTRA-HEADER=b.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I implemented a simple check to see if a key has already been set and added a warning for the user to check its env vars: 6a8a19d

Not sure about the logging part since I could not really find a clear way the logging is implemented in the code. Maybe you could have a closer look at that.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@trc-ikeskin thanks! this is exactly what I had in mind. Logging in this context just meant fmt.Println or using the logger which you've done.

I still think we should guard against people overwriting the headers that are set earlier in this same function, such as Authorization or X-Scope-OrgID. We can do that here or in a follow up PR. Up to you.

resp, err := c.http.Do(req.WithContext(ctx))
if err != nil {
return -1, err
Expand Down
Loading