Merge pull request #364 from skynet2/encrypt

feat: add encrypt & decrypt for aws

go.mod

@@ -8,6 +8,7 @@ go 1.18
 
 require (
 	github.com/aws/aws-sdk-go v1.43.9
+	github.com/aws/aws-sdk-go-v2 v1.17.3
 	github.com/aws/aws-sdk-go-v2/service/kms v1.20.0
 	github.com/btcsuite/btcd v0.22.1
 	github.com/golang/mock v1.6.0
@@ -29,7 +30,6 @@ require (
 
 require (
 	github.com/VictoriaMetrics/fastcache v1.5.7 // indirect
-	github.com/aws/aws-sdk-go-v2 v1.17.3 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/configsources v1.1.27 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.4.21 // indirect
 	github.com/aws/smithy-go v1.13.5 // indirect

pkg/aws/opts.go

@@ -12,6 +12,7 @@ import (
 
 type opts struct {
 	keyAliasPrefix string
+	awsClient      awsClient
 }
 
 // NewOpts create new opts populated with environment variable.
@@ -34,3 +35,8 @@ type Opts func(opts *opts)
 func WithKeyAliasPrefix(prefix string) Opts {
 	return func(opts *opts) { opts.keyAliasPrefix = prefix }
 }
+
+// WithAWSClient sets custom AWS client.
+func WithAWSClient(client awsClient) Opts {
+	return func(opts *opts) { opts.awsClient = client }
+}

pkg/aws/service.go

@@ -4,11 +4,14 @@ Copyright SecureKey Technologies Inc. All Rights Reserved.
 SPDX-License-Identifier: Apache-2.0
 */
 
+//go:generate mockgen -destination service_mocks.go -package aws -source=service.go
+
 package aws
 
 import (
 	"context"
 	"crypto/elliptic"
+	"crypto/rand"
 	"crypto/sha256"
 	"crypto/sha512"
 	"encoding/asn1"
@@ -26,7 +29,7 @@ import (
 	arieskms "github.com/hyperledger/aries-framework-go/pkg/kms"
 )
 
-type awsClient interface { //nolint:dupl
+type awsClient interface {
 	Sign(ctx context.Context, params *kms.SignInput, optFns ...func(*kms.Options)) (*kms.SignOutput, error)
 	GetPublicKey(ctx context.Context, params *kms.GetPublicKeyInput,
 		optFns ...func(*kms.Options)) (*kms.GetPublicKeyOutput, error)
@@ -37,11 +40,17 @@ type awsClient interface { //nolint:dupl
 		optFns ...func(*kms.Options)) (*kms.CreateKeyOutput, error)
 	CreateAlias(ctx context.Context, params *kms.CreateAliasInput,
 		optFns ...func(*kms.Options)) (*kms.CreateAliasOutput, error)
+	Encrypt(ctx context.Context, params *kms.EncryptInput, optFns ...func(*kms.Options)) (*kms.EncryptOutput, error)
+	Decrypt(ctx context.Context, params *kms.DecryptInput, optFns ...func(*kms.Options)) (*kms.DecryptOutput, error)
 }
 
 type metricsProvider interface {
 	SignCount()
+	EncryptCount()
+	DecryptCount()
 	SignTime(value time.Duration)
+	EncryptTime(value time.Duration)
+	DecryptTime(value time.Duration)
 	ExportPublicKeyCount()
 	ExportPublicKeyTime(value time.Duration)
 	VerifyCount()
@@ -58,6 +67,8 @@ type Service struct {
 	client           awsClient
 	metrics          metricsProvider
 	healthCheckKeyID string
+	encryptionAlgo   types.EncryptionAlgorithmSpec
+	nonceLength      int
 }
 
 const (
@@ -79,23 +90,108 @@ var keySpecToCurve = map[types.KeySpec]elliptic.Curve{
 	types.KeySpecEccSecgP256k1: btcec.S256(),
 }
 
+const (
+	defaultNonceLength = 16
+)
+
 // New return aws service.
-func New(awsConfig *aws.Config, metrics metricsProvider,
-	healthCheckKeyID string, opts ...Opts) *Service {
+func New(
+	awsConfig *aws.Config,
+	metrics metricsProvider,
+	healthCheckKeyID string,
+	opts ...Opts,
+) *Service {
 	options := newOpts()
 
 	for _, opt := range opts {
 		opt(options)
 	}
 
+	client := options.awsClient
+	if client == nil {
+		client = kms.NewFromConfig(*awsConfig)
+	}
+
 	return &Service{
 		options:          options,
-		client:           kms.NewFromConfig(*awsConfig),
+		client:           client,
 		metrics:          metrics,
 		healthCheckKeyID: healthCheckKeyID,
+		encryptionAlgo:   types.EncryptionAlgorithmSpecRsaesOaepSha256,
+		nonceLength:      defaultNonceLength,
 	}
 }
 
+// Decrypt data.
+func (s *Service) Decrypt(_, aad, _ []byte, kh interface{}) ([]byte, error) {
+	startTime := time.Now()
+
+	defer func() {
+		if s.metrics != nil {
+			s.metrics.DecryptTime(time.Since(startTime))
+		}
+	}()
+
+	if s.metrics != nil {
+		s.metrics.DecryptCount()
+	}
+
+	keyID, err := s.getKeyID(kh.(string))
+	if err != nil {
+		return nil, err
+	}
+
+	input := &kms.DecryptInput{
+		CiphertextBlob:      aad,
+		EncryptionAlgorithm: s.encryptionAlgo,
+		KeyId:               aws.String(keyID),
+	}
+
+	resp, err := s.client.Decrypt(context.Background(), input)
+	if err != nil {
+		return nil, err
+	}
+
+	return resp.Plaintext, nil
+}
+
+// Encrypt data.
+func (s *Service) Encrypt(
+	msg []byte,
+	_ []byte,
+	kh interface{},
+) ([]byte, []byte, error) {
+	startTime := time.Now()
+
+	defer func() {
+		if s.metrics != nil {
+			s.metrics.EncryptTime(time.Since(startTime))
+		}
+	}()
+
+	if s.metrics != nil {
+		s.metrics.EncryptCount()
+	}
+
+	keyID, err := s.getKeyID(kh.(string))
+	if err != nil {
+		return nil, nil, err
+	}
+
+	input := &kms.EncryptInput{
+		KeyId:               aws.String(keyID),
+		Plaintext:           msg,
+		EncryptionAlgorithm: s.encryptionAlgo,
+	}
+
+	resp, err := s.client.Encrypt(context.Background(), input)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return resp.CiphertextBlob, generateNonce(s.nonceLength), nil
+}
+
 // Sign data.
 func (s *Service) Sign(msg []byte, kh interface{}) ([]byte, error) { //nolint: funlen
 	startTime := time.Now()
@@ -310,6 +406,13 @@ func (s *Service) getKeyID(keyURI string) (string, error) {
 	return r[4], nil
 }
 
+func generateNonce(length int) []byte {
+	key := make([]byte, length)
+	_, _ = rand.Read(key) //nolint: errcheck
+
+	return key
+}
+
 func hashMessage(message []byte, algorithm types.SigningAlgorithmSpec) ([]byte, error) {
 	var digest hash.Hash