aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/rapi/router.go at 1435a9b64dd03ec7122d36c85dbbd69ec307548e · aws/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
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package rapi

import (
	"net/http"

	"github.com/go-chi/chi/v5"

	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/appctx"
	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/core"
	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/rapi/handler"
	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/rapi/middleware"
	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/rapi/rendering"
	"github.com/aws/aws-lambda-runtime-interface-emulator/internal/lambda-managed-instances/telemetry"
)

type runtimeRequestHandler interface {
	handler.RuntimeNextHandler
	handler.RuntimeResponseHandler
	handler.RuntimeErrorHandler
}

func NewRouter(appCtx appctx.ApplicationContext, registrationService core.RegistrationService, renderingService *rendering.EventRenderingService, runtimeReqHandler runtimeRequestHandler) http.Handler {
	router := chi.NewRouter()
	router.Use(middleware.AppCtxMiddleware(appCtx))
	router.Use(middleware.AccessLogMiddleware())
	router.Use(middleware.RuntimeReleaseMiddleware())

	router.Get("/ping", http.MaxBytesHandler(handler.NewPingHandler(), 0).ServeHTTP)

	router.Get("/runtime/invocation/next",
		http.MaxBytesHandler(handler.NewInvocationNextHandler(registrationService, runtimeReqHandler), 0).ServeHTTP)

	router.Post("/runtime/invocation/{awsrequestid}/response",
		handler.NewInvocationResponseHandler(runtimeReqHandler).ServeHTTP)

	router.Post("/runtime/invocation/{awsrequestid}/error",
		http.MaxBytesHandler(handler.NewInvocationErrorHandler(runtimeReqHandler), requestBodyLimitBytes).ServeHTTP)

	router.Post("/runtime/init/error", http.MaxBytesHandler(handler.NewInitErrorHandler(registrationService), requestBodyLimitBytes).ServeHTTP)
	return router
}

func ExtensionsRouter(appCtx appctx.ApplicationContext, registrationService core.RegistrationService, renderingService *rendering.EventRenderingService) http.Handler {
	router := chi.NewRouter()
	router.Use(middleware.AccessLogMiddleware())
	router.Use(middleware.AppCtxMiddleware(appCtx))

	registerHandler := handler.NewAgentRegisterHandler(registrationService)
	router.Post("/extension/register",
		registerHandler.ServeHTTP)

	router.Get("/extension/event/next",
		middleware.AgentUniqueIdentifierHeaderValidator(
			http.MaxBytesHandler(handler.NewAgentNextHandler(registrationService, renderingService), 0)).ServeHTTP)

	router.Post("/extension/init/error",
		middleware.AgentUniqueIdentifierHeaderValidator(
			handler.NewAgentInitErrorHandler(registrationService)).ServeHTTP)

	router.Post("/extension/exit/error",
		middleware.AgentUniqueIdentifierHeaderValidator(
			handler.NewAgentExitErrorHandler(registrationService)).ServeHTTP)

	return router
}

func TelemetryAPIRouter(registrationService core.RegistrationService, telemetrySubscriptionAPI telemetry.SubscriptionAPI) http.Handler {
	router := chi.NewRouter()
	router.Use(middleware.AccessLogMiddleware())

	router.Put("/telemetry",
		middleware.AgentUniqueIdentifierHeaderValidator(
			handler.NewRuntimeTelemetrySubscriptionHandler(registrationService, telemetrySubscriptionAPI)).ServeHTTP)

	return router
}