feat: CAS readiness probe support (#55)

Signed-off-by: Miguel Martinez Trivino <[email protected]>

Author: migmartri

app/artifact-cas/internal/server/grpc.go

@@ -19,6 +19,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"regexp"
 
 	v1 "github.com/chainloop-dev/chainloop/app/artifact-cas/api/cas/v1"
 	"github.com/chainloop-dev/chainloop/app/artifact-cas/internal/conf"
@@ -27,6 +28,7 @@ import (
 	"github.com/getsentry/sentry-go"
 	"github.com/go-kratos/kratos/v2/errors"
 	jwtMiddleware "github.com/go-kratos/kratos/v2/middleware/auth/jwt"
+	"github.com/go-kratos/kratos/v2/middleware/selector"
 	jwt "github.com/golang-jwt/jwt/v4"
 	"google.golang.org/genproto/googleapis/bytestream"
 
@@ -62,11 +64,14 @@ func NewGRPCServer(c *conf.Server, authConf *conf.Auth, byteService *service.Byt
 			),
 			logging.Server(logger),
 			// NOTE: JWT middleware only works for unary requests
-			// below you can see a reimplementation of the middleware as a stream interceptor
-			jwtMiddleware.Server(
-				loadPublicKey(rawKey),
-				jwtMiddleware.WithSigningMethod(casJWT.SigningMethod),
-				jwtMiddleware.WithClaims(func() jwt.Claims { return &casJWT.Claims{} })),
+			// below you can see a re-implementation of the middleware as a stream interceptor
+			// If we require a logged in user we
+			selector.Server(
+				jwtMiddleware.Server(
+					loadPublicKey(rawKey),
+					jwtMiddleware.WithSigningMethod(casJWT.SigningMethod),
+					jwtMiddleware.WithClaims(func() jwt.Claims { return &casJWT.Claims{} })),
+			).Match(requireAuthentication()).Build(),
 			validate.Validator(),
 		),
 
@@ -92,13 +97,24 @@ func NewGRPCServer(c *conf.Server, authConf *conf.Auth, byteService *service.Byt
 
 	bytestream.RegisterByteStreamServer(srv.Server, byteService)
 	v1.RegisterResourceServiceServer(srv.Server, rSvc)
+	v1.RegisterStatusServiceServer(srv.Server, service.NewStatusService(Version))
 
 	// Register and set metrics to 0
 	grpc_prometheus.Register(srv.Server)
 
 	return srv, nil
 }
 
+func requireAuthentication() selector.MatchFunc {
+	// Skip authentication on the status grpc service
+	const skipRegexp = "(cas.v1.StatusService/.*)"
+
+	return func(ctx context.Context, operation string) bool {
+		r := regexp.MustCompile(skipRegexp)
+		return !r.MatchString(operation)
+	}
+}
+
 // load key for verification
 func loadPublicKey(rawKey []byte) jwt.Keyfunc {
 	return func(token *jwt.Token) (interface{}, error) {

app/artifact-cas/internal/server/grpc_test.go

@@ -138,6 +138,23 @@ func TestJWTAuthFunc(t *testing.T) {
 	}
 }
 
+func TestRequireAuthentication(t *testing.T) {
+	testCases := []struct {
+		operation string
+		matches   bool
+	}{
+		{"/cas.v1.Resource/List", true},
+		{"/cas.v1.Bytestream/List", true},
+		{"/cas.v1.StatusService/Infoz", false},
+		{"/cas.v1.StatusService/Statusz", false},
+	}
+
+	matchFunc := requireAuthentication()
+	for _, op := range testCases {
+		assert.Equal(t, matchFunc(context.Background(), op.operation), op.matches)
+	}
+}
+
 func loadTestPublicKey(path string) jwt.Keyfunc {
 	rawKey, _ := os.ReadFile(path)
 	return func(token *jwt.Token) (interface{}, error) {

app/controlplane/cmd/wire.go

@@ -47,6 +47,7 @@ func wireApp(*conf.Bootstrap, credentials.ReaderWriter, log.Logger) (*app, func(
 			wire.Bind(new(biz.CASClient), new(*biz.CASClientUseCase)),
 			oci.NewBackendProvider,
 			serviceOpts,
+			wire.Value([]biz.CASClientOpts{}),
 			wire.FieldsOf(new(*conf.Bootstrap), "Server", "Auth", "Data", "CasServer"),
 			newApp,
 		),

app/controlplane/cmd/wire_gen.go

@@ -17,6 +17,7 @@ package biz
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 
@@ -29,9 +30,13 @@ import (
 )
 
 type CASClientUseCase struct {
+	// to generate temporary credentials
 	credsProvider *CASCredentialsUseCase
+	// configuration to generate the client
 	casServerConf *conf.Bootstrap_CASServer
-	logger        *log.Helper
+	// factory to generate the client
+	casClientFactory CASClientFactory
+	logger           *log.Helper
 }
 
 type CASUploader interface {
@@ -45,11 +50,40 @@ type CASDownloader interface {
 type CASClient interface {
 	CASUploader
 	CASDownloader
-	Configured() bool
 }
 
-func NewCASClientUseCase(credsProvider *CASCredentialsUseCase, config *conf.Bootstrap_CASServer, l log.Logger) *CASClientUseCase {
-	return &CASClientUseCase{credsProvider, config, servicelogger.ScopedHelper(l, "biz/cas-client")}
+type CASClientFactory func(conf *conf.Bootstrap_CASServer, token string) (casclient.DownloaderUploader, error)
+type CASClientOpts func(u *CASClientUseCase)
+
+func WithClientFactory(f CASClientFactory) CASClientOpts {
+	return func(c *CASClientUseCase) {
+		c.casClientFactory = f
+	}
+}
+
+func NewCASClientUseCase(credsProvider *CASCredentialsUseCase, config *conf.Bootstrap_CASServer, l log.Logger, opts ...CASClientOpts) *CASClientUseCase {
+	// generate a client from the given configuration
+	defaultCasClientFactory := func(conf *conf.Bootstrap_CASServer, token string) (casclient.DownloaderUploader, error) {
+		conn, err := grpcconn.New(conf.GetGrpc().GetAddr(), token, conf.GetInsecure())
+		if err != nil {
+			return nil, fmt.Errorf("failed to create grpc connection: %w", err)
+		}
+
+		return casclient.New(conn), nil
+	}
+
+	uc := &CASClientUseCase{
+		credsProvider:    credsProvider,
+		casServerConf:    config,
+		logger:           servicelogger.ScopedHelper(l, "biz/cas-client"),
+		casClientFactory: defaultCasClientFactory,
+	}
+
+	for _, opt := range opts {
+		opt(uc)
+	}
+
+	return uc
 }
 
 // The secretID is embedded in the JWT token and is used to identify the secret by the CAS server
@@ -90,35 +124,31 @@ func (uc *CASClientUseCase) Download(ctx context.Context, secretID string, w io.
 }
 
 // create a client with a temporary set of credentials for a specific operation
-func (uc *CASClientUseCase) casAPIClient(secretID string, role casJWT.Role) (*casclient.Client, error) {
+func (uc *CASClientUseCase) casAPIClient(secretID string, role casJWT.Role) (casclient.DownloaderUploader, error) {
 	token, err := uc.credsProvider.GenerateTemporaryCredentials(secretID, role)
 	if err != nil {
 		return nil, fmt.Errorf("failed to generate temporary credentials: %w", err)
 	}
 
 	// Initialize connection to CAS server
-	return casClient(uc.casServerConf, token)
+	return uc.casClientFactory(uc.casServerConf, token)
 }
 
-func casClient(conf *conf.Bootstrap_CASServer, token string) (*casclient.Client, error) {
-	conn, err := grpcconn.New(conf.GetGrpc().GetAddr(), token, conf.GetInsecure())
-	if err != nil {
-		return nil, fmt.Errorf("failed to create grpc connection: %w", err)
-	}
-
-	return casclient.New(conn), nil
-}
-
-// If the CAS client configuration is present and valid
-func (uc *CASClientUseCase) Configured() bool {
+// If the CAS server can be reached and reports readiness
+func (uc *CASClientUseCase) IsReady(ctx context.Context) (bool, error) {
 	if uc.casServerConf == nil {
-		return false
+		return false, errors.New("missing CAS server configuration")
 	}
 
 	err := uc.casServerConf.ValidateAll()
 	if err != nil {
-		uc.logger.Infow("msg", "Invalid CAS client configuration", "err", err.Error())
+		return false, fmt.Errorf("invalid CAS client configuration: %w", err)
+	}
+
+	c, err := uc.casClientFactory(uc.casServerConf, "")
+	if err != nil {
+		return false, fmt.Errorf("failed to create CAS client: %w", err)
 	}
 
-	return err == nil
+	return c.IsReady(ctx)
 }

app/controlplane/internal/biz/casclient.go

@@ -0,0 +1,84 @@
+//
+// Copyright 2023 The Chainloop Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package biz_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/chainloop-dev/chainloop/app/controlplane/internal/biz"
+	"github.com/chainloop-dev/chainloop/app/controlplane/internal/conf"
+	"github.com/chainloop-dev/chainloop/internal/casclient"
+	"github.com/chainloop-dev/chainloop/internal/casclient/mocks"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+)
+
+func TestIsReady(t *testing.T) {
+	validConf := &conf.Bootstrap_CASServer{
+		Grpc: &conf.Server_GRPC{Addr: "localhost:1111"},
+	}
+
+	testCases := []struct {
+		name     string
+		config   *conf.Bootstrap_CASServer
+		casReady bool
+		want     bool
+		wantErr  bool
+	}{
+		{
+			name:    "missing configuration",
+			config:  &conf.Bootstrap_CASServer{},
+			wantErr: true,
+		},
+		{
+			name:    "invalid configuration",
+			config:  &conf.Bootstrap_CASServer{Grpc: &conf.Server_GRPC{}},
+			wantErr: true,
+		},
+		{
+			name:    "not ready configuration",
+			config:  validConf,
+			wantErr: false,
+		},
+		{
+			name:     "ready configuration",
+			config:   validConf,
+			casReady: true,
+			want:     true,
+			wantErr:  false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			clientProvider := func(conf *conf.Bootstrap_CASServer, token string) (casclient.DownloaderUploader, error) {
+				c := mocks.NewDownloaderUploader(t)
+				c.On("IsReady", mock.Anything).Return(tc.casReady, nil)
+				return c, nil
+			}
+			uc := biz.NewCASClientUseCase(nil, tc.config, nil, biz.WithClientFactory(clientProvider))
+
+			got, err := uc.IsReady(context.Background())
+			if tc.wantErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}