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"
)

const (
	dataStorageGrpcAddress = "localhost:7777"
	keyStorageGrpcAddress  = "localhost:7877"
	certFile               = "../cert/server.crt"
	keyFile                = "../cert/server.key"
)

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

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

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

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

	privateKey := &ksapi.PrivateKey{
		Name:               "privkey1",
		PublicKeyReference: "pubkey1",
		Content:            []byte{0, 1, 4},
	}

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

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

	publicKey := &ksapi.PublicKey{
		Name:        "pubkey1",
		Content:     []byte{3, 4, 2},
		Certificate: []byte{5, 1, 6},
	}

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

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

func TestMain(m *testing.M) {
	// fire the gRPC server in a goroutine
	go func() {
		err := dataStorageServer.StartGRPCServer(dataStorageGrpcAddress, certFile, keyFile)
		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, certFile, keyFile, 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)
}