Mbed-TLS-Integration-Examples/matrix/matrix.c

#include "matrix.h"
#include "mbedtls/platform.h"
#include "mbedtls/build_info.h"
#include "mbedtls/debug.h"
#include "mbedtls/error.h"

#include <string.h>
#include <stdio.h>
#include <unistd.h>

#define SERVER_PORT "443"
#define SERVER_NAME "matrix.org"
#define DEBUG_LEVEL 0
#define CA_CERT_PATH "/etc/ssl/certs/ca-certificates.crt"
#define BUFFER_SIZE 8192

static void my_debug(void *ctx, int level,
                     const char *file, int line, const char *str)
{
    ((void) level);
    mbedtls_fprintf((FILE *) ctx, "%s:%04d: %s", file, line, str); // smth's up
    fflush((FILE *) ctx);
}

// connect to matrix server (TLS handshake)
int matrix_connect(matrix_connection_t *conn)
{
    int ret;

    mbedtls_net_init(&conn->server_fd);
    mbedtls_ssl_init(&conn->ssl);
    mbedtls_ssl_config_init(&conn->conf);
    mbedtls_x509_crt_init(&conn->cacert);

    mbedtls_printf(" - Connecting to %s:%s...", SERVER_NAME, SERVER_PORT);
    fflush(stdout);

    if ((ret = mbedtls_net_connect(&conn->server_fd, SERVER_NAME,
                                   SERVER_PORT, MBEDTLS_NET_PROTO_TCP)) != 0) {
        mbedtls_printf(" ! Connection error: -0x%x\n", -ret);
        return -1;
    }

    if ((ret = mbedtls_net_set_block(&conn->server_fd)) != 0) {
        mbedtls_printf(" ! failed\n");
        return -1;
    }

    if ((ret = mbedtls_ssl_config_defaults(&conn->conf,
                                           MBEDTLS_SSL_IS_CLIENT,
                                           MBEDTLS_SSL_TRANSPORT_STREAM,
                                           MBEDTLS_SSL_PRESET_DEFAULT)) != 0) {
        mbedtls_printf(" ! failed\n");
        return -1;
    }

    // Load CA certificates
    mbedtls_x509_crt_parse_file(&conn->cacert, CA_CERT_PATH);
    mbedtls_ssl_conf_ca_chain(&conn->conf, &conn->cacert, NULL);
    mbedtls_ssl_conf_authmode(&conn->conf, MBEDTLS_SSL_VERIFY_OPTIONAL);
    
    mbedtls_debug_set_threshold(DEBUG_LEVEL);
    mbedtls_ssl_conf_dbg(&conn->conf, my_debug, stdout);

    if ((ret = mbedtls_ssl_setup(&conn->ssl, &conn->conf)) != 0) {
        mbedtls_printf(" ! failed\n");
        return -1;
    }

    if ((ret = mbedtls_ssl_set_hostname(&conn->ssl, SERVER_NAME)) != 0) {
        mbedtls_printf(" ! failed\n");
        return -1;
    }

    mbedtls_ssl_set_bio(&conn->ssl, &conn->server_fd, mbedtls_net_send, mbedtls_net_recv, NULL);

    while ((ret = mbedtls_ssl_handshake(&conn->ssl)) != 0) {
        if (ret != MBEDTLS_ERR_SSL_WANT_READ && ret != MBEDTLS_ERR_SSL_WANT_WRITE) {
            mbedtls_printf(" ! Handshake error: -0x%x\n", -ret);
            return -1;
        }
    }

    mbedtls_printf(" ok\n");
    return 0;
}

// HTTP request/response
int matrix_send_request(matrix_connection_t *conn,
                        const char *req, char *res, size_t res_len)
{
    ssize_t bytes_read;
    unsigned char buf[BUFFER_SIZE];
    int retry_count = 0;
    int retry_delay = 100; // ms
    size_t res_size = 0;

    // send request
    size_t req_len = strlen(req);
    if (mbedtls_ssl_write(&conn->ssl, (const unsigned char*)req, req_len) < 0) {
        mbedtls_printf("  ! Failed to send request\n");
        return -1;
    }

    memset(res, 0, res_len); // clear the buffer
    
    // read response
    while (1) {
        bytes_read = mbedtls_ssl_read(&conn->ssl, buf, sizeof(buf) - 1);

        if (bytes_read > 0) {
            if (res_size + bytes_read < res_len) {
                memcpy(res + res_size, buf, bytes_read);
                res_size += bytes_read;
            }
            retry_count = 0;
            retry_delay = 100;
            
        } else if (bytes_read == 0) {
            break;  // connection closed
            
        } else if (bytes_read == MBEDTLS_ERR_SSL_RECEIVED_NEW_SESSION_TICKET) {
            usleep(100000);
            continue;
            
        } else if (bytes_read == MBEDTLS_ERR_SSL_WANT_READ || bytes_read == MBEDTLS_ERR_SSL_WANT_WRITE) {
            if (retry_count < 100) {
                usleep(retry_delay * 1000);
                retry_count++;
                if (retry_delay < 5000)
                    retry_delay *= 2;
                continue;
            } else {
                break;
            }
        } else {
            break;
        }
    }

    return 0;
}

// extract the token
int matrix_extract_access_token(const char *res, char *token, size_t token_len)
{
    const char *start = strstr(res, "\"access_token\":");
    if (!start) {
        mbedtls_printf(" no access_token found in response\n");
        return -1;
    }
    
    start = strchr(start, ':');
    if (!start) return -1;
    
    start = strchr(start, '"');
    if (!start) return -1;
    start++;
    
    const char *end = strchr(start, '"');
    if (!end) return -1;
    
    size_t len = end - start;
    if (len >= token_len) {
        mbedtls_printf(" token too long\n");
        return -1;
    }
    
    strncpy(token, start, len);
    token[len] = '\0';
    
    return 0;
}

// login to matrix and get access token
int matrix_login(matrix_connection_t *conn, const char *username, 
                const char *password, char *access_token, size_t token_len)
{
    char login_request[2048];
    char response[4096];

    size_t json_len = snprintf(NULL, 0, "{\"type\":\"m.login.password\",\"user\":\"%s\",\"password\":\"%s\"}", 
                              username, password);
    snprintf(login_request, sizeof(login_request), 
             "POST /_matrix/client/v3/login HTTP/1.1\r\n"
             "Host: matrix.org\r\n"
             "Content-Type: application/json\r\n"
             "Content-Length: %zu\r\n"
             "Connection: close\r\n\r\n"
             "{\"type\":\"m.login.password\",\"user\":\"%s\",\"password\":\"%s\"}",
             json_len, username, password);

    mbedtls_printf(" - Logging in as: %s ...", username);
    
    if (matrix_send_request(conn, login_request, response, sizeof(response)) != 0) {
        return -1;
    }

    // extract the token
    if (matrix_extract_access_token(response, access_token, token_len) != 0) {
        mbedtls_printf(" ! Login failed\n");
        return -1;
    }

    mbedtls_printf(" ok (got the token)\n");
    return 0;
}

// get all rooms the user has joined
int matrix_get_rooms(matrix_connection_t *conn, const char *access_token)
{
    char get_rooms_request[512];
    char response[16384];
    const char *json_start;

    snprintf(get_rooms_request, sizeof(get_rooms_request),
             "GET /_matrix/client/v3/joined_rooms HTTP/1.1\r\n"
             "Host: matrix.org\r\n"
             "Authorization: Bearer %s\r\n"
             "Connection: close\r\n\r\n",
             access_token);

    mbedtls_printf(" - Fetching room list... ");
    
    if (matrix_send_request(conn, get_rooms_request, response, sizeof(response)) != 0) {
        return -1;
    }

    // skip headers and strt with JSON body
    json_start = strstr(response, "\r\n\r\n");
    if (!json_start) {
        json_start = strstr(response, "\n\n");
        if (!json_start) {
            mbedtls_printf(" ! Invalid response format\n");
            return -1;
        }
        json_start += 2;
    } else {
        json_start += 4;
    }

    // parse joined_rooms list
    const char *room_list = strstr(json_start, "\"joined_rooms\":");
    if (room_list) {
        const char *start = strchr(room_list, '[');
        const char *end = strchr(start, ']');
        
        if (start && end) {
            mbedtls_printf(" ok\n");
            mbedtls_printf(" \nAll rooms:\n\n");
            
            const char *pos = start + 1;
            int room_count = 0;
            
            while (pos < end) {
                const char *room_start = strchr(pos, '"');
                if (!room_start || room_start >= end) break;
                room_start++;
                
                const char *room_end = strchr(room_start, '"');
                if (!room_end || room_end >= end) break;
                
                room_count++;
                mbedtls_printf("  %d. ", room_count);
                for (const char *p = room_start; p < room_end; p++) {
                    mbedtls_printf("%c", *p);
                }
                mbedtls_printf("\n");
                
                pos = room_end + 1;
            }
            
            if (room_count == 0) {
                mbedtls_printf("  (No rooms joined)\n");
            } else {
                mbedtls_printf("\n  Total: %d rooms\n", room_count);
            }
        }
    } else {
        mbedtls_printf(" ! Failed to parse room list\n");
    }

    return 0;
}

// logout & close the connection
int matrix_logout(matrix_connection_t *conn, const char *access_token)
{
    char logout_request[512];
    char response[1024];

    snprintf(logout_request, sizeof(logout_request),
             "POST /_matrix/client/v3/logout HTTP/1.1\r\n"
             "Host: matrix.org\r\n"
             "Authorization: Bearer %s\r\n"
             "Connection: close\r\n\r\n",
             access_token);

    mbedtls_printf(" - Logging out...");
    fflush(stdout);

    if (matrix_send_request(conn, logout_request, response, sizeof(response)) != 0) {
        mbedtls_printf(" ! Logout request failed\n");
        return -1;
    }

    size_t resp_len = strlen(response);
    
    // check response status
    if (strstr(response, "HTTP/1.1 200") || strstr(response, "200 OK")) {
        mbedtls_printf(" ok (access token invalidated on server)\n");
    } else if (strstr(response, "HTTP/1.1 401")) {
        mbedtls_printf(" (unauthorized)\n");
        mbedtls_printf(" ! Invalid or expired token\n");
    } else if (resp_len > 0) {
        mbedtls_printf(" (unexpected response)\n");
    } else {
        mbedtls_printf(" (no response)\n");
    }

    mbedtls_printf(" - Closing TLS connection...");
    fflush(stdout);
    mbedtls_ssl_close_notify(&conn->ssl);
    mbedtls_net_free(&conn->server_fd);
    mbedtls_ssl_free(&conn->ssl);
    mbedtls_ssl_config_free(&conn->conf);
    mbedtls_x509_crt_free(&conn->cacert);
    mbedtls_printf(" ok\n");

    return 0;
}

// and finally do the orderly cleanup
void matrix_disconnect(matrix_connection_t *conn)
{
    mbedtls_ssl_close_notify(&conn->ssl);
    mbedtls_net_free(&conn->server_fd);
    mbedtls_ssl_free(&conn->ssl);
    mbedtls_ssl_config_free(&conn->conf);
    mbedtls_x509_crt_free(&conn->cacert);
}