/*
 *
 * Copyright 2018 gRPC authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package conn

import (
	"testing"

	core "google.golang.org/grpc/credentials/alts/internal"
)

// getGCMCryptoPair outputs a client/server pair on aes128gcmRekey.
func getRekeyCryptoPair(key []byte, counter []byte, t *testing.T) (ALTSRecordCrypto, ALTSRecordCrypto) {
	client, err := NewAES128GCMRekey(core.ClientSide, key)
	if err != nil {
		t.Fatalf("NewAES128GCMRekey(ClientSide, key) = %v", err)
	}
	server, err := NewAES128GCMRekey(core.ServerSide, key)
	if err != nil {
		t.Fatalf("NewAES128GCMRekey(ServerSide, key) = %v", err)
	}
	// set counter if provided.
	if counter != nil {
		if CounterSide(counter) == core.ClientSide {
			client.(*aes128gcmRekey).outCounter = CounterFromValue(counter, overflowLenAES128GCMRekey)
			server.(*aes128gcmRekey).inCounter = CounterFromValue(counter, overflowLenAES128GCMRekey)
		} else {
			server.(*aes128gcmRekey).outCounter = CounterFromValue(counter, overflowLenAES128GCMRekey)
			client.(*aes128gcmRekey).inCounter = CounterFromValue(counter, overflowLenAES128GCMRekey)
		}
	}
	return client, server
}

func testRekeyEncryptRoundtrip(client ALTSRecordCrypto, server ALTSRecordCrypto, t *testing.T) {
	// Encrypt.
	const plaintext = "This is plaintext."
	var err error
	buf := []byte(plaintext)
	buf, err = client.Encrypt(buf[:0], buf)
	if err != nil {
		t.Fatal("Encrypting with client-side context: unexpected error", err, "\n",
			"Plaintext:", []byte(plaintext))
	}

	// Encrypt a second message.
	const plaintext2 = "This is a second plaintext."
	buf2 := []byte(plaintext2)
	buf2, err = client.Encrypt(buf2[:0], buf2)
	if err != nil {
		t.Fatal("Encrypting with client-side context: unexpected error", err, "\n",
			"Plaintext:", []byte(plaintext2))
	}

	// Decryption fails: cannot decrypt second message before first.
	if got, err := server.Decrypt(nil, buf2); err == nil {
		t.Error("Decrypting client-side ciphertext with a client-side context unexpectedly succeeded; want unexpected counter error:\n",
			"  Original plaintext:", []byte(plaintext2), "\n",
			"  Ciphertext:", buf2, "\n",
			"  Decrypted plaintext:", got)
	}

	// Decryption fails: wrong counter space.
	if got, err := client.Decrypt(nil, buf); err == nil {
		t.Error("Decrypting client-side ciphertext with a client-side context unexpectedly succeeded; want counter space error:\n",
			"  Original plaintext:", []byte(plaintext), "\n",
			"  Ciphertext:", buf, "\n",
			"  Decrypted plaintext:", got)
	}

	// Decrypt first message.
	ciphertext := append([]byte(nil), buf...)
	buf, err = server.Decrypt(buf[:0], buf)
	if err != nil || string(buf) != plaintext {
		t.Fatal("Decrypting client-side ciphertext with a server-side context did not produce original content:\n",
			"  Original plaintext:", []byte(plaintext), "\n",
			"  Ciphertext:", ciphertext, "\n",
			"  Decryption error:", err, "\n",
			"  Decrypted plaintext:", buf)
	}

	// Decryption fails: replay attack.
	if got, err := server.Decrypt(nil, buf); err == nil {
		t.Error("Decrypting client-side ciphertext with a client-side context unexpectedly succeeded; want unexpected counter error:\n",
			"  Original plaintext:", []byte(plaintext), "\n",
			"  Ciphertext:", buf, "\n",
			"  Decrypted plaintext:", got)
	}
}

// Test encrypt and decrypt on roundtrip messages for aes128gcmRekey.
func (s) TestAES128GCMRekeyEncryptRoundtrip(t *testing.T) {
	// Test for aes128gcmRekey.
	key := make([]byte, 44)
	client, server := getRekeyCryptoPair(key, nil, t)
	testRekeyEncryptRoundtrip(client, server, t)
}