dsfx/internal/lib/handshake/handshake_test.go

package handshake_test

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"fmt"
	"net"
	"os"
	"sync"
	"testing"

	"numenor-labs.us/dsfx/dsfx/internal/lib/crypto/identity"
	"numenor-labs.us/dsfx/dsfx/internal/lib/handshake"
)

func TestHandshake(t *testing.T) {
	ctx := context.Background()

	// alice, represented by an ed25519 key pair.
	alice, _ := identity.Generate()
	// bob, also represented by an ed25519 key pair.
	bob, _ := identity.Generate()

	var (
		// the secret that alice should arrive at on her own
		aliceSecret []byte
		// any errors produce by alice
		aliceErr error
		// alice's public key as discovered by bob
		discoveredAlicePublicKey ed25519.PublicKey

		// the secret that bob should arrive at on his own
		bobSecret []byte
		// any errors produce by bob
		bobErr error
	)

	// Create a network pipe to simulate a network connection between alice and bob.
	client, server := net.Pipe()
	defer client.Close()
	defer server.Close()

	// Run the handshake in parallel so both sides can proceed concurrently.
	// Since they're talking through the network pipe, the entire process should
	// be simulated as if it were a real network connection.
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		aliceSecret, aliceErr = handshake.Initiate(ctx, client, alice, identity.ToPublicKey(bob))
		wg.Done()
	}()
	go func() {
		discoveredAlicePublicKey, bobSecret, bobErr = handshake.Accept(ctx, server, bob)
		wg.Done()
	}()
	wg.Wait()

	// Neither alice nor bob should have encountered any errors.
	if aliceErr != nil {
		t.Errorf("alice error: %v", aliceErr)
		return
	}
	if bobErr != nil {
		t.Errorf("bob error: %v", bobErr)
		return
	}
	// Both alice and bob should have arrived at a shared secret.
	if aliceSecret == nil || bobSecret == nil {
		t.Errorf("handshake failed: sessions are nil")
		return
	}
	// Alice and bob should have arrived at the SAME shared secret.
	if !bytes.Equal(aliceSecret, bobSecret) {
		t.Errorf("handshake failed: sessions are not equal")
		return
	}
	// Bob should have discovered alice's public key.
	if !identity.ToPublicKey(alice).Equal(discoveredAlicePublicKey) {
		t.Errorf("handshake failed: discovered public key is not equal to alice's public key")
		return
	}
}

func BenchmarkHandshake(b *testing.B) {
	for b.Loop() {
		if err := runSimulation(); err != nil {
			fmt.Fprint(os.Stderr, err)
		}
	}
}

func runSimulation() error {
	ctx := context.Background()

	// alice, represented by an ed25519 key pair.
	alice, _ := identity.Generate()
	// bob, also represented by an ed25519 key pair.
	bob, _ := identity.Generate()

	var (
		// the secret that alice should arrive at on her own
		// aliceSecret []byte
		// any errors produce by alice
		aliceErr error
		// alice's public key as discovered by bob
		// discoveredAlicePublicKey *ecdsa.PublicKey

		// the secret that bob should arrive at on his own
		// bobSecret []byte
		// any errors produce by bob
		bobErr error
	)

	// Create a network pipe to simulate a network connection between alice and bob.
	client, server := net.Pipe()
	defer client.Close()
	defer server.Close()

	// Run the handshake in parallel so both sides can proceed concurrently.
	// Since they're talking through the network pipe, the entire process should
	// be simulated as if it were a real network connection.
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		_, aliceErr = handshake.Initiate(ctx, client, alice, identity.ToPublicKey(bob))
		wg.Done()
	}()
	go func() {
		_, _, bobErr = handshake.Accept(ctx, server, bob)
		wg.Done()
	}()
	wg.Wait()
	if aliceErr != nil {
		return aliceErr
	}
	if bobErr != nil {
		return bobErr
	}
	return nil
}