APIGateway＋Lambdaでミドルウェアを使う

背景

Go言語でlambdaを使ってサービス提供するときに、複数エンドポイント共通の処理を入れたくなった。
各エンドポイントに共通の処理を書き込むと、可読性が下がるためミドルウェアとして実装してみた。

参考:
https://github.com/mefellows/vesper
https://github.com/aws/aws-lambda-go/blob/main/lambda/handler.go

前提

APIGatewayProxyRequestとAPIGatewayProxyResponseを利用している。

実装

いろいろと既存実装の制約あってreflectを利用してごちゃごちゃしているけど、実行速度考えたらできる限りreflect使わないで、エンドポイントのハンドラーの引数をAPIGatewayProxyRequest、戻り値をAPIGatewayProxyResponseに限定してあげた方がいいかも。

type LambdaFunc func(context.Context, events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error)

type middlewareFunc func(next LambdaFunc) LambdaFunc

type M struct {
	middleware []middlewareFunc
}

func New() *M {
	return &M{}
}

// これで自作したmiddleware用の関数を設定する
func (m *M) Use(middleware middlewareFunc) {
	m.middleware = append(m.middleware, middleware)
}

// lambda.Startをラップして設定したmiddlewareが実行されるようにしている。
func (m *M) Start(handler interface{}) {
	h := apply(NewHandler(handler), m.middleware...)

	lambda.Start(h)
}

func apply(f LambdaFunc, middleware ...middlewareFunc) LambdaFunc {
	if len(middleware) == 0 {
		return f
	}
	return middleware[0](apply(f, middleware[1:]...))
}

// Start()で引数にinterface{}を受け取るため、それをLambdaFunc型にするための処理
// 引数と戻り値の型チェックも一応している
func NewHandler(handlerFunc interface{}) LambdaFunc {
	errorLambdaFunc := func(err error) func(context.Context, events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
		return func(context.Context, events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
			return events.APIGatewayProxyResponse{}, err
		}
	}

	handler := reflect.ValueOf(handlerFunc)
	handlerType := reflect.TypeOf(handlerFunc)
	if handlerType.Kind() != reflect.Func {
		return errorLambdaFunc(fmt.Errorf("handler kind %s is not %s", handlerType.Kind(), reflect.Func))
	}

	takesContext := handlerTakesContext(handlerType)

	var tIn reflect.Type
	if (!takesContext && handlerType.NumIn() == 1) || handlerType.NumIn() == 2 {
		reqType := reflect.TypeOf(events.APIGatewayProxyRequest{})
		tIn = handlerType.In(handlerType.NumIn() - 1)
		if tIn.Kind() != reqType.Kind() {
			return errorLambdaFunc(fmt.Errorf("handler kind %s is not %s", tIn.Kind(), reqType.Kind()))
		}
	}

	return LambdaFunc(func(ctx context.Context, payload events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
		var args []reflect.Value

		if takesContext {
			args = append(args, reflect.ValueOf(ctx))
		}

		if tIn != nil {
			args = append(args, reflect.ValueOf(payload))
		}

		response := handler.Call(args)

		var err error
		if len(response) > 0 {
			if errVal, ok := response[len(response)-1].Interface().(error); ok {
				err = errVal
			}
		}
		APIGatewayProxyResponseType := reflect.TypeOf(events.APIGatewayProxyResponse{})
		var val events.APIGatewayProxyResponse
		if len(response) > 1 {
			if !response[0].CanConvert(APIGatewayProxyResponseType) {
				return events.APIGatewayProxyResponse{}, (fmt.Errorf("handler response kind %s can not convert to %s", response[0].Kind(), APIGatewayProxyResponseType))
			}
			val, _ = response[0].Convert(APIGatewayProxyResponseType).Interface().(events.APIGatewayProxyResponse)
		}

		return val, err
	})
}

func handlerTakesContext(handlerType reflect.Type) bool {
	if handlerType.NumIn() > 0 {
		contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
		argumentType := handlerType.In(0)
		return argumentType.Implements(contextType)
	}
	return false
}

補足

func (m *M) Start(handler interface{}) {
	h := apply(NewHandler(handler), m.middleware...)

	lambda.Start(h)
}

この部分は、エンドポイントのハンドラーが、func(context.Context, events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error)に固定できるならinterface{}で受け取る必要はない。
今回は、events.APIGatewayProxyRequestをDefined Typeにしてしまっているところがあったため、柔軟性を持たせるためにinterface{}で受け取っている。