|
1 | 1 | package cmd
|
2 | 2 |
|
3 | 3 | import (
|
4 |
| - "encoding/json" |
| 4 | + "context" |
5 | 5 | "fmt"
|
6 |
| - "io" |
7 | 6 | "net/http"
|
8 | 7 | "os"
|
9 | 8 | "time"
|
10 | 9 |
|
| 10 | + "github.com/aws/aws-lambda-go/lambda" |
11 | 11 | log "github.com/sirupsen/logrus"
|
12 | 12 | "github.com/spf13/cobra"
|
13 | 13 |
|
14 | 14 | "github.com/gitpod-io/enterprise-deployment-toolkit/gitpod-network-check/pkg/lambda_types"
|
15 | 15 | )
|
16 | 16 |
|
17 |
| -var lambdaHandlerCmd = &cobra.Command{ |
18 |
| - Use: "lambda-handler", |
19 |
| - Short: "Internal command to execute network checks within AWS Lambda (reads JSON request from stdin, writes JSON response to stdout)", |
20 |
| - Hidden: true, // Hide this command from user help output |
21 |
| - PersistentPreRun: func(cmd *cobra.Command, args []string) { |
22 |
| - // override parent, as we don't care about the config or other flags |
23 |
| - }, |
24 |
| - RunE: func(cmd *cobra.Command, args []string) error { |
25 |
| - // Lambda environment might not have sophisticated logging setup, print directly |
26 |
| - fmt.Fprintln(os.Stderr, "Lambda Handler: Starting execution.") |
27 |
| - |
28 |
| - // Read request payload from stdin |
29 |
| - stdinBytes, err := io.ReadAll(os.Stdin) |
30 |
| - if err != nil { |
31 |
| - fmt.Fprintf(os.Stderr, "Lambda Handler: Error reading stdin: %v\n", err) |
32 |
| - return fmt.Errorf("error reading stdin: %w", err) |
33 |
| - } |
| 17 | +// Core logic for handling the Lambda event |
| 18 | +// This function is called by the aws-lambda-go library. |
| 19 | +func handleLambdaEvent(ctx context.Context, request lambda_types.CheckRequest) (lambda_types.CheckResponse, error) { |
| 20 | + log.Infof("Lambda Handler: Received check request for %d endpoints.", len(request.Endpoints)) |
34 | 21 |
|
35 |
| - var request lambda_types.CheckRequest |
36 |
| - err = json.Unmarshal(stdinBytes, &request) |
37 |
| - if err != nil { |
38 |
| - fmt.Fprintf(os.Stderr, "Lambda Handler: Error unmarshalling request JSON: %v\n", err) |
39 |
| - fmt.Fprintf(os.Stderr, "Lambda Handler: Received input: %s\n", string(stdinBytes)) |
40 |
| - return fmt.Errorf("error unmarshalling request: %w", err) |
41 |
| - } |
| 22 | + response := lambda_types.CheckResponse{ |
| 23 | + Results: make(map[string]lambda_types.CheckResult), |
| 24 | + } |
42 | 25 |
|
43 |
| - fmt.Fprintf(os.Stderr, "Lambda Handler: Received check request for %d endpoints.\n", len(request.Endpoints)) |
| 26 | + client := &http.Client{ |
| 27 | + Timeout: 10 * time.Second, // Consider making this configurable if needed |
| 28 | + } |
44 | 29 |
|
45 |
| - response := lambda_types.CheckResponse{ |
46 |
| - Results: make(map[string]lambda_types.CheckResult), |
47 |
| - } |
| 30 | + // Perform checks |
| 31 | + for name, url := range request.Endpoints { |
| 32 | + log.Debugf("Lambda Handler: Checking endpoint: %s (%s)", name, url) |
48 | 33 |
|
49 |
| - client := &http.Client{ |
50 |
| - Timeout: 10 * time.Second, // Slightly longer timeout for Lambda environment? |
| 34 | + // Use the context provided by the Lambda runtime |
| 35 | + req, err := http.NewRequestWithContext(ctx, "GET", url, nil) |
| 36 | + if err != nil { |
| 37 | + response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("failed to create request: %v", err)} |
| 38 | + log.Warnf(" -> Failed (request creation): %v", err) |
| 39 | + continue |
51 | 40 | }
|
52 | 41 |
|
53 |
| - // Perform checks (similar logic to the previous dedicated handler) |
54 |
| - for name, url := range request.Endpoints { |
55 |
| - fmt.Fprintf(os.Stderr, "Lambda Handler: Checking endpoint: %s (%s)\n", name, url) |
56 |
| - // Use context from command if needed, otherwise background context is fine here |
57 |
| - req, err := http.NewRequestWithContext(cmd.Context(), "GET", url, nil) |
58 |
| - if err != nil { |
59 |
| - response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("failed to create request: %v", err)} |
60 |
| - fmt.Fprintf(os.Stderr, " -> Failed (request creation): %v\n", err) |
61 |
| - continue |
62 |
| - } |
63 |
| - |
64 |
| - resp, err := client.Do(req) |
65 |
| - if err != nil { |
66 |
| - response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("HTTP request failed: %v", err)} |
67 |
| - fmt.Fprintf(os.Stderr, " -> Failed (HTTP request): %v\n", err) |
| 42 | + resp, err := client.Do(req) |
| 43 | + if err != nil { |
| 44 | + response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("HTTP request failed: %v", err)} |
| 45 | + log.Warnf(" -> Failed (HTTP request): %v", err) |
| 46 | + } else { |
| 47 | + resp.Body.Close() // Ensure body is closed |
| 48 | + if resp.StatusCode >= 200 && resp.StatusCode < 300 { |
| 49 | + response.Results[name] = lambda_types.CheckResult{Success: true} |
| 50 | + log.Debugf(" -> Success (Status: %d)", resp.StatusCode) |
68 | 51 | } else {
|
69 |
| - resp.Body.Close() // Ensure body is closed |
70 |
| - if resp.StatusCode >= 200 && resp.StatusCode < 300 { |
71 |
| - response.Results[name] = lambda_types.CheckResult{Success: true} |
72 |
| - fmt.Fprintf(os.Stderr, " -> Success (Status: %d)\n", resp.StatusCode) |
73 |
| - } else { |
74 |
| - response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("unexpected status code: %d", resp.StatusCode)} |
75 |
| - fmt.Fprintf(os.Stderr, " -> Failed (Status: %d)\n", resp.StatusCode) |
76 |
| - } |
| 52 | + response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("unexpected status code: %d", resp.StatusCode)} |
| 53 | + log.Warnf(" -> Failed (Status: %d)", resp.StatusCode) |
77 | 54 | }
|
78 | 55 | }
|
| 56 | + } |
79 | 57 |
|
80 |
| - // Marshal response payload to stdout |
81 |
| - responseBytes, err := json.Marshal(response) |
82 |
| - if err != nil { |
83 |
| - fmt.Fprintf(os.Stderr, "Lambda Handler: Error marshalling response JSON: %v\n", err) |
84 |
| - return fmt.Errorf("error marshalling response: %w", err) |
85 |
| - } |
86 |
| - |
87 |
| - _, err = fmt.Fprint(os.Stdout, string(responseBytes)) |
88 |
| - if err != nil { |
89 |
| - fmt.Fprintf(os.Stderr, "Lambda Handler: Error writing response to stdout: %v\n", err) |
90 |
| - return fmt.Errorf("error writing response: %w", err) |
91 |
| - } |
| 58 | + log.Info("Lambda Handler: Check processing complete.") |
| 59 | + // The lambda library handles marshalling the response and deals with errors. |
| 60 | + // We return the response struct and nil error if processing logic itself didn't fail critically. |
| 61 | + return response, nil |
| 62 | +} |
92 | 63 |
|
93 |
| - fmt.Fprintln(os.Stderr, "Lambda Handler: Execution complete.") |
94 |
| - return nil |
| 64 | +// lambdaHandlerCmd is the Cobra command invoked when the binary is run with the "lambda-handler" argument. |
| 65 | +// This happens inside the AWS Lambda environment via the bootstrap script. |
| 66 | +var lambdaHandlerCmd = &cobra.Command{ |
| 67 | + Use: "lambda-handler", |
| 68 | + Short: "Internal command used by AWS Lambda runtime to execute network checks", |
| 69 | + Hidden: true, // Hide this command from user help output |
| 70 | + PersistentPreRun: func(cmd *cobra.Command, args []string) { |
| 71 | + // override parent, as we don't care about the config or other flags when run by lambda |
| 72 | + // Ensure logs go to stderr (Lambda standard) |
| 73 | + log.SetOutput(os.Stderr) |
| 74 | + // Optionally set log level from env var if needed, e.g., os.Getenv("LOG_LEVEL") |
| 75 | + // Consider setting a default level appropriate for Lambda execution. |
| 76 | + log.SetLevel(log.InfoLevel) // Example: Set a default level |
| 77 | + }, |
| 78 | + RunE: func(cmd *cobra.Command, args []string) error { |
| 79 | + // The aws-lambda-go library takes over execution when lambda.Start is called. |
| 80 | + // It handles reading events, invoking the handler, and writing responses. |
| 81 | + log.Info("Lambda Handler: Starting AWS Lambda handler loop.") |
| 82 | + lambda.Start(handleLambdaEvent) |
| 83 | + // lambda.Start blocks and never returns unless there's a critical error during initialization |
| 84 | + log.Error("Lambda Handler: lambda.Start returned unexpectedly (should not happen)") |
| 85 | + return fmt.Errorf("lambda.Start returned unexpectedly") |
95 | 86 | },
|
96 |
| - // Disable flag parsing for this internal command as it gets input via stdin |
| 87 | + // Disable flag parsing for this internal command as input comes from Lambda event payload |
97 | 88 | DisableFlagParsing: true,
|
98 | 89 | }
|
99 | 90 |
|
100 | 91 | func init() {
|
101 |
| - // Note: We don't add this to networkCheckCmd directly in init() here |
102 |
| - // because it might interfere with normal flag parsing if not careful. |
103 |
| - // It will be added in the main Execute() function or similar central place. |
104 |
| - // For now, just define the command struct. |
105 |
| - // We also need to ensure logging doesn't interfere with stdout JSON output. |
106 |
| - // Maybe configure logging to stderr specifically for this command? |
107 |
| - lambdaHandlerCmd.PersistentPreRun = func(cmd *cobra.Command, args []string) { |
108 |
| - // Ensure logs go to stderr for this command to keep stdout clean for JSON |
109 |
| - log.SetOutput(os.Stderr) |
110 |
| - } |
111 |
| - |
112 |
| - NetworkCheckCmd.AddCommand(lambdaHandlerCmd) // Register the hidden lambda handler command |
| 92 | + // Register the hidden lambda handler command |
| 93 | + // It's invoked by the Lambda runtime via the bootstrap script |
| 94 | + NetworkCheckCmd.AddCommand(lambdaHandlerCmd) |
113 | 95 | }
|
0 commit comments