<commit_before>package kms

import (
	"context"
	"crypto/sha512"
	"fmt"

	apiv1 "cloud.google.com/go/kms/apiv1"
	"github.com/grafeas/voucher/signer"
	kms_pb "google.golang.org/genproto/googleapis/cloud/kms/v1"
)

const (
	AlgoSHA512 = "SHA512"
	APIPath    = "//cloudkms.googleapis.com/v1"
)

type Key struct {
	Path string
	Algo string
}

// Signer is an AttestationSigner that uses Google's Cloud KMS to sign attestations
// Only supports SHA512 digests.
type Signer struct {
	keys   map[string]Key
	client *apiv1.KeyManagementClient
}

func NewSigner(keys map[string]Key) (*Signer, error) {
	client, err := apiv1.NewKeyManagementClient(context.Background())
	if err != nil {
		return nil, err
	}

	return &Signer{
		keys:   keys,
		client: client,
	}, nil
}

func (s *Signer) Sign(checkName, body string) (string, string, error) {
	key, ok := s.keys[checkName]
	if !ok {
		return "", "", signer.ErrNoKeyForCheck
	}

	if key.Algo != AlgoSHA512 {
		return "", "", fmt.Errorf("unable to hash algorithm %q, must be SHA512", key.Algo)
	}

	resp, err := s.client.AsymmetricSign(context.Background(), &kms_pb.AsymmetricSignRequest{
		Name: key.Path,
		Digest: &kms_pb.Digest{
			Digest: &kms_pb.Digest_Sha512{Sha512: sha512Digest([]byte(body))},
		},
	})

	if err != nil {
		return "", "", err
	}

	return string(resp.Signature), fmt.Sprintf(APIPath+"/%v", key.Path), nil
}

// Close closes the KMS signer's connections.
func (s *Signer) Close() error {
	return s.client.Close()
}

func sha512Digest(data []byte) []byte {
	output := make([]byte, sha512.Size)
	sha := sha512.Sum512(data)
	copy(output, sha[:])
	return output
}<commit_msg>Add sha256 and sha384 support for kms.<commit_after>package kms

import (
	"context"
	"crypto/sha256"
	"crypto/sha512"
	"fmt"
	"hash"

	apiv1 "cloud.google.com/go/kms/apiv1"
	"github.com/grafeas/voucher/signer"
	kms_pb "google.golang.org/genproto/googleapis/cloud/kms/v1"
)

const (
	AlgoSHA256 = "SHA256"
	AlgoSHA384 = "SHA384"
	AlgoSHA512 = "SHA512"
	APIPath    = "//cloudkms.googleapis.com/v1"
)

type Key struct {
	Path string
	Algo string
}

// Signer is an AttestationSigner that uses Google's Cloud KMS to sign attestations
// Only supports SHA512 digests.
type Signer struct {
	keys   map[string]Key
	client *apiv1.KeyManagementClient
}

func NewSigner(keys map[string]Key) (*Signer, error) {
	client, err := apiv1.NewKeyManagementClient(context.Background())
	if err != nil {
		return nil, err
	}

	return &Signer{
		keys:   keys,
		client: client,
	}, nil
}

func (s *Signer) Sign(checkName, body string) (string, string, error) {
	key, ok := s.keys[checkName]
	if !ok {
		return "", "", signer.ErrNoKeyForCheck
	}

	var digest hash.Hash
	switch key.Algo {
	case AlgoSHA256:
		digest = sha256.New()
	case AlgoSHA384:
		digest = sha512.New384()
	case AlgoSHA512:
		digest = sha512.New()
	default:
		return "", "", fmt.Errorf("Unsupported digest algorithm %v", key.Algo)
	}

	if _, err := digest.Write([]byte(body)); err != nil {
		return "", "", err
	}

	var d kms_pb.Digest
	switch key.Algo {
	case AlgoSHA256:
		d.Digest = &kms_pb.Digest_Sha256{
			Sha256: digest.Sum(nil),
		}
	case AlgoSHA384:
		d.Digest = &kms_pb.Digest_Sha384{
			Sha384: digest.Sum(nil),
		}
	case AlgoSHA512:
		d.Digest = &kms_pb.Digest_Sha512{
			Sha512: digest.Sum(nil),
		}
	default:
		return "", "", fmt.Errorf("Unsupported digest algorithm %v", key.Algo)
	}

	resp, err := s.client.AsymmetricSign(context.Background(), &kms_pb.AsymmetricSignRequest{
		Name: key.Path,
		Digest: &d,
	})

	if err != nil {
		return "", "", err
	}

	return string(resp.Signature), fmt.Sprintf(APIPath+"/%v", key.Path), nil
}

// Close closes the KMS signer's connections.
func (s *Signer) Close() error {
	return s.client.Close()
}
{'fraction_non_alphanumeric': 0.14710610932475884, 'fraction_numerical': 0.03536977491961415, 'mean_word_length': 3.9906417112299466, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5483739', 'n_tokens_mistral': 1597, 'n_tokens_neox': 1446, 'n_words': 330}