Remove the Lambdaify handler to simplify code

punmechanic · Nov 12, 2024 · 40b1f72 · 40b1f72
1 parent ada05bf
commit 40b1f72
Show file tree

Hide file tree

Showing 6 changed files with 89 additions and 217 deletions.
diff --git a/internal/api/http.go b/internal/api/http.go
@@ -1,11 +1,11 @@
 package api
 
 import (
-	"net/http"
 	"strings"
 
 	"log/slog"
 
+	"github.com/aws/aws-lambda-go/events"
 	"golang.org/x/oauth2"
 )
 
@@ -14,24 +14,26 @@ import (
 // []any is returned instead of []slog.Attr to make it easier to supply the attributes to slog functions using spread, for example:
 //
 //	slog.Error(msg, RequestAttrs(r)...)
-func RequestAttrs(r *http.Request) []any {
-	attrs := []any{
-		slog.String("origin_ip_address", r.RemoteAddr),
-	}
+func RequestAttrs(r events.ALBTargetGroupRequest) []any {
+	var attrs []any
 
-	if v := r.Header.Get("x-amzn-trace-id"); v != "" {
+	if v, ok := r.Headers["x-amzn-trace-id"]; ok {
 		attrs = append(attrs, slog.String("amz_request_id", v))
 	}
 
-	if v := r.Header.Get("x-forwarded-for"); v != "" {
+	if v, ok := r.Headers["x-forwarded-for"]; ok {
 		attrs = append(attrs, slog.String("x_forwarded_for", v))
 	}
 
 	return attrs
 }
 
-func requestTokenSource(r *http.Request) (oauth2.TokenSource, bool) {
-	headerValue := r.Header.Get("authorization")
+func requestTokenSource(r events.ALBTargetGroupRequest) (oauth2.TokenSource, bool) {
+	headerValue, ok := r.Headers["authorization"]
+	if !ok {
+		return nil, false
+	}
+
 	parts := strings.Split(headerValue, " ")
 	if len(parts) != 2 {
 		return nil, false

diff --git a/internal/api/json.go b/internal/api/json.go
@@ -2,30 +2,30 @@ package api
 
 import (
 	"encoding/json"
-	"net/http"
 
 	"log/slog"
+
+	"github.com/aws/aws-lambda-go/events"
 )
 
-func ServeJSON[T any](w http.ResponseWriter, data T) {
+func ServeJSON[T any](w *events.ALBTargetGroupResponse, data T) {
 	buf, err := json.Marshal(data)
 	if err != nil {
 		// Nothing to be done here
 		slog.Error("could not marshal JSON: %s", "error", err)
 		return
 	}
 
-	// Nothing we can do to respond to the error message here either
-	w.Header().Set("Content-Type", "application/json")
-	_, err = w.Write(buf)
-	slog.Error("could not write JSON to the client: %s", "error", err)
+	w.Headers["Content-Type"] = "application/json"
+	w.Body = string(buf)
 }
 
-type JSONError struct {
-	Message string `json:"error"`
-}
+func ServeJSONError(w *events.ALBTargetGroupResponse, statusCode int, msg string) {
+	var jsonError struct {
+		Message string `json:"error"`
+	}
 
-func ServeJSONError(w http.ResponseWriter, statusCode int, msg string) {
-	w.WriteHeader(statusCode)
-	ServeJSON[JSONError](w, JSONError{msg})
+	jsonError.Message = msg
+	w.StatusCode = statusCode
+	ServeJSON(w, jsonError)
 }
diff --git a/internal/api/serverless_functions.go b/internal/api/serverless_functions.go
@@ -7,6 +7,8 @@ import (
 
 	"log/slog"
 
+	"github.com/aws/aws-lambda-go/events"
+	"github.com/aws/aws-lambda-go/lambda"
 	"github.com/coreos/go-oidc"
 	"github.com/okta/okta-sdk-golang/v2/okta"
 )
@@ -20,58 +22,74 @@ type OktaService interface {
 	ListApplicationsForUser(ctx context.Context, user string) ([]*okta.AppLink, error)
 }
 
-func ServeUserApplications(okta OktaService, idp *oidc.Provider) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ctx := r.Context()
-		requestAttrs := RequestAttrs(r)
-		ts, ok := requestTokenSource(r)
-		if !ok {
-			slog.Error("no bearer token present", requestAttrs...)
-			ServeJSONError(w, http.StatusUnauthorized, "unauthorized")
-			return
-		}
+type ServeUserApplicationsHandler struct {
+	Okta OktaService
+	Idp  *oidc.Provider
+}
+
+func (s ServeUserApplicationsHandler) Handle(ctx context.Context, r events.ALBTargetGroupRequest) (w events.ALBTargetGroupResponse, err error) {
+	requestAttrs := RequestAttrs(r)
+	ts, ok := requestTokenSource(r)
+	if !ok {
+		slog.Error("no bearer token present", requestAttrs...)
+		ServeJSONError(&w, http.StatusUnauthorized, "unauthorized")
+		return
+	}
 
-		info, err := idp.UserInfo(ctx, ts)
-		if err != nil {
-			if errors.Is(err, ErrBadRequest) {
-				slog.Error("okta indicated the request was poorly formed", requestAttrs...)
-				ServeJSONError(w, http.StatusInternalServerError, "internal error when talking to the Okta API")
-				return
-			}
-
-			requestAttrs = append(requestAttrs, slog.String("error", err.Error()))
-			slog.Error("okta rejected id token", requestAttrs...)
-			ServeJSONError(w, http.StatusForbidden, "unauthorized")
+	info, err := s.Idp.UserInfo(ctx, ts)
+	if err != nil {
+		if errors.Is(err, ErrBadRequest) {
+			slog.Error("okta indicated the request was poorly formed", requestAttrs...)
+			ServeJSONError(&w, http.StatusInternalServerError, "internal error when talking to the Okta API")
 			return
 		}
 
-		var claims Claims
-		if err := info.Claims(&claims); err != nil {
-			slog.Error("failed to parse claims from Okta userinfo endpoint", requestAttrs...)
-			ServeJSONError(w, http.StatusInternalServerError, "internal error when talking to the Okta API")
-		}
+		requestAttrs = append(requestAttrs, slog.String("error", err.Error()))
+		slog.Error("okta rejected id token", requestAttrs...)
+		ServeJSONError(&w, http.StatusForbidden, "unauthorized")
+		return
+	}
 
-		requestAttrs = append(requestAttrs, slog.String("username", claims.PreferredUsername))
-		applications, err := okta.ListApplicationsForUser(ctx, claims.PreferredUsername)
-		if err != nil {
-			requestAttrs = append(requestAttrs, slog.String("error", err.Error()))
-			slog.Error("failed to fetch applications", requestAttrs...)
-			ServeJSONError(w, http.StatusBadGateway, "upstream error")
-			return
-		}
+	var claims Claims
+	if err := info.Claims(&claims); err != nil {
+		slog.Error("failed to parse claims from Okta userinfo endpoint", requestAttrs...)
+		ServeJSONError(&w, http.StatusInternalServerError, "internal error when talking to the Okta API")
+	}
 
-		var accounts []Application
-		for _, app := range applications {
-			if app.AppName == "amazon_aws" {
-				accounts = append(accounts, Application{
-					ID:   app.AppInstanceId,
-					Name: app.Label,
-				})
-			}
+	requestAttrs = append(requestAttrs, slog.String("username", claims.PreferredUsername))
+	applications, err := s.Okta.ListApplicationsForUser(ctx, claims.PreferredUsername)
+	if err != nil {
+		requestAttrs = append(requestAttrs, slog.String("error", err.Error()))
+		slog.Error("failed to fetch applications", requestAttrs...)
+		ServeJSONError(&w, http.StatusBadGateway, "upstream error")
+		return
+	}
+
+	var accounts []Application
+	for _, app := range applications {
+		if app.AppName == "amazon_aws" {
+			accounts = append(accounts, Application{
+				ID:   app.AppInstanceId,
+				Name: app.Label,
+			})
 		}
+	}
+
+	requestAttrs = append(requestAttrs, slog.Int("application_count", len(accounts)))
+	slog.Info("served applications", requestAttrs...)
+	ServeJSON(&w, accounts)
+	return
+}
+
+func (s ServeUserApplicationsHandler) Handler() lambda.Handler {
+	return lambda.NewHandler(s.Handle)
+}
+
+func ServeUserApplications(okta OktaService, idp *oidc.Provider) lambda.Handler {
+	h := ServeUserApplicationsHandler{
+		Okta: okta,
+		Idp:  idp,
+	}
 
-		requestAttrs = append(requestAttrs, slog.Int("application_count", len(accounts)))
-		slog.Info("served applications", requestAttrs...)
-		ServeJSON(w, accounts)
-	})
+	return h.Handler()
 }
diff --git a/internal/lambdaify.go b/internal/lambdaify.go
diff --git a/internal/lambdaify_test.go b/internal/lambdaify_test.go
diff --git a/lambda/list_applications/main.go b/lambda/list_applications/main.go
@@ -9,7 +9,6 @@ import (
 
 	"github.com/aws/aws-lambda-go/lambda"
 	"github.com/coreos/go-oidc"
-	"github.com/riotgames/key-conjurer/internal"
 	"github.com/riotgames/key-conjurer/internal/api"
 )
 
@@ -33,5 +32,6 @@ func main() {
 	}
 
 	slog.Info("running list_applications_v2 Lambda")
-	lambda.Start(internal.Lambdaify(api.ServeUserApplications(service, idp)))
+
+	lambda.Start(api.ServeUserApplications(service, idp))
 }