cocos/pkg/clients/grpc/connect_test.go

// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"math/big"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/absmach/supermq/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/ultravioletrs/cocos/pkg/clients"
	"github.com/ultravioletrs/cocos/pkg/tls"
)

func TestNewClient(t *testing.T) {
	caCertFile, clientCertFile, clientKeyFile, err := createCertificatesFiles()
	require.NoError(t, err)

	policyFile, err := os.CreateTemp("", "attestation_policy.json")
	require.NoError(t, err)
	_, err = policyFile.WriteString("{}")
	require.NoError(t, err)
	err = policyFile.Close()
	require.NoError(t, err)

	t.Cleanup(func() {
		os.Remove(caCertFile)
		os.Remove(clientCertFile)
		os.Remove(clientKeyFile)
		os.Remove(policyFile.Name())
	})

	tests := []struct {
		name     string
		cfg      clients.StandardClientConfig
		agentCfg clients.AttestedClientConfig
		wantErr  bool
		err      error
	}{
		{
			name: "Success without TLS",
			cfg: clients.StandardClientConfig{
				URL: "localhost:7001",
			},
			wantErr: false,
			err:     nil,
		},
		{
			name: "Success with TLS",
			cfg: clients.StandardClientConfig{
				URL:          "localhost:7001",
				ServerCAFile: caCertFile,
			},
			wantErr: false,
			err:     nil,
		},
		{
			name: "Success with mTLS",
			cfg: clients.StandardClientConfig{
				URL:          "localhost:7001",
				ServerCAFile: caCertFile,
				ClientCert:   clientCertFile,
				ClientKey:    clientKeyFile,
			},
			wantErr: false,
			err:     nil,
		},
		{
			name: "Success agent client with mTLS",
			agentCfg: clients.AttestedClientConfig{
				StandardClientConfig: clients.StandardClientConfig{
					URL:          "localhost:7001",
					ServerCAFile: caCertFile,
					ClientCert:   clientCertFile,
					ClientKey:    clientKeyFile,
				},
			},
			wantErr: false,
			err:     nil,
		},
		{
			name: "Success agent client with aTLS",
			agentCfg: clients.AttestedClientConfig{
				StandardClientConfig: clients.StandardClientConfig{
					URL:          "localhost:7001",
					ServerCAFile: caCertFile,
					ClientCert:   clientCertFile,
					ClientKey:    clientKeyFile,
				},
				AttestedTLS:       true,
				AttestationPolicy: policyFile.Name(),
			},
			wantErr: false,
			err:     nil,
		},
		{
			name: "Failed agent client with aTLS",
			agentCfg: clients.AttestedClientConfig{
				StandardClientConfig: clients.StandardClientConfig{
					URL:          "localhost:7001",
					ServerCAFile: caCertFile,
					ClientCert:   clientCertFile,
					ClientKey:    clientKeyFile,
				},
				AttestedTLS:       true,
				AttestationPolicy: "no such file",
			},
			wantErr: true,
			err:     fmt.Errorf("failed to stat attestation policy file"),
		},
		{
			name: "Fail with invalid ServerCAFile",
			cfg: clients.StandardClientConfig{
				URL:          "localhost:7001",
				ServerCAFile: "nonexistent.pem",
			},
			wantErr: true,
			err:     tls.ErrFailedToLoadRootCA,
		},
		{
			name: "Fail with invalid ClientCert",
			cfg: clients.StandardClientConfig{
				URL:          "localhost:7001",
				ServerCAFile: caCertFile,
				ClientCert:   "nonexistent.pem",
				ClientKey:    clientKeyFile,
			},
			wantErr: true,
			err:     tls.ErrFailedToLoadClientCertKey,
		},
		{
			name: "Fail with invalid ClientKey",
			cfg: clients.StandardClientConfig{
				URL:          "localhost:7001",
				ServerCAFile: caCertFile,
				ClientCert:   clientCertFile,
				ClientKey:    "nonexistent.pem",
			},
			wantErr: true,
			err:     tls.ErrFailedToLoadClientCertKey,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var client Client
			if strings.Contains(tt.name, "agent client") {
				client, err = NewClient(tt.agentCfg)
			} else {
				client, err = NewClient(tt.cfg)
			}
			assert.True(t, errors.Contains(err, tt.err), fmt.Sprintf("expected error %v, got %v", tt.err, err))
			if tt.wantErr {
				assert.Error(t, err)
				assert.Nil(t, client)
			} else {
				assert.NoError(t, err)
				assert.NotNil(t, client)
				assert.NoError(t, client.Close())
			}
		})
	}
}

func TestClientSecure(t *testing.T) {
	tests := []struct {
		name     string
		secure   tls.Security
		expected string
	}{
		{
			name:     "Without TLS",
			secure:   tls.WithoutTLS,
			expected: "without TLS",
		},
		{
			name:     "With TLS",
			secure:   tls.WithTLS,
			expected: "with TLS",
		},
		{
			name:     "With mTLS",
			secure:   tls.WithMTLS,
			expected: "with mTLS",
		},
		{
			name:     "With aTLS",
			secure:   tls.WithATLS,
			expected: "with aTLS",
		},
		{
			name:     "With maTLS",
			secure:   tls.WithMATLS,
			expected: "with maTLS",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c := &client{security: tt.secure}
			assert.Equal(t, tt.expected, c.Secure())
		})
	}
}

func createCertificatesFiles() (string, string, string, error) {
	caKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return "", "", "", err
	}

	caTemplate := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			Organization: []string{"Test Org"},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour * 24),
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
		BasicConstraintsValid: true,
		IsCA:                  true,
	}

	caCertDER, err := x509.CreateCertificate(rand.Reader, &caTemplate, &caTemplate, &caKey.PublicKey, caKey)
	if err != nil {
		return "", "", "", err
	}

	caCertFile, err := createTempFile(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: caCertDER}))
	if err != nil {
		return "", "", "", err
	}

	clientKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return "", "", "", err
	}

	clientTemplate := x509.Certificate{
		SerialNumber: big.NewInt(2),
		Subject: pkix.Name{
			Organization: []string{"Test Org"},
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().Add(time.Hour * 24),
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
	}

	clientCertDER, err := x509.CreateCertificate(rand.Reader, &clientTemplate, &caTemplate, &clientKey.PublicKey, caKey)
	if err != nil {
		return "", "", "", err
	}

	clientCertFile, err := createTempFile(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: clientCertDER}))
	if err != nil {
		return "", "", "", err
	}

	clientKeyFile, err := createTempFile(pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(clientKey)}))
	if err != nil {
		return "", "", "", err
	}

	return caCertFile, clientCertFile, clientKeyFile, nil
}

func createTempFile(data []byte) (string, error) {
	file, err := createTempFileHandle()
	if err != nil {
		return "", err
	}

	_, err = file.Write(data)
	if err != nil {
		return "", err
	}

	err = file.Close()
	if err != nil {
		return "", err
	}

	return file.Name(), nil
}

func createTempFileHandle() (*os.File, error) {
	return os.CreateTemp("", "test")
}