server_test.go

/*
Copyright (c) 2018 Vereign AG [https://www.vereign.com]

This is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

package server

import (
	"log"
	"os"
	"testing"
	"time"

	dataStorageServer "code.vereign.com/code/data-storage-agent/server"
	"code.vereign.com/code/data-storage-agent/utils"
	"code.vereign.com/code/viam-apis/authentication"
	ksapi "code.vereign.com/code/viam-apis/key-storage-agent/api"
	ksclient "code.vereign.com/code/viam-apis/key-storage-agent/client"
	timestamp "github.com/golang/protobuf/ptypes/timestamp"
)

const (
	dataStorageGrpcAddress    = "localhost:7777"
	keyStorageGrpcAddress     = "localhost:7877"
	certFilePath              = "../cert/server.crt"
	privateKeyFilePath        = "../cert/server.key"
	vereignCertFilePath       = "../cert/vereign_ca.cer"
	vereignPrivateKeyFilePath = "../cert/vereign_ca.key"
)

func TestSetAndGetKeys(t *testing.T) {
	dataStorageClient := utils.CreateClientFromUuidAndSession("viam-system", "viam-session", dataStorageGrpcAddress, certFilePath)

	keyStorageAuth := &authentication.Authentication{
		Uuid:    "some-uuid",
		Session: "some-session",
	}

	_, _, _ = dataStorageClient.RenewSession(keyStorageAuth.Uuid, keyStorageAuth.Session)

	keyStorageClient := &ksclient.KeyStorageClientImpl{}
	keyStorageClient.SetUpClient(keyStorageAuth, keyStorageGrpcAddress, certFilePath)
	defer keyStorageClient.CloseClient()

	uuid, statusList, _ := keyStorageClient.DoReserveKeyUUID()
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoReserveKeyUUID, returned error: %s.", status.Code+":"+status.Description)
		}
	}

	privateKey := &ksapi.Key{
		Content: []byte{0, 1, 4},
	}

	statusList, _ = keyStorageClient.DoSetKey(uuid, ksapi.KeyType_PRIVATE, privateKey)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoSetKey, returned error: %s.", status.Code+":"+status.Description)
			return
		}
	}

	privateKeyResult, statusList, _ := keyStorageClient.DoGetKey(uuid, ksapi.KeyType_PRIVATE)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoGetKey, returned error: %s.", status.Code+":"+status.Description)
			return
		}
	}
	if privateKeyResult.Content == nil || !utils.CompareByteArrays(privateKeyResult.Content, []byte{0, 1, 4}) {
		t.Errorf("DoGetKey, incorrect keyResult.Content, expected: %v, but was: %v",
			[]byte{0, 1, 2}, privateKeyResult.Content)
	}

	// Test setting the same key twice
	statusList, _ = keyStorageClient.DoSetKey(uuid, ksapi.KeyType_PRIVATE, privateKey)
	for _, status := range statusList {
		if status.StatusType != ksapi.StatusType_ERROR {
			t.Errorf("DoSetKey, expected error, but got success")
			return
		}
	}

	publicKey := &ksapi.Key{
		Content: []byte{3, 4, 2},
	}

	statusList, _ = keyStorageClient.DoSetKey(uuid, ksapi.KeyType_PUBLIC, publicKey)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoSetKey, returned error: %s.", status.Code+":"+status.Description)
		}
	}

	publicKeyResult, statusList, _ := keyStorageClient.DoGetKey(uuid, ksapi.KeyType_PUBLIC)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoGetKey, returned error: %s.", status.Code+":"+status.Description)
		}
	}
	if publicKeyResult.Content == nil || !utils.CompareByteArrays(publicKeyResult.Content, []byte{3, 4, 2}) {
		t.Errorf("DoGetKey, incorrect publicKeyResult.Content, expected: %v, but was: %v",
			[]byte{3, 4, 2}, publicKeyResult.Content)
	}

	statusList, _ = keyStorageClient.DoRevoke(uuid)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoRevoke, returned error: %s.", status.Code+":"+status.Description)
		}
	}
	revokedPrivateKeyResult, statusList, _ := keyStorageClient.DoGetKey(uuid, ksapi.KeyType_PRIVATE)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoGetKey, returned error: %s.", status.Code+":"+status.Description)
		}
	}
	if revokedPrivateKeyResult.Revoked == false {
		t.Errorf("DoRevoke, key was not revoked")
	}
	revokedPublicKeyResult, statusList, _ := keyStorageClient.DoGetKey(uuid, ksapi.KeyType_PUBLIC)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoGetKey, returned error: %s.", status.Code+":"+status.Description)
		}
	}
	if revokedPublicKeyResult.Revoked == false {
		t.Errorf("DoRevoke, key was not revoked")
	}
}
func TestGenerateKeyPairAndCertificate(t *testing.T) {
	dataStorageClient := utils.CreateClientFromUuidAndSession("viam-system", "viam-session", dataStorageGrpcAddress, certFilePath)

	keyStorageAuth := &authentication.Authentication{
		Uuid:    "some-uuid",
		Session: "some-session",
	}

	_, _, _ = dataStorageClient.RenewSession(keyStorageAuth.Uuid, keyStorageAuth.Session)

	keyStorageClient := &ksclient.KeyStorageClientImpl{}
	keyStorageClient.SetUpClient(keyStorageAuth, keyStorageGrpcAddress, certFilePath)
	defer keyStorageClient.CloseClient()

	uuid, statusList, _ := keyStorageClient.DoGenerateKeyPair(2048)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoGenerateKeyPair, returned error: %s.", status.Code+":"+status.Description)
			return
		}
	}
	if uuid == "" {
		t.Errorf("DoGenerateKeyPair, uuid is empty")
		return
	}

	privateKeyResult, statusList, _ := keyStorageClient.DoGetKey(uuid, ksapi.KeyType_PRIVATE)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoGetKey, returned error: %s.", status.Code+":"+status.Description)
			return
		}
	}
	if privateKeyResult.Content == nil {
		t.Errorf("DoGetKey, privateKeyResult.Content is nil")
	}

	publicKeyResult, statusList, _ := keyStorageClient.DoGetKey(uuid, ksapi.KeyType_PUBLIC)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoGetKey, returned error: %s.", status.Code+":"+status.Description)
			return
		}
	}
	if publicKeyResult.Content == nil {
		t.Errorf("DoGetKey, publicKeyResult.Content is nil")
	}

	nowTime := time.Now()
	nowTimeSeconds := time.Now().Unix()
	nowTimeNanoSeconds := int32(nowTime.Sub(time.Unix(nowTimeSeconds, 0)))
	notBefore := &timestamp.Timestamp{
		Seconds: nowTimeSeconds,
		Nanos:   nowTimeNanoSeconds,
	}
	timeAfterOneDay := nowTime.Add(1000000 * 3600 * 24)
	secondsAfterOneDay := time.Now().Unix()
	nanoSecondsAfterOneDay := int32(timeAfterOneDay.Sub(time.Unix(secondsAfterOneDay, 0)))
	notAfter := &timestamp.Timestamp{
		Seconds: secondsAfterOneDay,
		Nanos:   nanoSecondsAfterOneDay,
	}

	certificateData := &ksapi.GenerateCertificateRequest_CertificateData{
		Country:            "BG",
		Organization:       "company",
		OrganizationalUnit: "DEV",
		CommonName:         "CN",
		NotBefore:          notBefore,
		NotAfter:           notAfter,
		Host:               "abcde.com",
	}

	statusList, _ = keyStorageClient.DoGenerateCertificate(uuid, certificateData)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoGenerateCertificate, returned error: %s.", status.Code+":"+status.Description)
		}
	}

	certificate, statusList, _ := keyStorageClient.DoGetKey(uuid, ksapi.KeyType_CERTIFICATE)
	for _, status := range statusList {
		if status.StatusType == ksapi.StatusType_ERROR {
			t.Errorf("DoGetKey, returned error: %s.", status.Code+":"+status.Description)
		}
	}
	if certificate == nil || certificate.Content == nil {
		t.Errorf("DoGetKey, certificate is nil")
	}
}

func TestMain(m *testing.M) {
	// fire the gRPC server in a goroutine
	go func() {
		err := dataStorageServer.StartGRPCServer(dataStorageGrpcAddress, certFilePath, privateKeyFilePath)
		if err != nil {
			log.Fatalf("failed to start gRPC server: %s", err)
		}
	}()

	// wait a second for the server to start
	time.Sleep(time.Duration(1) * time.Second)

	// fire the gRPC server in a goroutine
	go func() {
		err := StartGRPCServer(keyStorageGrpcAddress, certFilePath, privateKeyFilePath,
			vereignCertFilePath, vereignPrivateKeyFilePath, dataStorageGrpcAddress)
		if err != nil {
			log.Fatalf("failed to start gRPC server: %s", err)
		}
	}()

	// wait a second for the server to start
	time.Sleep(time.Duration(1) * time.Second)

	retCode := m.Run()
	os.Exit(retCode)
}