aws-lambda-runtime-interface-emulator/lambda/rapi/middleware/middleware_test.go at dd9ee8288145aece13dfe1f604037ccfbdddbfd6 · mbfreder/aws-lambda-runtime-interface-emulator · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package middleware

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/go-chi/chi/v5"
	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"go.amzn.com/lambda/appctx"
	"go.amzn.com/lambda/extensions"
	"go.amzn.com/lambda/rapi/handler"
	"go.amzn.com/lambda/rapi/model"
)

type mockHandler struct{}

func (h *mockHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {}

func TestRuntimeReleaseMiddleware(t *testing.T) {
	appCtx := appctx.NewApplicationContext()
	router := chi.NewRouter()
	handler := &mockHandler{}
	router.Use(RuntimeReleaseMiddleware())
	router.Get("/", handler.ServeHTTP)

	userAgent := "foobar"

	responseRecorder := httptest.NewRecorder()
	responseBody := make([]byte, 100)
	request := httptest.NewRequest("GET", "/", bytes.NewReader(responseBody))
	request.Header.Set("User-Agent", userAgent)
	router.ServeHTTP(responseRecorder, appctx.RequestWithAppCtx(request, appCtx))

	assert.Equal(t, http.StatusOK, responseRecorder.Code)
	ctxRuntimeRelease, ok := appCtx.Load(appctx.AppCtxRuntimeReleaseKey)
	assert.True(t, ok)
	assert.Equal(t, userAgent, ctxRuntimeRelease)
}

func TestAgentUniqueIdentifierHeaderValidatorForbidden(t *testing.T) {
	router := chi.NewRouter()
	mockHandler := &mockHandler{}
	router.Get("/", AgentUniqueIdentifierHeaderValidator(mockHandler).ServeHTTP)
	responseBody := make([]byte, 100)
	var errorResponse model.ErrorResponse

	request := httptest.NewRequest("GET", "/", bytes.NewReader(responseBody))

	responseRecorder := httptest.NewRecorder()
	router.ServeHTTP(responseRecorder, request)
	assert.Equal(t, http.StatusForbidden, responseRecorder.Code)
	respBody, _ := io.ReadAll(responseRecorder.Body)
	json.Unmarshal(respBody, &errorResponse)
	assert.Equal(t, handler.ErrAgentIdentifierMissing, errorResponse.ErrorType)

	responseRecorder = httptest.NewRecorder()
	request.Header.Set(handler.LambdaAgentIdentifier, "invalid-unique-identifier")
	router.ServeHTTP(responseRecorder, request)
	assert.Equal(t, http.StatusForbidden, responseRecorder.Code)
	respBody, _ = io.ReadAll(responseRecorder.Body)
	json.Unmarshal(respBody, &errorResponse)
	assert.Equal(t, handler.ErrAgentIdentifierInvalid, errorResponse.ErrorType)
}

func TestAgentUniqueIdentifierHeaderValidatorSuccess(t *testing.T) {
	router := chi.NewRouter()
	mockHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		val, ok := r.Context().Value(handler.AgentIDCtxKey).(uuid.UUID)
		if !ok {
			assert.FailNow(t, "expected key not in request context")
		}
		assert.Equal(t, "85083764-ff1e-476f-ada1-d51f26e4f6be", val.String())
	})
	router.Get("/", AgentUniqueIdentifierHeaderValidator(mockHandler).ServeHTTP)
	responseBody := make([]byte, 100)
	request := httptest.NewRequest("GET", "/", bytes.NewReader(responseBody))
	ctx := context.Background()
	request = request.WithContext(ctx)

	responseRecorder := httptest.NewRecorder()
	responseRecorder.Code = http.StatusOK
	request.Header.Set(handler.LambdaAgentIdentifier, "85083764-ff1e-476f-ada1-d51f26e4f6be")
	router.ServeHTTP(responseRecorder, request)
	assert.Equal(t, http.StatusOK, responseRecorder.Code)
}

func TestAllowIfExtensionsEnabledPositive(t *testing.T) {
	router := chi.NewRouter()
	handler := &mockHandler{}
	router.Use(AllowIfExtensionsEnabled)
	router.Get("/", handler.ServeHTTP)

	responseRecorder := httptest.NewRecorder()
	responseBody := make([]byte, 100)

	extensions.Enable()
	defer extensions.Disable()

	router.ServeHTTP(responseRecorder, httptest.NewRequest("GET", "/", bytes.NewReader(responseBody)))

	assert.Equal(t, http.StatusOK, responseRecorder.Code)
}

func TestAllowIfExtensionsEnabledNegative(t *testing.T) {
	router := chi.NewRouter()
	handler := &mockHandler{}
	router.Use(AllowIfExtensionsEnabled)
	router.Get("/", handler.ServeHTTP)

	responseRecorder := httptest.NewRecorder()
	responseBody := make([]byte, 100)
	router.ServeHTTP(responseRecorder, httptest.NewRequest("GET", "/", bytes.NewReader(responseBody)))

	assert.Equal(t, http.StatusNotFound, responseRecorder.Code)
}