CBDDC/tests/ZB.MOM.WW.CBDDC.Network.Tests/SecureHandshakeTests.cs

using System;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using ZB.MOM.WW.CBDDC.Network.Security;
using Microsoft.Extensions.Logging.Abstractions;
using Xunit;

namespace ZB.MOM.WW.CBDDC.Network.Tests
{
    public class SecureHandshakeTests
    {
        /// <summary>
        /// Verifies handshake negotiation succeeds between initiator and responder services.
        /// </summary>
        [Fact]
        public async Task Handshake_Should_Succeed_Between_Two_Services()
        {
            // Arrange
            var clientStream = new PipeStream();
            var serverStream = new PipeStream();

            // Client writes to clientStream, server reads from clientStream
            // Server writes to serverStream, client reads from serverStream

            var clientSocket = new DuplexStream(serverStream, clientStream); // Read from server, Write to client
            var serverSocket = new DuplexStream(clientStream, serverStream); // Read from client, Write to server

            var clientService = new SecureHandshakeService(NullLogger<SecureHandshakeService>.Instance);
            var serverService = new SecureHandshakeService(NullLogger<SecureHandshakeService>.Instance);

            var cts = new CancellationTokenSource(TimeSpan.FromSeconds(5));

            // Act
            var clientTask = clientService.HandshakeAsync(clientSocket, isInitiator: true, myNodeId: "client", token: cts.Token);
            var serverTask = serverService.HandshakeAsync(serverSocket, isInitiator: false, myNodeId: "server", token: cts.Token);

            await Task.WhenAll(clientTask, serverTask);

            // Assert
            var clientState = clientTask.Result;
            var serverState = serverTask.Result;

            clientState.ShouldNotBeNull();
            serverState.ShouldNotBeNull();

            // Keys should match (Symmetric)
            clientState!.EncryptKey.ShouldBe(serverState!.DecryptKey);
            clientState.DecryptKey.ShouldBe(serverState.EncryptKey);
        }

        // Simulates a pipe. Writes go to buffer, Reads drain buffer.
        class SimplexStream : MemoryStream
        {
            // Simple approach: Use one MemoryStream as a shared buffer?
            // No, MemoryStream is not thread safe for concurrent Read/Write in this pipe manner really.
            // Better to use a producer/consumer stream but for simplicity let's use a basic blocking queue logic or just wait.
            // Actually, for unit tests, strictly ordered operations are better. But handshake is interactive.
            // We need a proper pipe.
        }

        // Let's use a simple PipeStream implementation using SemaphoreSlim for sync
        class PipeStream : Stream
        {
            private readonly MemoryStream _buffer = new MemoryStream();
            private readonly SemaphoreSlim _readSemaphore = new SemaphoreSlim(0);
            private readonly object _lock = new object();

            /// <inheritdoc />
            public override bool CanRead => true;
            /// <inheritdoc />
            public override bool CanSeek => false;
            /// <inheritdoc />
            public override bool CanWrite => true;
            /// <inheritdoc />
            public override long Length => _buffer.Length;
            /// <inheritdoc />
            public override long Position { get => _buffer.Position; set => throw new NotSupportedException(); }

            /// <inheritdoc />
            public override void Flush() { }

            /// <inheritdoc />
            public override int Read(byte[] buffer, int offset, int count) => throw new NotImplementedException("Use Async");

            /// <inheritdoc />
            public override async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
            {
                await _readSemaphore.WaitAsync(cancellationToken);
                lock (_lock)
                {
                    _buffer.Position = 0;
                    int read = _buffer.Read(buffer, offset, count);

                    // Compact buffer (inefficient but works for unit tests)
                    byte[] remaining = _buffer.ToArray().Skip(read).ToArray();
                    _buffer.SetLength(0);
                    _buffer.Write(remaining, 0, remaining.Length);

                    if (_buffer.Length > 0) _readSemaphore.Release(); // Signal if data remains

                    return read;
                }
            }

            /// <inheritdoc />
            public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();

            /// <inheritdoc />
            public override void SetLength(long value) => throw new NotSupportedException();

            /// <inheritdoc />
            public override void Write(byte[] buffer, int offset, int count)
            {
                lock (_lock)
                {
                    long pos = _buffer.Position;
                    _buffer.Seek(0, SeekOrigin.End);
                    _buffer.Write(buffer, offset, count);
                    _buffer.Position = pos;
                }
                _readSemaphore.Release();
            }
        }

        class DuplexStream : Stream
        {
            private readonly Stream _readSource;
            private readonly Stream _writeTarget;

            /// <summary>
            /// Initializes a new instance of the <see cref="DuplexStream"/> class.
            /// </summary>
            /// <param name="readSource">The underlying stream used for read operations.</param>
            /// <param name="writeTarget">The underlying stream used for write operations.</param>
            public DuplexStream(Stream readSource, Stream writeTarget)
            {
                _readSource = readSource;
                _writeTarget = writeTarget;
            }

            /// <inheritdoc />
            public override bool CanRead => true;
            /// <inheritdoc />
            public override bool CanSeek => false;
            /// <inheritdoc />
            public override bool CanWrite => true;
            /// <inheritdoc />
            public override long Length => 0;
            /// <inheritdoc />
            public override long Position { get => 0; set { } }

            /// <inheritdoc />
            public override void Flush() => _writeTarget.Flush();

            /// <inheritdoc />
            public override int Read(byte[] buffer, int offset, int count) => _readSource.Read(buffer, offset, count);

            /// <inheritdoc />
            public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
                => _readSource.ReadAsync(buffer, offset, count, cancellationToken);

            /// <inheritdoc />
            public override void Write(byte[] buffer, int offset, int count) => _writeTarget.Write(buffer, offset, count);

            /// <inheritdoc />
            public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
                => _writeTarget.WriteAsync(buffer, offset, count, cancellationToken);

            /// <inheritdoc />
            public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();

            /// <inheritdoc />
            public override void SetLength(long value) => throw new NotSupportedException();
        }
    }
}