/*      $OpenBSD: handshake_table.c,v 1.18 2022/12/01 13:49:12 tb Exp $ */
/*
 * Copyright (c) 2019 Theo Buehler <tb@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <err.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include "tls13_handshake.h"

#define MAX_FLAGS (UINT8_MAX + 1)

/*
 * From RFC 8446:
 *
 * Appendix A.  State Machine
 *
 *    This appendix provides a summary of the legal state transitions for
 *    the client and server handshakes.  State names (in all capitals,
 *    e.g., START) have no formal meaning but are provided for ease of
 *    comprehension.  Actions which are taken only in certain circumstances
 *    are indicated in [].  The notation "K_{send,recv} = foo" means "set
 *    the send/recv key to the given key".
 *
 * A.1.  Client
 *
 *                               START <----+
 *                Send ClientHello |        | Recv HelloRetryRequest
 *           [K_send = early data] |        |
 *                                 v        |
 *            /                 WAIT_SH ----+
 *            |                    | Recv ServerHello
 *            |                    | K_recv = handshake
 *        Can |                    V
 *       send |                 WAIT_EE
 *      early |                    | Recv EncryptedExtensions
 *       data |           +--------+--------+
 *            |     Using |                 | Using certificate
 *            |       PSK |                 v
 *            |           |            WAIT_CERT_CR
 *            |           |        Recv |       | Recv CertificateRequest
 *            |           | Certificate |       v
 *            |           |             |    WAIT_CERT
 *            |           |             |       | Recv Certificate
 *            |           |             v       v
 *            |           |              WAIT_CV
 *            |           |                 | Recv CertificateVerify
 *            |           +> WAIT_FINISHED <+
 *            |                  | Recv Finished
 *            \                  | [Send EndOfEarlyData]
 *                               | K_send = handshake
 *                               | [Send Certificate [+ CertificateVerify]]
 *     Can send                  | Send Finished
 *     app data   -->            | K_send = K_recv = application
 *     after here                v
 *                           CONNECTED
 *
 *    Note that with the transitions as shown above, clients may send
 *    alerts that derive from post-ServerHello messages in the clear or
 *    with the early data keys.  If clients need to send such alerts, they
 *    SHOULD first rekey to the handshake keys if possible.
 *
 */

struct child {
        enum tls13_message_type mt;
        uint8_t                 flag;
        uint8_t                 forced;
        uint8_t                 illegal;
};

static struct child stateinfo[][TLS13_NUM_MESSAGE_TYPES] = {
        [CLIENT_HELLO] = {
                {
                        .mt = SERVER_HELLO_RETRY_REQUEST,
                },
                {
                        .mt = SERVER_HELLO,
                        .flag = WITHOUT_HRR,
                },
        },
        [SERVER_HELLO_RETRY_REQUEST] = {
                {
                        .mt = CLIENT_HELLO_RETRY,
                },
        },
        [CLIENT_HELLO_RETRY] = {
                {
                        .mt = SERVER_HELLO,
                },
        },
        [SERVER_HELLO] = {
                {
                        .mt = SERVER_ENCRYPTED_EXTENSIONS,
                },
        },
        [SERVER_ENCRYPTED_EXTENSIONS] = {
                {
                        .mt = SERVER_CERTIFICATE_REQUEST,
                },
                {       .mt = SERVER_CERTIFICATE,
                        .flag = WITHOUT_CR,
                },
                {
                        .mt = SERVER_FINISHED,
                        .flag = WITH_PSK,
                },
        },
        [SERVER_CERTIFICATE_REQUEST] = {
                {
                        .mt = SERVER_CERTIFICATE,
                },
        },
        [SERVER_CERTIFICATE] = {
                {
                        .mt = SERVER_CERTIFICATE_VERIFY,
                },
        },
        [SERVER_CERTIFICATE_VERIFY] = {
                {
                        .mt = SERVER_FINISHED,
                },
        },
        [SERVER_FINISHED] = {
                {
                        .mt = CLIENT_FINISHED,
                        .forced = WITHOUT_CR | WITH_PSK,
                },
                {
                        .mt = CLIENT_CERTIFICATE,
                        .illegal = WITHOUT_CR | WITH_PSK,
                },
        },
        [CLIENT_CERTIFICATE] = {
                {
                        .mt = CLIENT_FINISHED,
                },
                {
                        .mt = CLIENT_CERTIFICATE_VERIFY,
                        .flag = WITH_CCV,
                },
        },
        [CLIENT_CERTIFICATE_VERIFY] = {
                {
                        .mt = CLIENT_FINISHED,
                },
        },
        [CLIENT_FINISHED] = {
                {
                        .mt = APPLICATION_DATA,
                },
        },
        [APPLICATION_DATA] = {
                {
                        .mt = 0,
                },
        },
};

const size_t     stateinfo_count = sizeof(stateinfo) / sizeof(stateinfo[0]);

void             build_table(enum tls13_message_type
                     table[MAX_FLAGS][TLS13_NUM_MESSAGE_TYPES],
                     struct child current, struct child end,
                     struct child path[], uint8_t flags, unsigned int depth);
size_t           count_handshakes(void);
void             edge(enum tls13_message_type start,
                     enum tls13_message_type end, uint8_t flag);
const char      *flag2str(uint8_t flag);
void             flag_label(uint8_t flag);
void             forced_edges(enum tls13_message_type start,
                     enum tls13_message_type end, uint8_t forced);
int              generate_graphics(void);
void             fprint_entry(FILE *stream,
                     enum tls13_message_type path[TLS13_NUM_MESSAGE_TYPES],
                     uint8_t flags);
void             fprint_flags(FILE *stream, uint8_t flags);
const char      *mt2str(enum tls13_message_type mt);
void             usage(void);
int              verify_table(enum tls13_message_type
                     table[MAX_FLAGS][TLS13_NUM_MESSAGE_TYPES], int print);

const char *
flag2str(uint8_t flag)
{
        const char *ret;

        if (flag & (flag - 1))
                errx(1, "more than one bit is set");

        switch (flag) {
        case INITIAL:
                ret = "INITIAL";
                break;
        case NEGOTIATED:
                ret = "NEGOTIATED";
                break;
        case WITHOUT_CR:
                ret = "WITHOUT_CR";
                break;
        case WITHOUT_HRR:
                ret = "WITHOUT_HRR";
                break;
        case WITH_PSK:
                ret = "WITH_PSK";
                break;
        case WITH_CCV:
                ret = "WITH_CCV";
                break;
        case WITH_0RTT:
                ret = "WITH_0RTT";
                break;
        default:
                ret = "UNKNOWN";
        }

        return ret;
}

const char *
mt2str(enum tls13_message_type mt)
{
        const char *ret;

        switch (mt) {
        case INVALID:
                ret = "INVALID";
                break;
        case CLIENT_HELLO:
                ret = "CLIENT_HELLO";
                break;
        case CLIENT_HELLO_RETRY:
                ret = "CLIENT_HELLO_RETRY";
                break;
        case CLIENT_END_OF_EARLY_DATA:
                ret = "CLIENT_END_OF_EARLY_DATA";
                break;
        case CLIENT_CERTIFICATE:
                ret = "CLIENT_CERTIFICATE";
                break;
        case CLIENT_CERTIFICATE_VERIFY:
                ret = "CLIENT_CERTIFICATE_VERIFY";
                break;
        case CLIENT_FINISHED:
                ret = "CLIENT_FINISHED";
                break;
        case SERVER_HELLO:
                ret = "SERVER_HELLO";
                break;
        case SERVER_HELLO_RETRY_REQUEST:
                ret = "SERVER_HELLO_RETRY_REQUEST";
                break;
        case SERVER_ENCRYPTED_EXTENSIONS:
                ret = "SERVER_ENCRYPTED_EXTENSIONS";
                break;
        case SERVER_CERTIFICATE:
                ret = "SERVER_CERTIFICATE";
                break;
        case SERVER_CERTIFICATE_VERIFY:
                ret = "SERVER_CERTIFICATE_VERIFY";
                break;
        case SERVER_CERTIFICATE_REQUEST:
                ret = "SERVER_CERTIFICATE_REQUEST";
                break;
        case SERVER_FINISHED:
                ret = "SERVER_FINISHED";
                break;
        case APPLICATION_DATA:
                ret = "APPLICATION_DATA";
                break;
        case TLS13_NUM_MESSAGE_TYPES:
                ret = "TLS13_NUM_MESSAGE_TYPES";
                break;
        default:
                ret = "UNKNOWN";
                break;
        }

        return ret;
}

void
fprint_flags(FILE *stream, uint8_t flags)
{
        int first = 1, i;

        if (flags == 0) {
                fprintf(stream, "%s", flag2str(flags));
                return;
        }

        for (i = 0; i < 8; i++) {
                uint8_t set = flags & (1U << i);

                if (set) {
                        fprintf(stream, "%s%s", first ? "" : " | ",
                            flag2str(set));
                        first = 0;
                }
        }
}

void
fprint_entry(FILE *stream,
    enum tls13_message_type path[TLS13_NUM_MESSAGE_TYPES], uint8_t flags)
{
        int i;

        fprintf(stream, "\t[");
        fprint_flags(stream, flags);
        fprintf(stream, "] = {\n");

        for (i = 0; i < TLS13_NUM_MESSAGE_TYPES; i++) {
                if (path[i] == 0)
                        break;
                fprintf(stream, "\t\t%s,\n", mt2str(path[i]));
        }
        fprintf(stream, "\t},\n");
}

void
edge(enum tls13_message_type start, enum tls13_message_type end,
    uint8_t flag)
{
        printf("\t%s -> %s", mt2str(start), mt2str(end));
        flag_label(flag);
        printf(";\n");
}

void
flag_label(uint8_t flag)
{
        if (flag)
                printf(" [label=\"%s\"]", flag2str(flag));
}

void
forced_edges(enum tls13_message_type start, enum tls13_message_type end,
    uint8_t forced)
{
        uint8_t forced_flag, i;

        if (forced == 0)
                return;

        for (i = 0; i < 8; i++) {
                forced_flag = forced & (1U << i);
                if (forced_flag)
                        edge(start, end, forced_flag);
        }
}

int
generate_graphics(void)
{
        enum tls13_message_type start, end;
        unsigned int            child;
        uint8_t                 flag;
        uint8_t                 forced;

        printf("digraph G {\n");
        printf("\t%s [shape=box];\n", mt2str(CLIENT_HELLO));
        printf("\t%s [shape=box];\n", mt2str(APPLICATION_DATA));

        for (start = CLIENT_HELLO; start < APPLICATION_DATA; start++) {
                for (child = 0; stateinfo[start][child].mt != 0; child++) {
                        end = stateinfo[start][child].mt;
                        flag = stateinfo[start][child].flag;
                        forced = stateinfo[start][child].forced;

                        if (forced == 0)
                                edge(start, end, flag);
                        else
                                forced_edges(start, end, forced);
                }
        }

        printf("}\n");
        return 0;
}

extern enum tls13_message_type  handshakes[][TLS13_NUM_MESSAGE_TYPES];
extern size_t                   handshake_count;

size_t
count_handshakes(void)
{
        size_t  ret = 0, i;

        for (i = 0; i < handshake_count; i++) {
                if (handshakes[i][0] != INVALID)
                        ret++;
        }

        return ret;
}

void
build_table(enum tls13_message_type table[MAX_FLAGS][TLS13_NUM_MESSAGE_TYPES],
    struct child current, struct child end, struct child path[], uint8_t flags,
    unsigned int depth)
{
        unsigned int i;

        if (depth >= TLS13_NUM_MESSAGE_TYPES - 1)
                errx(1, "recursed too deeply");

        /* Record current node. */
        path[depth++] = current;
        flags |= current.flag;

        /* If we haven't reached the end, recurse over the children. */
        if (current.mt != end.mt) {
                for (i = 0; stateinfo[current.mt][i].mt != 0; i++) {
                        struct child child = stateinfo[current.mt][i];
                        int forced = stateinfo[current.mt][i].forced;
                        int illegal = stateinfo[current.mt][i].illegal;

                        if ((forced == 0 || (forced & flags)) &&
                            (illegal == 0 || !(illegal & flags)))
                                build_table(table, child, end, path, flags,
                                    depth);
                }
                return;
        }

        if (flags == 0)
                errx(1, "path does not set flags");

        if (table[flags][0] != 0)
                errx(1, "path traversed twice");

        for (i = 0; i < depth; i++)
                table[flags][i] = path[i].mt;
}

int
verify_table(enum tls13_message_type table[MAX_FLAGS][TLS13_NUM_MESSAGE_TYPES],
    int print)
{
        int     success = 1, i;
        size_t  num_valid, num_found = 0;
        uint8_t flags = 0;

        do {
                if (table[flags][0] == 0)
                        continue;

                num_found++;

                for (i = 0; i < TLS13_NUM_MESSAGE_TYPES; i++) {
                        if (table[flags][i] != handshakes[flags][i]) {
                                fprintf(stderr,
                                    "incorrect entry %d of handshake ", i);
                                fprint_flags(stderr, flags);
                                fprintf(stderr, "\n");
                                success = 0;
                        }
                }

                if (print)
                        fprint_entry(stdout, table[flags], flags);
        } while(++flags != 0);

        num_valid = count_handshakes();
        if (num_valid != num_found) {
                fprintf(stderr,
                    "incorrect number of handshakes: want %zu, got %zu.\n",
                    num_valid, num_found);
                success = 0;
        }

        return success;
}

void
usage(void)
{
        fprintf(stderr, "usage: handshake_table [-C | -g]\n");
        exit(1);
}

int
main(int argc, char *argv[])
{
        static enum tls13_message_type
            hs_table[MAX_FLAGS][TLS13_NUM_MESSAGE_TYPES] = {
                [INITIAL] = {
                        CLIENT_HELLO,
                        SERVER_HELLO_RETRY_REQUEST,
                        CLIENT_HELLO_RETRY,
                        SERVER_HELLO,
                },
        };
        struct child    start = {
                .mt = CLIENT_HELLO,
        };
        struct child    end = {
                .mt = APPLICATION_DATA,
        };
        struct child    path[TLS13_NUM_MESSAGE_TYPES] = {{0}};
        uint8_t         flags = NEGOTIATED;
        unsigned int    depth = 0;
        int             ch, graphviz = 0, print = 0;

        while ((ch = getopt(argc, argv, "Cg")) != -1) {
                switch (ch) {
                case 'C':
                        print = 1;
                        break;
                case 'g':
                        graphviz = 1;
                        break;
                default:
                        usage();
                }
        }
        argc -= optind;
        argv += optind;

        if (argc != 0)
                usage();

        if (graphviz && print)
                usage();

        if (graphviz)
                return generate_graphics();

        build_table(hs_table, start, end, path, flags, depth);
        if (!verify_table(hs_table, print))
                return 1;

        return 0;
}