added bearertoken extension encryption

extension/bearertokenauthextension/bearertokenauth.go

@@ -5,7 +5,11 @@
 
 import (
 	"context"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/md5"
 	"crypto/subtle"
+	"encoding/hex"
 	"errors"
 	"fmt"
 	"net/http"
@@ -53,6 +57,9 @@
 
 	shutdownCH chan struct{}
 
+	encryptionKey  string
+	encryptionType string
+
 	filename string
 	logger   *zap.Logger
 }
@@ -144,6 +151,35 @@
 	}
 }
 
+func (b *bearerTokenAuth) aesDecrypt(ciphered []byte, key string) ([]byte, error) {
+
+	byteinput := []byte(key)
+	md5Hash := md5.Sum(byteinput)
+	hashedKey := hex.EncodeToString(md5Hash[:])
+
+	aesBlock, err := aes.NewCipher([]byte(hashedKey))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create AES cipher: %w", err)
+	}
+
+	gcmInstance, err := cipher.NewGCM(aesBlock)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create GCM instance: %w", err)
+	}
+
+	nonceSize := gcmInstance.NonceSize()
+	//fmt.Printf("ciphered length: %d, nonceSize: %d\n", len(ciphered), nonceSize)
+	if len(ciphered) < nonceSize {
+		return nil, fmt.Errorf("ciphered data is too short: expected at least %d bytes, got %d", nonceSize, len(ciphered))
+	}
+	nonce, ciphered := ciphered[:nonceSize], ciphered[nonceSize:]
+	decodedBytes, err := gcmInstance.Open(nil, nonce, ciphered, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decrypt data: %w", err)
+	}
+	return decodedBytes, nil
+}
+
 // Reloads token from file
 func (b *bearerTokenAuth) refreshToken() {
 	b.logger.Info("refresh token", zap.String("filename", b.filename))
@@ -153,6 +189,15 @@
 		return
 	}
 
+	// If AES encryption is specified, decrypt the token
+	if b.encryptionType == "AES" {
+		tokenData, err = b.aesDecrypt(tokenData, b.encryptionKey)
+		if err != nil {
+			b.logger.Error(err.Error())
+			return
+		}
+	}
+
 	tokens := strings.Split(string(tokenData), "\n")
 	for i, token := range tokens {
 		tokens[i] = strings.TrimSpace(token)

extension/bearertokenauthextension/bearertokenauth_test.go

@@ -4,7 +4,13 @@
 package bearertokenauthextension
 
 import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/md5"
+	"crypto/rand"
+	"encoding/hex"
 	"fmt"
+	"io"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -446,3 +452,69 @@
 
 	assert.NoError(t, bauth.Shutdown(t.Context()))
 }
+
+func encryptAES(plaintext []byte, key string) (string, error) {
+	// Generate MD5 hash of the key
+	byteinput := []byte(key)
+	md5Hash := md5.Sum(byteinput)
+	hashedKey := hex.EncodeToString(md5Hash[:])
+
+	// Create AES cipher block
+	aesBlock, err := aes.NewCipher([]byte(hashedKey))
+	if err != nil {
+		return "", err
+	}
+
+	// Create GCM instance
+	gcm, err := cipher.NewGCM(aesBlock)
+	if err != nil {
+		return "", err
+	}
+
+	// Generate nonce
+	nonce := make([]byte, gcm.NonceSize())
+	_, _ = io.ReadFull(rand.Reader, nonce)
+	ciphertext := gcm.Seal(nonce, nonce, plaintext, nil)
+	return string(ciphertext), nil
+}
+
+func TestAESDecrypt(t *testing.T) {
+	cfg := createDefaultConfig().(*Config)
+	cfg.EncryptionType = "AES"
+	cfg.EncryptionKey = "encryptionkey"
+
+	encryptedText, err := encryptAES([]byte("testdata"), cfg.EncryptionKey)
+	assert.NoError(t, err)
+
+	invalidEncryptedText, err := encryptAES([]byte("testdata"), "invalidencryptionkey")
+	assert.NoError(t, err)
+
+	bauth := newBearerTokenAuth(cfg, nil)
+	assert.NotNil(t, bauth)
+	// Test cases
+	tests := []struct {
+		name             string
+		plaintext        string
+		preEncryptedText string
+		key              string
+		expectedError    bool
+	}{
+		{"Valid Decryption", "testdata", encryptedText, cfg.EncryptionKey, false},
+		{"Invalid Decryption", "testdata", invalidEncryptedText, cfg.EncryptionKey, true},
+		{"Invalid Key", "testdata", encryptedText, "invalidencryptionkey", true},
+		{"Empty Ciphertext", "testdata", "", cfg.EncryptionKey, true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+
+			decryptedText, err := bauth.aesDecrypt([]byte(tt.preEncryptedText), tt.key)
+			if tt.expectedError {
+				assert.Error(t, err)
+				return
+			}
+			assert.NoError(t, err)
+			assert.Equal(t, tt.plaintext, string(decryptedText))
+		})
+	}
+}

extension/bearertokenauthextension/config.go

@@ -27,14 +27,22 @@ type Config struct {
 	// Filename points to a file that contains the bearer token(s) to use for every RPC.
 	Filename string `mapstructure:"filename,omitempty"`
 
+	// EncryptionType specifies the encryption type used to encrypt the bearer token(s) in the file.
+	EncryptionType string `mapstructure:"encryption_type,omitempty"`
+
+	// EncryptionKey specifies the encryption key used to encrypt the bearer token(s) in the file.
+	EncryptionKey string `mapstructure:"encryption_key,omitempty"`
+
 	// prevent unkeyed literal initialization
 	_ struct{}
 }
 
 var (
-	_                         component.Config = (*Config)(nil)
-	errNoTokenProvided                         = errors.New("no bearer token provided")
-	errTokensAndTokenProvided                  = errors.New("either tokens or token should be provided, not both")
+	_                            component.Config = (*Config)(nil)
+	errNoTokenProvided                            = errors.New("no bearer token provided")
+	errTokensAndTokenProvided                     = errors.New("either tokens or token should be provided, not both")
+	errEncyptionTypeNotSupported                  = errors.New("encryption type not supported")
+	errInvalidAESEncryptionKey                    = errors.New("invalid AES encryption key")
 )
 
 // Validate checks if the extension configuration is valid
@@ -45,5 +53,11 @@ func (cfg *Config) Validate() error {
 	if cfg.BearerToken != "" && len(cfg.Tokens) > 0 {
 		return errTokensAndTokenProvided
 	}
+	if cfg.EncryptionType == "AES" && cfg.EncryptionKey == "" {
+		return errInvalidAESEncryptionKey
+	}
+	if cfg.EncryptionType != "AES" && cfg.EncryptionType != "" {
+		return errEncyptionTypeNotSupported
+	}
 	return nil
 }