Middy Middleware Integration

/**
 * Function used to cleanup Powertools for AWS resources when a Middy middleware
 * returns early and terminates the middleware chain.
 *
 * When a middleware returns early, all middleware lifecycle functions that come after it
 * are not executed. This means that if a middleware was relying on certain logic to be run
 * during the 'after' or 'onError' lifecycle functions, that logic will not be executed.
 *
 * This is the case for Powertools middlewares which rely on these lifecycle functions
 * to perform cleanup operations like closing the current segment in the tracer or
 * flushing any stored metrics.
 *
 * When authoring a middleware that might return early, use this function to cleanup
 * Powertools resources. This function will check if any cleanup function is present
 * in the request.internal object and execute it.
 *
 * @param request - The Middy request object
 */
function cleanupMiddlewares(request: MiddyLikeRequest): Promise<void>;

/**
 * Key to store the tracer instance in the request.internal object.
 * Value: 'powertools-for-aws.tracer'
 */
const TRACER_KEY: string;

/**
 * Key to store the metrics instance in the request.internal object.
 * Value: 'powertools-for-aws.metrics'
 */
const METRICS_KEY: string;

/**
 * Key to store the logger instance in the request.internal object.
 * Value: 'powertools-for-aws.logger'
 */
const LOGGER_KEY: string;

/**
 * Key to store the idempotency instance in the request.internal object.
 * Value: 'powertools-for-aws.idempotency'
 */
const IDEMPOTENCY_KEY: string;

/**
 * Interface representing a Middy request object.
 */
interface MiddyLikeRequest {
  /** The Lambda event */
  event: unknown;
  /** The Lambda context */
  context: unknown;
  /** The response to be returned */
  response: unknown;
  /** Any error that occurred */
  error: Error | null;
  /** Internal storage for middleware state and cleanup functions */
  internal: Record<string, unknown>;
}

/**
 * Full Middy request type with additional properties.
 */
type Request = MiddyLikeRequest;

/**
 * Type for Middy middleware function.
 */
type MiddlewareFn = (request: Request) => Promise<unknown>;

/**
 * Type for Middy middleware object with lifecycle methods.
 */
type MiddlewareLikeObj = {
  before?: MiddlewareFn;
  after?: MiddlewareFn;
  onError?: MiddlewareFn;
};

/**
 * Type for cleanup function stored in request.internal.
 */
type CleanupFunction = (request: Request) => Promise<void>;

import middy from '@middy/core';
import { cleanupMiddlewares } from '@aws-lambda-powertools/commons';
import type { MiddyLikeRequest } from '@aws-lambda-powertools/commons/types';

// Example middleware that returns early
const myCustomMiddleware = (): middy.MiddlewareObj => {
  const before = async (request: MiddyLikeRequest): Promise<undefined | string> => {
    // Check some condition
    if (request.event.httpMethod === 'GET') {
      // Cleanup Powertools resources before returning early
      await cleanupMiddlewares(request);

      // Return early with response
      return 'GET method not supported';
    }
  };

  return {
    before,
  };
};

// Use the middleware
const handler = middy(async (event, context) => {
  return { statusCode: 200, body: 'Success' };
}).use(myCustomMiddleware());

import middy from '@middy/core';
import { cleanupMiddlewares } from '@aws-lambda-powertools/commons';

const authMiddleware = (): middy.MiddlewareObj => {
  const before = async (request) => {
    const token = request.event.headers?.authorization;

    if (!token) {
      // Cleanup before returning unauthorized response
      await cleanupMiddlewares(request);

      return {
        statusCode: 401,
        body: JSON.stringify({ message: 'Unauthorized' }),
      };
    }

    // Validate token...
    const isValid = validateToken(token);

    if (!isValid) {
      await cleanupMiddlewares(request);

      return {
        statusCode: 403,
        body: JSON.stringify({ message: 'Forbidden' }),
      };
    }
  };

  return { before };
};

import {
  TRACER_KEY,
  METRICS_KEY,
  LOGGER_KEY,
} from '@aws-lambda-powertools/commons';
import type { MiddyLikeRequest } from '@aws-lambda-powertools/commons/types';

// Example: Storing cleanup functions
const myMiddleware = () => {
  const before = async (request: MiddyLikeRequest) => {
    // Store a cleanup function that will be called by cleanupMiddlewares
    request.internal[TRACER_KEY] = async (req) => {
      console.log('Cleaning up tracer');
      // Perform cleanup...
    };

    request.internal[METRICS_KEY] = async (req) => {
      console.log('Flushing metrics');
      // Flush metrics...
    };
  };

  return { before };
};

import middy from '@middy/core';
import { cleanupMiddlewares } from '@aws-lambda-powertools/commons';

const rateLimitMiddleware = (maxRequests: number): middy.MiddlewareObj => {
  let requestCount = 0;

  const before = async (request) => {
    requestCount++;

    if (requestCount > maxRequests) {
      // Cleanup Powertools before returning rate limit response
      await cleanupMiddlewares(request);

      return {
        statusCode: 429,
        body: JSON.stringify({ message: 'Too Many Requests' }),
      };
    }
  };

  return { before };
};

const handler = middy(async (event, context) => {
  return { statusCode: 200, body: 'Success' };
}).use(rateLimitMiddleware(100));

import middy from '@middy/core';
import { cleanupMiddlewares } from '@aws-lambda-powertools/commons';

const validateRequestMiddleware = (schema): middy.MiddlewareObj => {
  const before = async (request) => {
    try {
      // Validate request against schema
      const isValid = schema.validate(request.event);

      if (!isValid) {
        await cleanupMiddlewares(request);

        return {
          statusCode: 400,
          body: JSON.stringify({ message: 'Invalid request' }),
        };
      }
    } catch (error) {
      await cleanupMiddlewares(request);

      return {
        statusCode: 400,
        body: JSON.stringify({ message: 'Validation error' }),
      };
    }
  };

  return { before };
};