DodoSSH/tests/DodoSSH.Client.App.Tests/SoftwareTpm.cs

using System.Buffers.Binary;
using System.Net.Sockets;
using Tpm2Lib;

namespace DodoSSH.Client.App.Tests;

/// <summary>
/// Talks to a software TPM over its TCP data channel, so the real store can be driven without hardware.
/// </summary>
/// <remarks>
/// <para>
/// <b>TSS.Net ships a TcpTpmDevice and it cannot be used here.</b> It also opens the simulator's
/// <em>platform</em> socket and exchanges a handshake that the Microsoft reference simulator answers and
/// swtpm does not, so <c>Connect()</c> blocks forever — with no error and no timeout. That was found by
/// watching it hang, and it is the whole reason this file exists.
/// </para>
/// <para>
/// The data channel is the framing <c>tcti-mssim</c> uses and swtpm does implement: a command word, a
/// locality byte, a length and the bytes, answered by a length, the bytes and an acknowledgement word.
/// Power management is unnecessary when swtpm is started with <c>--flags not-need-init,startup-clear</c>,
/// which is the invocation <see cref="LinuxDeviceKeyStoreTests"/> documents.
/// </para>
/// </remarks>
internal sealed class SoftwareTpm : Tpm2Device
{
    /// <summary>tcti-mssim's "here is a command for the TPM".</summary>
    private const uint SendCommand = 8;

    private readonly string host;
    private readonly int port;

    private TcpClient? client;
    private NetworkStream? stream;

    internal SoftwareTpm(string host, int port)
    {
        this.host = host;
        this.port = port;
    }

    public override void Connect()
    {
        client = new TcpClient();
        client.Connect(host, port);
        stream = client.GetStream();
    }

    public override void Close()
    {
        stream?.Dispose();
        client?.Dispose();

        stream = null;
        client = null;
    }

    public override void DispatchCommand(CommandModifier active, byte[] inBuf, out byte[] outBuf)
    {
        ArgumentNullException.ThrowIfNull(active);
        ArgumentNullException.ThrowIfNull(inBuf);

        if (stream is null)
        {
            throw new InvalidOperationException("Not connected.");
        }

        Span<byte> header = stackalloc byte[9];

        BinaryPrimitives.WriteUInt32BigEndian(header, SendCommand);
        header[4] = active.ActiveLocality;
        BinaryPrimitives.WriteUInt32BigEndian(header[5..], (uint)inBuf.Length);

        stream.Write(header);
        stream.Write(inBuf, 0, inBuf.Length);
        stream.Flush();

        outBuf = ReadExactly((int)ReadUInt32());

        // The trailing acknowledgement word. It has to be consumed even though nothing reads it, or every
        // response after this one is misframed by four bytes.
        _ = ReadUInt32();
    }

    protected override void Dispose(bool disposing)
    {
        if (disposing)
        {
            Close();
        }

        base.Dispose(disposing);
    }

    private uint ReadUInt32() => BinaryPrimitives.ReadUInt32BigEndian(ReadExactly(4));

    private byte[] ReadExactly(int count)
    {
        var buffer = new byte[count];
        var read = 0;

        while (read < count)
        {
            var got = stream!.Read(buffer, read, count - read);

            if (got <= 0)
            {
                throw new IOException("The TPM closed the connection.");
            }

            read += got;
        }

        return buffer;
    }
}