Merge pull request #3 from DefangLabs/edw/add-listener-rule

Add listener rule when invoked the first time

Author: edwardrf

acme/update.go

@@ -10,9 +10,9 @@ import (
 	"log"
 	"os"
 
-	awsalb "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
 	"github.com/DefangLabs/cloudacme/aws/acm"
 	"github.com/DefangLabs/cloudacme/aws/alb"
+	awsalb "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
 	"github.com/mholt/acmez"
 	"go.uber.org/zap"
 )
@@ -33,14 +33,43 @@ func UpdateAcmeCertificate(ctx context.Context, albArn, domain string, solver ac
 		return fmt.Errorf("failed to get account key: %w", err)
 	}
 
+	certToUpdate, _, err := GetExistingCertificate(ctx, albArn, domain)
+	if err != nil {
+		return fmt.Errorf("failed to get existing certificate: %w", err)
+	}
+
+	acmeDirectory := os.Getenv("ACME_DIRECTORY")
+	if acmeDirectory == "" {
+		acmeDirectory = DefaultAcmeDirectory
+	}
+
+	acmeClient := Acme{
+		Directory:  acmeDirectory,
+		AccountKey: accountKey,
+		Logger:     logger,
+		AlbArn:     albArn,
+		HttpSolver: solver,
+	}
+
+	key, chain, err := acmeClient.GetCertificate(ctx, []string{domain})
+	if err != nil {
+		return fmt.Errorf("failed to get certificates: %w", err)
+	}
+
+	if err := acm.ImportCertificate(ctx, key, chain, certToUpdate); err != nil {
+		return fmt.Errorf("error importing certificate: %w", err)
+	}
+	return nil
+}
+
+func GetExistingCertificate(ctx context.Context, albArn, domain string) (string, *x509.Certificate, error) {
 	// Find the certificate to update from all the certificates attached to the ALB
 	certArns, err := alb.GetAlbCerts(ctx, albArn)
 	if err != nil {
-		return fmt.Errorf("failed to get ALB certificates: %w", err)
+		return "", nil, fmt.Errorf("failed to get ALB certificates: %w", err)
 	}
 
 	var getCertErrs []error
-	certToUpdate := ""
 	for _, certArn := range certArns {
 		certPem, err := acm.GetCertificate(ctx, certArn)
 		if err != nil {
@@ -60,37 +89,25 @@ func UpdateAcmeCertificate(ctx context.Context, albArn, domain string, solver ac
 		if cert.Subject.CommonName == domain {
 			// TODO: check the issuer and expiration date
 			// TODO: should we check SANs? probably not, as byod domain are added as SNI single domain certs
-			certToUpdate = certArn
-			break
+			return certArn, cert, nil
 		}
 	}
-	if certToUpdate == "" {
-		if len(getCertErrs) == 0 {
-			return fmt.Errorf("no certificate matching %v found", domain)
-		}
-		return fmt.Errorf("failed to get certificate: %w", errors.Join(getCertErrs...))
-	}
-
-	acmeDirectory := os.Getenv("ACME_DIRECTORY")
-	if acmeDirectory == "" {
-		acmeDirectory = DefaultAcmeDirectory
-	}
+	return "", nil, fmt.Errorf("no certificate matching %v found: %w", domain, errors.Join(getCertErrs...))
+}
 
-	acmeClient := Acme{
-		Directory:  acmeDirectory,
-		AccountKey: accountKey,
-		Logger:     logger,
-		AlbArn:     albArn,
-		HttpSolver: solver,
+func SetupHttpRule(ctx context.Context, albArn, lambdaArn string, ruleCond alb.RuleCondition) error {
+	listener, err := alb.GetListener(ctx, albArn, awsalb.ProtocolEnumHttp, 80)
+	if err != nil {
+		return fmt.Errorf("cannot get http listener: %w", err)
 	}
 
-	key, chain, err := acmeClient.GetCertificate(ctx, []string{domain})
+	targetGroupArn, err := alb.GetLambdaTargetGroup(ctx, lambdaArn)
 	if err != nil {
-		return fmt.Errorf("failed to get certificates: %w", err)
+		return fmt.Errorf("cannot get target group for lambda %v: %w", lambdaArn, err)
 	}
 
-	if err := acm.ImportCertificate(ctx, key, chain, certToUpdate); err != nil {
-		return fmt.Errorf("error importing certificate: %w", err)
+	if err := alb.AddListenerTriggerTargetGroupRule(ctx, *listener.ListenerArn, ruleCond, targetGroupArn); err != nil {
+		return fmt.Errorf("failed to create listener static rule: %w", err)
 	}
 	return nil
 }

aws/alb/updatealb.go

@@ -4,8 +4,10 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"log"
 	"sort"
 	"strconv"
+	"strings"
 
 	"github.com/DefangLabs/cloudacme/aws"
 	elbv2 "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
@@ -132,6 +134,79 @@ func AddListenerStaticRule(ctx context.Context, listenerArn string, ruleCond Rul
 	return nil
 }
 
+func AddListenerTriggerTargetGroupRule(ctx context.Context, listenerArn string, ruleCond RuleCondition, targetArn string) error {
+	svc := elbv2.NewFromConfig(aws.LoadConfig())
+
+	priority, err := GetNextAvailablePriority(ctx, listenerArn)
+	if err != nil {
+		return err
+	}
+
+	input := &elbv2.CreateRuleInput{
+		Actions: []types.Action{
+			{
+				Type:           types.ActionTypeEnumForward,
+				TargetGroupArn: &targetArn,
+			},
+		},
+		Conditions: []types.RuleCondition{
+			{
+				Field:             ptr.String("path-pattern"),
+				PathPatternConfig: &types.PathPatternConditionConfig{Values: ruleCond.PathPattern},
+			},
+			{
+				Field:            ptr.String("host-header"),
+				HostHeaderConfig: &types.HostHeaderConditionConfig{Values: ruleCond.HostHeader},
+			},
+		},
+		ListenerArn: &listenerArn,
+		Priority:    ptr.Int32(priority),
+	}
+
+	if _, err := svc.CreateRule(ctx, input); err != nil {
+		return err
+	}
+	return nil
+}
+
+func GetLambdaTargetGroup(ctx context.Context, lambdaArn string) (string, error) {
+	svc := elbv2.NewFromConfig(aws.LoadConfig())
+	paginator := elbv2.NewDescribeTargetGroupsPaginator(svc, &elbv2.DescribeTargetGroupsInput{})
+
+	log.Printf("Searching for target group for lambda %s", lambdaArn)
+	for paginator.HasMorePages() {
+		page, err := paginator.NextPage(ctx)
+		if err != nil {
+			return "", fmt.Errorf("failed to list target groups: %w", err)
+		}
+
+		for _, tg := range page.TargetGroups {
+			log.Printf("Checking target group %s of type %s", *tg.TargetGroupArn, tg.TargetType)
+			if tg.TargetType != types.TargetTypeEnumLambda {
+				continue
+			}
+
+			// Check registered targets for this target group
+			targetsOut, err := svc.DescribeTargetHealth(ctx, &elbv2.DescribeTargetHealthInput{
+				TargetGroupArn: tg.TargetGroupArn,
+			})
+			if err != nil {
+				return "", fmt.Errorf("describe target health failed for %s: %w", *tg.TargetGroupArn, err)
+			}
+
+			for _, desc := range targetsOut.TargetHealthDescriptions {
+				if desc.Target != nil && desc.Target.Id != nil {
+					log.Printf("Checking target %s with status %s", *desc.Target.Id, desc.TargetHealth.State)
+					if strings.HasPrefix(lambdaArn, *desc.Target.Id) {
+						return *tg.TargetGroupArn, nil
+					}
+				}
+			}
+		}
+	}
+	return "", fmt.Errorf("no target group found for lambda %s", lambdaArn)
+}
+
 func GetNextAvailablePriority(ctx context.Context, listenerArn string) (int32, error) {
 	rules, err := GetAllRules(ctx, listenerArn)
 	if err != nil {

cmd/lambda/main.go

@@ -3,18 +3,21 @@ package main
 import (
 	"context"
 	"crypto/tls"
+	"crypto/x509"
 	"errors"
 	"fmt"
 	"log"
 	"net/http"
 	"net/url"
+	"strings"
 	"time"
 
-	"github.com/aws/aws-lambda-go/events"
-	"github.com/aws/aws-lambda-go/lambda"
 	"github.com/DefangLabs/cloudacme/acme"
 	"github.com/DefangLabs/cloudacme/aws/alb"
 	"github.com/DefangLabs/cloudacme/solver"
+	"github.com/aws/aws-lambda-go/events"
+	"github.com/aws/aws-lambda-go/lambda"
+	"github.com/aws/aws-lambda-go/lambdacontext"
 )
 
 var version = "dev" // to be set by ldflags
@@ -34,7 +37,28 @@ func HandleEvent(ctx context.Context, evt Event) (any, error) {
 	if evt.HTTPMethod != "" {
 		return HandleALBEvent(ctx, evt.ALBTargetGroupRequest)
 	} else {
-		return nil, HandleEventBridgeEvent(ctx, evt.CertificateRenewalEvent)
+		_, cert, err := acme.GetExistingCertificate(ctx, evt.AlbArn, evt.Domain)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get existing certificate: %w", err)
+		}
+
+		ownArn := ""
+		if lc, ok := lambdacontext.FromContext(ctx); ok {
+			ownArn = lc.InvokedFunctionArn
+		}
+		if ownArn == "" {
+			return nil, errors.New("unable to determine own Lambda ARN from context")
+		}
+
+		if !IsLetsEncryptCertificate(cert) {
+			log.Printf("Certificate for domain %s is not issued by Let's Encrypt, initial run, setup load balancer rule for acme lambda", evt.Domain)
+			return nil, acme.SetupHttpRule(ctx, evt.AlbArn, ownArn, alb.RuleCondition{
+				HostHeader:  []string{evt.Domain},
+				PathPattern: []string{"/"},
+			})
+		} else {
+			return nil, HandleScheduledRenewalEvent(ctx, evt.CertificateRenewalEvent)
+		}
 	}
 }
 
@@ -124,7 +148,7 @@ func getHttpsRedirectURL(evt events.ALBTargetGroupRequest) string {
 	return fmt.Sprintf("https://%s%s%s", evt.Headers["host"], evt.Path, params)
 }
 
-func HandleEventBridgeEvent(ctx context.Context, evt CertificateRenewalEvent) error {
+func HandleScheduledRenewalEvent(ctx context.Context, evt CertificateRenewalEvent) error {
 	log.Printf("Handling Certificate Renewal Event: %+v", evt)
 
 	albSolver := solver.AlbHttp01Solver{
@@ -139,6 +163,22 @@ func HandleEventBridgeEvent(ctx context.Context, evt CertificateRenewalEvent) er
 	return nil
 }
 
+func IsLetsEncryptCertificate(cert *x509.Certificate) bool {
+	// Check Issuer Organization
+	for _, org := range cert.Issuer.Organization {
+		if strings.Contains(strings.ToLower(org), "let's encrypt") {
+			return true
+		}
+	}
+
+	// Fallback: check Common Name
+	if strings.Contains(strings.ToLower(cert.Issuer.CommonName), "let's encrypt") {
+		return true
+	}
+
+	return false
+}
+
 func main() {
 	lambda.Start(HandleEvent)
 }