// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

// SPDX-License-Identifier: Apache-2.0

package rapi

import (

"net/http"

"go.amzn.com/lambda/appctx"

"go.amzn.com/lambda/rapi/handler"

"go.amzn.com/lambda/rapi/middleware"

"go.amzn.com/lambda/telemetry"

"github.com/go-chi/chi"

"go.amzn.com/lambda/core"

"go.amzn.com/lambda/rapi/rendering"

)

// NewRouter returns a new instance of chi router implementing

// Runtime API specification.

func NewRouter(appCtx appctx.ApplicationContext, registrationService core.RegistrationService, renderingService *rendering.EventRenderingService) http.Handler {

router := chi.NewRouter()

router.Use(middleware.AppCtxMiddleware(appCtx))

router.Use(middleware.AccessLogMiddleware())

router.Use(middleware.RuntimeReleaseMiddleware())

// To respect Hyrum's Law, keeping /ping API even though

// we no longer use it ourselves.

// http://www.hyrumslaw.com/

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

router.Get("/runtime/invocation/next",

handler.NewInvocationNextHandler(registrationService, renderingService).ServeHTTP)

// Note, request validation must happen before state

// transition. State machine transitions are irreversible

// at the moment.

router.Post("/runtime/invocation/{awsrequestid}/response",

middleware.AwsRequestIDValidator(

handler.NewInvocationResponseHandler(registrationService)).ServeHTTP)

router.Post("/runtime/invocation/{awsrequestid}/error",

middleware.AwsRequestIDValidator(

handler.NewInvocationErrorHandler(registrationService)).ServeHTTP)

router.Post("/runtime/init/error",

handler.NewInitErrorHandler(registrationService).ServeHTTP)

return router

}

// ExtensionsRouter returns a new instance of chi router implementing

// Extensions Runtime API specification.

func ExtensionsRouter(appCtx appctx.ApplicationContext, registrationService core.RegistrationService, renderingService *rendering.EventRenderingService) http.Handler {

router := chi.NewRouter()

router.Use(middleware.AccessLogMiddleware())

router.Use(middleware.AllowIfExtensionsEnabled)

router.Use(middleware.AppCtxMiddleware(appCtx))

registerHandler := handler.NewAgentRegisterHandler(registrationService)

router.Post("/extension/register",

registerHandler.ServeHTTP)

router.Get("/extension/event/next",

middleware.AgentUniqueIdentifierHeaderValidator(

handler.NewAgentNextHandler(registrationService, renderingService)).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

}

// LogsAPIRouter returns a new instance of chi router implementing

// Logs API specification.

func LogsAPIRouter(registrationService core.RegistrationService, logsSubscriptionAPI telemetry.LogsSubscriptionAPI) http.Handler {

router := chi.NewRouter()

router.Use(middleware.AccessLogMiddleware())

router.Use(middleware.AllowIfExtensionsEnabled)

router.Put("/logs",

middleware.AgentUniqueIdentifierHeaderValidator(

handler.NewRuntimeLogsHandler(registrationService, logsSubscriptionAPI)).ServeHTTP)

return router

}

// LogsAPIStubRouter returns a new instance of chi router implementing

// a stub of Logs API that always returns a non-committal response to

// prevent customer code from crashing when Logs API is disabled locally

func LogsAPIStubRouter() http.Handler {

router := chi.NewRouter()

router.Put("/logs", handler.NewRuntimeLogsStubHandler().ServeHTTP)

return router

}

func CredentialsAPIRouter(credentialsService core.CredentialsService) http.Handler {

router := chi.NewRouter()

router.Get("/credentials", handler.NewCredentialsHandler(credentialsService).ServeHTTP)

return router

}