implement grpc health check service

chart/iam-runtime-infratographer/templates/_container.tpl

@@ -19,6 +19,16 @@ envFrom:
 env:
  {{- toYaml . | nindent 2 }}
 {{- end }}
+{{- if $values.livenessProbe.enabled }}
+{{- with omit $values.livenessProbe "enabled" }}
+livenessProbe: {{- toYaml . | nindent 2 }}
+{{- end }}
+{{- end }}
+{{- if $values.readinessProbe.enabled }}
+{{- with omit $values.readinessProbe "enabled" }}
+readinessProbe: {{- toYaml . | nindent 2 }}
+{{- end }}
+{{- end }}
 volumeMounts:
   - name: {{ include "iam-runtime-infratographer.resource.fullname" (dict "suffix" "config" "context" $) | quote }}
     mountPath: /etc/iam-runtime-infratographer/

chart/iam-runtime-infratographer/values.yaml

@@ -127,3 +127,19 @@ securityContext:
   readOnlyRootFilesystem: true
   runAsNonRoot: true
   runAsUser: 65532
+
+livenessProbe:
+  # -- enables liveness probe.
+  enabled: true
+  grpc:
+    # -- sets the grpc health service port.
+    port: 4784
+  timeoutSeconds: 10
+
+readinessProbe:
+  # -- enables readiness probe.
+  enabled: true
+  grpc:
+    # -- sets the grpc health service port.
+    port: 4784
+  timeoutSeconds: 10

go.mod

@@ -13,6 +13,7 @@ require (
 	github.com/hashicorp/go-retryablehttp v0.7.7
 	github.com/labstack/echo/v4 v4.12.0
 	github.com/metal-toolbox/iam-runtime v0.4.1
+	github.com/nats-io/nats.go v1.36.0
 	github.com/spf13/cobra v1.8.1
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.19.0
@@ -52,7 +53,6 @@ require (
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
-	github.com/nats-io/nats.go v1.36.0 // indirect
 	github.com/nats-io/nkeys v0.4.7 // indirect
 	github.com/nats-io/nuid v1.0.1 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.2 // indirect

internal/accesstoken/tokensource.go

@@ -2,17 +2,25 @@ package accesstoken
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net/url"
 	"os"
 	"path/filepath"
 	"strings"
 
 	"go.infratographer.com/iam-runtime-infratographer/internal/jwt"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/clientcredentials"
 )
 
+const tracerName = "go.infratographer.com/iam-runtime-infratographer/internal/accesstoken"
+
+var tracer = otel.GetTracerProvider().Tracer(tracerName)
+
 func (c Config) toTokenSource(ctx context.Context) (oauth2.TokenSource, error) {
 	source, err := c.Source.toTokenSource(ctx)
 	if err != nil {
@@ -105,12 +113,71 @@ func (c AccessTokenExchangeConfig) toTokenSource(ctx context.Context, upstream o
 	return newExchangeTokenSource(ctx, c, upstream)
 }
 
+// HealthyTokenSource extends oauth2.TokenSource implementing the HealthChecker interface.
+type HealthyTokenSource interface {
+	oauth2.TokenSource
+
+	// HealthCheck returns nil when the service is healthy.
+	HealthCheck(ctx context.Context) error
+}
+
+type healthyTokenSource struct {
+	oauth2.TokenSource
+}
+
+// HealthCheck returns nil when the service is healthy.
+func (s *healthyTokenSource) HealthCheck(ctx context.Context) error {
+	ctx, span := tracer.Start(ctx, "HealthCheck")
+	defer span.End()
+
+	errCh := make(chan error, 1)
+
+	go func() {
+		_, err := s.TokenSource.Token()
+		errCh <- err
+
+		close(errCh)
+	}()
+
+	select {
+	case <-ctx.Done():
+		span.SetStatus(codes.Error, ctx.Err().Error())
+		span.RecordError(ctx.Err())
+		span.SetAttributes(attribute.String("healthcheck.outcome", "unhealthy"))
+
+		return ctx.Err()
+	case err := <-errCh:
+		if err != nil {
+			if errors.Is(err, ErrAccessTokenSourceNotEnabled) {
+				span.SetAttributes(attribute.String("healthcheck.outcome", "disabled"))
+
+				return nil
+			} else {
+				span.SetStatus(codes.Error, err.Error())
+				span.RecordError(err)
+				span.SetAttributes(attribute.String("healthcheck.outcome", "unhealthy"))
+
+				return err
+			}
+		}
+	}
+
+	span.SetAttributes(attribute.String("healthcheck.outcome", "healthy"))
+
+	return nil
+}
+
 // NewTokenSource initializes a new token source from the provided config.
 // If the config has Enabled false, then a disabled token source is returned.
-func NewTokenSource(ctx context.Context, cfg Config) (oauth2.TokenSource, error) {
+func NewTokenSource(ctx context.Context, cfg Config) (HealthyTokenSource, error) {
 	if !cfg.Enabled {
-		return &disabledTokenSource{}, nil
+		return &healthyTokenSource{&disabledTokenSource{}}, nil
+	}
+
+	ts, err := cfg.toTokenSource(ctx)
+	if err != nil {
+		return nil, err
 	}
 
-	return cfg.toTokenSource(ctx)
+	return &healthyTokenSource{ts}, nil
 }

internal/eventsx/errors.go

@@ -2,5 +2,10 @@ package eventsx
 
 import "errors"
 
-// ErrPublishNotEnabled represents an error state where an event publish was attempted despite not being enabled
-var ErrPublishNotEnabled = errors.New("event publishing is not enabled")
+var (
+	// ErrPublishNotEnabled represents an error state where an event publish was attempted despite not being enabled
+	ErrPublishNotEnabled = errors.New("event publishing is not enabled")
+
+	// ErrPublisherNotConnected is returned when the underlying connection status is not CONNECTED.
+	ErrPublisherNotConnected = errors.New("event publisher is not connected")
+)

internal/eventsx/publisher.go

@@ -2,14 +2,26 @@ package eventsx
 
 import (
 	"context"
+	"fmt"
 
+	"github.com/nats-io/nats.go"
 	"go.infratographer.com/x/events"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
 )
 
+const tracerName = "go.infratographer.com/iam-runtime-infratographer/internal/eventsx"
+
+var tracer = otel.GetTracerProvider().Tracer(tracerName)
+
 // Publisher represents something that sends relationships to permissions-api via NATS.
 type Publisher interface {
 	// PublishAuthRelationship is similar to events.Publisher.PublishAuthRelationship, but with no topic.
 	PublishAuthRelationshipRequest(ctx context.Context, message events.AuthRelationshipRequest) (events.Message[events.AuthRelationshipResponse], error)
+
+	// HealthCheck returns nil when the service is healthy.
+	HealthCheck(ctx context.Context) error
 }
 
 type publisher struct {
@@ -25,6 +37,30 @@ func (p publisher) PublishAuthRelationshipRequest(ctx context.Context, message e
 	return p.innerPub.PublishAuthRelationshipRequest(ctx, p.topic, message)
 }
 
+// HealthCheck returns nil when the service is healthy.
+func (p publisher) HealthCheck(ctx context.Context) error {
+	_, span := tracer.Start(ctx, "HealthCheck")
+	defer span.End()
+
+	if !p.enabled {
+		span.SetAttributes(attribute.String("healthcheck.outcome", "disabled"))
+
+		return nil
+	}
+
+	conn := p.innerPub.(*events.NATSConnection).Source().(*nats.Conn)
+	if conn.Status() != nats.CONNECTED {
+		span.SetStatus(codes.Error, fmt.Sprintf("status not connected: %s", conn.Status()))
+		span.SetAttributes(attribute.String("healthcheck.outcome", "unhealthy"))
+
+		return fmt.Errorf("%w: status: %s", ErrPublisherNotConnected, conn.Status())
+	}
+
+	span.SetAttributes(attribute.String("healthcheck.outcome", "healthy"))
+
+	return nil
+}
+
 // NewPublisher creates a new events publisher from the given config.
 func NewPublisher(cfg Config) (Publisher, error) {
 	if !cfg.Enabled {

internal/jwt/validator.go

@@ -1,25 +1,44 @@
 package jwt
 
 import (
+	"context"
+	"errors"
 	"net/http"
 	"net/url"
 
 	"github.com/MicahParks/jwkset"
 	"github.com/MicahParks/keyfunc/v3"
 	"github.com/golang-jwt/jwt/v5"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
+)
+
+const tracerName = "go.infratographer.com/iam-runtime-infratographer/internal/jwt"
+
+var (
+	// ErrIssuerKeysMissing is returned by the health check when no issuer keys exist in the store.
+	ErrIssuerKeysMissing = errors.New("issuer keys missing")
+
+	tracer = otel.GetTracerProvider().Tracer(tracerName)
 )
 
 // Validator represents a JWT validator.
 type Validator interface {
 	// ValidateToken checks that the given token is valid (i.e., is well-formed with a valid
 	// signature and future expiry). On success, it returns a map of claims describing the subject.
 	ValidateToken(string) (string, map[string]any, error)
+
+	// HealthCheck returns nil when the service is healthy.
+	HealthCheck(ctx context.Context) error
 }
 
 type validator struct {
 	kf     jwt.Keyfunc
 	parser *jwt.Parser
+
+	keyStorage jwkset.Storage
 }
 
 // NewValidator creates a validator with the given configuration.
@@ -61,6 +80,8 @@ func NewValidator(config Config) (Validator, error) {
 	out := &validator{
 		kf:     kf.Keyfunc,
 		parser: parser,
+
+		keyStorage: storage,
 	}
 
 	return out, nil
@@ -81,3 +102,39 @@ func (v *validator) ValidateToken(tokenString string) (string, map[string]any, e
 
 	return sub, mapClaims, nil
 }
+
+// HealthCheck returns nil when the service is healthy.
+func (v *validator) HealthCheck(ctx context.Context) error {
+	ctx, span := tracer.Start(ctx, "HealthCheck")
+	defer span.End()
+
+	span.SetAttributes(attribute.String("healthcheck.outcome", "unhealthy"))
+
+	keys, err := v.keyStorage.KeyReadAll(ctx)
+	if err != nil {
+		span.SetStatus(codes.Error, err.Error())
+		span.RecordError(err)
+
+		return err
+	}
+
+	if len(keys) == 0 {
+		span.SetStatus(codes.Error, ErrIssuerKeysMissing.Error())
+		span.RecordError(ErrIssuerKeysMissing)
+
+		return ErrIssuerKeysMissing
+	}
+
+	for _, key := range keys {
+		if err := key.Validate(); err != nil {
+			span.SetStatus(codes.Error, err.Error())
+			span.RecordError(err)
+
+			return err
+		}
+	}
+
+	span.SetAttributes(attribute.String("healthcheck.outcome", "healthy"))
+
+	return nil
+}