Merge pull request #9121 from valeriosetti/issue8963

Add client-server build to all.sh
diff --git a/tests/Makefile b/tests/Makefile
index 7fb4f35..3cafccf 100644
--- a/tests/Makefile
+++ b/tests/Makefile
@@ -182,6 +182,7 @@
 	rm -f include/test/instrument_record_status.h
 	rm -f include/alt-extra/*/*_alt.h
 	rm -rf libtestdriver1
+	rm -rf libpsaclient libpsaserver
 	rm -f ../library/libtestdriver1.a
 else
 	if exist *.c del /Q /F *.c
diff --git a/tests/psa-client-server/psasim/Makefile b/tests/psa-client-server/psasim/Makefile
index a84483c..db0c412 100644
--- a/tests/psa-client-server/psasim/Makefile
+++ b/tests/psa-client-server/psasim/Makefile
@@ -1,23 +1,62 @@
-CFLAGS ?= -Wall -Werror -std=c99 -D_XOPEN_SOURCE=1 -D_POSIX_C_SOURCE=200809L
+CFLAGS += -Wall -Werror -std=c99 -D_XOPEN_SOURCE=1 -D_POSIX_C_SOURCE=200809L
 
 ifeq ($(DEBUG),1)
-	CFLAGS += -DDEBUG -O0 -g
+CFLAGS += -DDEBUG -O0 -g
 endif
 
-.PHONY: all lib test run
+LIBPSACLIENT := -Llibpsaclient/ -lmbedcrypto -lmbedx509 -lmbedtls
+LIBPSASERVER := -Llibpsaserver/ -lmbedcrypto
 
-all: lib test
+MBEDTLS_ROOT_PATH = ../../..
+COMMON_INCLUDE := -I./include -I$(MBEDTLS_ROOT_PATH)/include
 
-lib:
-	$(MAKE) -C src CFLAGS="$(CFLAGS)"
+TEST_BIN = 	test/psa_client \
+			test/psa_partition
 
-test: lib
-	$(MAKE) -C test CFLAGS="$(CFLAGS)"
+GENERATED_H_FILES =	include/psa_manifest/manifest.h \
+					include/psa_manifest/pid.h \
+					include/psa_manifest/sid.h
+
+PSA_CLIENT_SRC = src/psa_ff_client.c \
+				 src/client.c
+
+PARTITION_SERVER_BOOTSTRAP = src/psa_ff_bootstrap_TEST_PARTITION.c
+
+PSA_SERVER_SRC = $(PARTITION_SERVER_BOOTSTRAP) \
+				 src/psa_ff_server.c
+
+.PHONY: all clean libpsaclient libpsaserver
+
+all: $(TEST_BIN)
+
+test/seedfile:
+	dd if=/dev/urandom of=./test/seedfile bs=64 count=1
+
+test/psa_client: $(PSA_CLIENT_SRC) $(GENERATED_H_FILES)
+	$(CC) $(COMMON_INCLUDE) $(CFLAGS) $(PSA_CLIENT_SRC) $(LIBPSACLIENT) $(LDFLAGS) -o $@
+
+test/psa_partition: $(PSA_SERVER_SRC) $(GENERATED_H_FILES) test/seedfile
+	$(CC) $(COMMON_INCLUDE) $(CFLAGS) $(PSA_SERVER_SRC) $(LIBPSASERVER) $(LDFLAGS) -o $@
+
+$(PARTITION_SERVER_BOOTSTRAP) $(GENERATED_H_FILES): src/manifest.json src/server.c
+	tools/psa_autogen.py src/manifest.json
+
+# Build MbedTLS libraries (crypto, x509 and tls) and copy them locally to
+# build client/server applications.
+#
+# Note: these rules assume that mbedtls_config.h is already configured by all.sh.
+# If not using all.sh then the user must do it manually.
+libpsaclient libpsaserver:
+	$(MAKE) -C $(MBEDTLS_ROOT_PATH)/library CFLAGS="$(CFLAGS)" LDFLAGS="$(LDFLAGS)" libmbedcrypto.a libmbedx509.a libmbedtls.a
+	mkdir -p $@
+	cp $(MBEDTLS_ROOT_PATH)/library/libmbed*.a $@/
+	$(MAKE) -C $(MBEDTLS_ROOT_PATH) clean
 
 clean:
-	rm -f $(PSA_LIB) $(PSA_LIB_OBJS)
-	$(MAKE) -C test clean
-	$(MAKE) -C src clean
-
-run: test
-	cd test && ./run_test.sh
+	rm -f $(TEST_BIN)
+	rm -f $(PARTITION_SERVER_BOOTSTRAP)
+	rm -rf libpsaclient libpsaserver
+	rm -rf include/psa_manifest
+	rm -f test/psa_service_* test/psa_notify_*
+	rm -f test/*.log
+	rm -f test/seedfile
diff --git a/tests/psa-client-server/psasim/include/psa/client.h b/tests/psa-client-server/psasim/include/client.h
similarity index 97%
rename from tests/psa-client-server/psasim/include/psa/client.h
rename to tests/psa-client-server/psasim/include/client.h
index d1af993..d48498e 100644
--- a/tests/psa-client-server/psasim/include/psa/client.h
+++ b/tests/psa-client-server/psasim/include/client.h
@@ -15,7 +15,9 @@
 #include <stdint.h>
 #include <stddef.h>
 
-#include "psa/error.h"
+#include "psa/crypto.h"
+
+#include "error_ext.h"
 /*********************** PSA Client Macros and Types *************************/
 
 #define PSA_FRAMEWORK_VERSION  (0x0100)
diff --git a/tests/psa-client-server/psasim/include/psa/common.h b/tests/psa-client-server/psasim/include/common.h
similarity index 96%
rename from tests/psa-client-server/psasim/include/psa/common.h
rename to tests/psa-client-server/psasim/include/common.h
index d0205d2..ee5b5a3 100644
--- a/tests/psa-client-server/psasim/include/psa/common.h
+++ b/tests/psa-client-server/psasim/include/common.h
@@ -27,7 +27,6 @@
 
 #define NON_SECURE (1 << 30)
 
-typedef int32_t psa_status_t;
 typedef int32_t psa_handle_t;
 
 #define PSA_MAX_IOVEC (4u)
diff --git a/tests/psa-client-server/psasim/include/error_ext.h b/tests/psa-client-server/psasim/include/error_ext.h
new file mode 100644
index 0000000..6c82b8a
--- /dev/null
+++ b/tests/psa-client-server/psasim/include/error_ext.h
@@ -0,0 +1,19 @@
+/* PSA status codes used by psasim. */
+
+/*
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
+ */
+
+#ifndef PSA_ERROR_H
+#define PSA_ERROR_H
+
+#include <stdint.h>
+
+#include "common.h"
+
+#define PSA_ERROR_PROGRAMMER_ERROR      ((psa_status_t) -129)
+#define PSA_ERROR_CONNECTION_REFUSED    ((psa_status_t) -130)
+#define PSA_ERROR_CONNECTION_BUSY       ((psa_status_t) -131)
+
+#endif
diff --git a/tests/psa-client-server/psasim/include/psasim/init.h b/tests/psa-client-server/psasim/include/init.h
similarity index 94%
rename from tests/psa-client-server/psasim/include/psasim/init.h
rename to tests/psa-client-server/psasim/include/init.h
index 9496fc2..de95d90 100644
--- a/tests/psa-client-server/psasim/include/psasim/init.h
+++ b/tests/psa-client-server/psasim/include/init.h
@@ -6,7 +6,7 @@
  */
 
 #include <stdint.h>
-#include <psa/service.h>
+#include <service.h>
 void raise_signal(psa_signal_t signal);
 void __init_psasim(const char **array,
                    int size,
diff --git a/tests/psa-client-server/psasim/include/psa/lifecycle.h b/tests/psa-client-server/psasim/include/lifecycle.h
similarity index 100%
rename from tests/psa-client-server/psasim/include/psa/lifecycle.h
rename to tests/psa-client-server/psasim/include/lifecycle.h
diff --git a/tests/psa-client-server/psasim/include/psa/error.h b/tests/psa-client-server/psasim/include/psa/error.h
deleted file mode 100644
index 44fc0b1..0000000
--- a/tests/psa-client-server/psasim/include/psa/error.h
+++ /dev/null
@@ -1,38 +0,0 @@
-/* PSA status codes used by psasim. */
-
-/*
- *  Copyright The Mbed TLS Contributors
- *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
- */
-
-#ifndef PSA_ERROR_H
-#define PSA_ERROR_H
-
-#include <stdint.h>
-
-#include "psa/common.h"
-
-#define PSA_SUCCESS                     ((psa_status_t) 0)
-
-#define PSA_ERROR_PROGRAMMER_ERROR      ((psa_status_t) -129)
-#define PSA_ERROR_CONNECTION_REFUSED    ((psa_status_t) -130)
-#define PSA_ERROR_CONNECTION_BUSY       ((psa_status_t) -131)
-#define PSA_ERROR_GENERIC_ERROR         ((psa_status_t) -132)
-#define PSA_ERROR_NOT_PERMITTED         ((psa_status_t) -133)
-#define PSA_ERROR_NOT_SUPPORTED         ((psa_status_t) -134)
-#define PSA_ERROR_INVALID_ARGUMENT      ((psa_status_t) -135)
-#define PSA_ERROR_INVALID_HANDLE        ((psa_status_t) -136)
-#define PSA_ERROR_BAD_STATE             ((psa_status_t) -137)
-#define PSA_ERROR_BUFFER_TOO_SMALL      ((psa_status_t) -138)
-#define PSA_ERROR_ALREADY_EXISTS        ((psa_status_t) -139)
-#define PSA_ERROR_DOES_NOT_EXIST        ((psa_status_t) -140)
-#define PSA_ERROR_INSUFFICIENT_MEMORY   ((psa_status_t) -141)
-#define PSA_ERROR_INSUFFICIENT_STORAGE  ((psa_status_t) -142)
-#define PSA_ERROR_INSUFFICIENT_DATA     ((psa_status_t) -143)
-#define PSA_ERROR_SERVICE_FAILURE       ((psa_status_t) -144)
-#define PSA_ERROR_COMMUNICATION_FAILURE ((psa_status_t) -145)
-#define PSA_ERROR_STORAGE_FAILURE       ((psa_status_t) -146)
-#define PSA_ERROR_HARDWARE_FAILURE      ((psa_status_t) -147)
-#define PSA_ERROR_INVALID_SIGNATURE     ((psa_status_t) -149)
-
-#endif
diff --git a/tests/psa-client-server/psasim/include/psa/service.h b/tests/psa-client-server/psasim/include/service.h
similarity index 99%
rename from tests/psa-client-server/psasim/include/psa/service.h
rename to tests/psa-client-server/psasim/include/service.h
index c8c0024..cbcb918 100644
--- a/tests/psa-client-server/psasim/include/psa/service.h
+++ b/tests/psa-client-server/psasim/include/service.h
@@ -15,7 +15,9 @@
 #include <stdint.h>
 #include <stddef.h>
 
-#include "psa/common.h"
+#include "common.h"
+
+#include "psa/crypto.h"
 
 /********************** PSA Secure Partition Macros and Types ****************/
 
diff --git a/tests/psa-client-server/psasim/include/psa/util.h b/tests/psa-client-server/psasim/include/util.h
similarity index 96%
rename from tests/psa-client-server/psasim/include/psa/util.h
rename to tests/psa-client-server/psasim/include/util.h
index c3669a1..558149f 100644
--- a/tests/psa-client-server/psasim/include/psa/util.h
+++ b/tests/psa-client-server/psasim/include/util.h
@@ -5,7 +5,7 @@
  *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
  */
 
-#include "psa/service.h"
+#include "service.h"
 
 #define PRINT(fmt, ...) \
     fprintf(stdout, fmt "\n", ##__VA_ARGS__)
diff --git a/tests/psa-client-server/psasim/src/Makefile b/tests/psa-client-server/psasim/src/Makefile
deleted file mode 100644
index fc6ba25..0000000
--- a/tests/psa-client-server/psasim/src/Makefile
+++ /dev/null
@@ -1,17 +0,0 @@
-INCLUDE = -I../include/
-PSA_LIB = libpsaff.a
-
-PSA_LIB_OBJS = client.o service.o
-
-.PHONY: all lib
-
-all: $(PSA_LIB)
-
-%.o: %.c
-	$(CC) $(INCLUDE) $(CFLAGS) -c $< -o $@
-
-$(PSA_LIB): $(PSA_LIB_OBJS)
-	$(AR) rcs $(PSA_LIB) client.o service.o
-
-clean:
-	rm -f $(PSA_LIB) $(PSA_LIB_OBJS)
diff --git a/tests/psa-client-server/psasim/src/client.c b/tests/psa-client-server/psasim/src/client.c
index 5a3986e..550a6e8 100644
--- a/tests/psa-client-server/psasim/src/client.c
+++ b/tests/psa-client-server/psasim/src/client.c
@@ -1,392 +1,54 @@
-/* PSA firmware framework client API */
+/* psasim test client */
 
 /*
  *  Copyright The Mbed TLS Contributors
  *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
  */
 
-#include <stdint.h>
-#include <stdlib.h>
-#include <stddef.h>
-#include <assert.h>
 #include <stdio.h>
-#include <string.h>
-#include <strings.h>
-#include <inttypes.h>
-#include <sys/types.h>
-#include <sys/ipc.h>
-#include <sys/msg.h>
+#include <unistd.h>
 
-#include "psa/client.h"
-#include "psa/common.h"
-#include "psa/error.h"
-#include "psa/util.h"
+/* Includes from psasim */
+#include <client.h>
+#include <util.h>
+#include "psa_manifest/sid.h"
+#include "psa_functions_codes.h"
 
-typedef struct internal_handle {
-    int server_qid;
-    int client_qid;
-    int internal_server_qid;
-    int valid;
-} internal_handle_t;
+/* Includes from mbedtls */
+#include "mbedtls/version.h"
+#include "psa/crypto.h"
 
-typedef struct vectors {
-    const psa_invec *in_vec;
-    size_t in_len;
-    psa_outvec *out_vec;
-    size_t out_len;
-} vectors_t;
+#define CLIENT_PRINT(fmt, ...) \
+    PRINT("Client: " fmt, ##__VA_ARGS__)
 
-/* Note that this implementation is functional and not secure */
-int __psa_ff_client_security_state = NON_SECURE;
-
-/* Access to this global is not thread safe */
-#define MAX_HANDLES 32
-static internal_handle_t handles[MAX_HANDLES] = { { 0 } };
-
-static int get_next_free_handle()
+int main()
 {
-    /* Never return handle 0 as it's a special null handle */
-    for (int i = 1; i < MAX_HANDLES; i++) {
-        if (handles[i].valid == 0) {
-            return i;
-        }
-    }
-    return -1;
-}
+    char mbedtls_version[18];
+    // psa_invec invecs[1];
+    // psa_outvec outvecs[1];
+    psa_status_t status;
 
-static int handle_is_valid(psa_handle_t handle)
-{
-    if (handle > 0 && handle < MAX_HANDLES) {
-        if (handles[handle].valid == 1) {
-            return 1;
-        }
+    mbedtls_version_get_string_full(mbedtls_version);
+    CLIENT_PRINT("%s", mbedtls_version);
+
+    CLIENT_PRINT("My PID: %d", getpid());
+
+    CLIENT_PRINT("PSA version: %u", psa_version(PSA_SID_CRYPTO_SID));
+    psa_handle_t h = psa_connect(PSA_SID_CRYPTO_SID, 1);
+
+    if (h < 0) {
+        CLIENT_PRINT("Couldn't connect %d", h);
+        return 1;
     }
-    ERROR("ERROR: Invalid handle");
+
+    status = psa_call(h, PSA_CRYPTO_INIT, NULL, 0, NULL, 0);
+    CLIENT_PRINT("PSA_CRYPTO_INIT returned: %d", status);
+
+    CLIENT_PRINT("Closing handle");
+    psa_close(h);
+
+    if (status != PSA_SUCCESS) {
+        return 1;
+    }
     return 0;
 }
-
-static int get_queue_info(char *path, int *cqid, int *sqid)
-{
-
-    key_t server_queue_key;
-    int rx_qid, server_qid;
-
-    INFO("Attempting to contact a RoT service queue");
-
-    if ((rx_qid = msgget(IPC_PRIVATE, 0660)) == -1) {
-        ERROR("msgget: rx_qid");
-        return -1;
-    }
-
-    if ((server_queue_key = ftok(path, PROJECT_ID)) == -1) {
-        ERROR("ftok");
-        return -2;
-    }
-
-    if ((server_qid = msgget(server_queue_key, 0)) == -1) {
-        ERROR("msgget: server_qid");
-        return -3;
-    }
-
-    *cqid = rx_qid;
-    *sqid = server_qid;
-
-    return 0;
-}
-
-static psa_status_t process_response(int rx_qid, vectors_t *vecs, int type,
-                                     int *internal_server_qid)
-{
-
-    struct message response, request;
-    psa_status_t ret = PSA_ERROR_CONNECTION_REFUSED;
-    size_t invec_seek[4] = { 0 };
-    size_t data_size;
-    psa_status_t invec, outvec; /* TODO: Should these be size_t ? */
-
-    assert(internal_server_qid > 0);
-
-    while (1) {
-        data_size = 0;
-        invec = 0;
-        outvec = 0;
-
-        // read response from server
-        if (msgrcv(rx_qid, &response, sizeof(struct message_text), 0, 0) == -1) {
-            ERROR("   msgrcv failed");
-            return ret;
-        }
-
-        // process return message from server
-        switch (response.message_type) {
-            case PSA_REPLY:
-                memcpy(&ret, response.message_text.buf, sizeof(psa_status_t));
-                INFO("   Message received from server: %d", ret);
-                if (type == PSA_IPC_CONNECT && ret > 0) {
-                    *internal_server_qid = ret;
-                    INFO("   ASSSIGNED q ID %d", *internal_server_qid);
-                    ret = PSA_SUCCESS;
-                }
-                return ret;
-                break;
-            case READ_REQUEST:
-                /* read data request */
-                request.message_type = READ_RESPONSE;
-
-                assert(vecs != 0);
-
-                memcpy(&invec, response.message_text.buf, sizeof(psa_status_t));
-                memcpy(&data_size, response.message_text.buf+sizeof(size_t), sizeof(size_t));
-                INFO("   Partition asked for %lu bytes from invec %d", data_size, invec);
-
-                /* need to add more checks here */
-                assert(invec >= 0 && invec < PSA_MAX_IOVEC);
-
-                if (data_size > MAX_FRAGMENT_SIZE) {
-                    data_size = MAX_FRAGMENT_SIZE;
-                }
-
-                /* send response */
-                INFO("   invec_seek[invec] is %lu", invec_seek[invec]);
-                INFO("   Reading from offset %p", vecs->in_vec[invec].base + invec_seek[invec]);
-                memcpy(request.message_text.buf,
-                       (vecs->in_vec[invec].base + invec_seek[invec]),
-                       data_size);
-
-                /* update invec base TODO: check me */
-                invec_seek[invec] = invec_seek[invec] + data_size;
-
-                INFO("   Sending message of type %li", request.message_type);
-                INFO("       with content %s", request.message_text.buf);
-
-                if (msgsnd(*internal_server_qid, &request,
-                           sizeof(int) + sizeof(uint32_t) + data_size, 0) == -1) {
-                    ERROR("Internal error: failed to respond to read request");
-                }
-                break;
-            case WRITE_REQUEST:
-                assert(vecs != 0);
-
-                request.message_type = WRITE_RESPONSE;
-
-                memcpy(&outvec, response.message_text.buf, sizeof(psa_status_t));
-                memcpy(&data_size, response.message_text.buf + sizeof(size_t), sizeof(size_t));
-                INFO("   Partition wants to write %lu bytes to outvec %d", data_size, outvec);
-
-                assert(outvec >= 0 && outvec < PSA_MAX_IOVEC);
-
-                /* copy memory into message and send back amount written */
-                size_t sofar = vecs->out_vec[outvec].len;
-                memcpy(vecs->out_vec[outvec].base + sofar,
-                       response.message_text.buf+(sizeof(size_t)*2), data_size);
-                INFO("   Data size is %lu", data_size);
-                vecs->out_vec[outvec].len += data_size;
-
-                INFO("   Sending message of type %li", request.message_type);
-
-                /* send response */
-                if (msgsnd(*internal_server_qid, &request, sizeof(int) + data_size, 0) == -1) {
-                    ERROR("Internal error: failed to respond to write request");
-                }
-                break;
-            case SKIP_REQUEST:
-                memcpy(&invec, response.message_text.buf, sizeof(psa_status_t));
-                memcpy(&data_size, response.message_text.buf+sizeof(size_t), sizeof(size_t));
-                INFO("   Partition asked to skip %lu bytes in invec %d", data_size, invec);
-                assert(invec >= 0 && invec < PSA_MAX_IOVEC);
-                /* update invec base TODO: check me */
-                invec_seek[invec] = invec_seek[invec] + data_size;
-                break;
-
-            default:
-                FATAL("   ERROR: unknown internal message type: %ld",
-                      response.message_type);
-                return ret;
-        }
-    }
-}
-
-static psa_status_t send(int rx_qid, int server_qid, int *internal_server_qid,
-                         int32_t type, uint32_t minor_version, vectors_t *vecs)
-{
-    {
-        psa_status_t ret = PSA_ERROR_CONNECTION_REFUSED;
-        size_t request_msg_size = (sizeof(int) + sizeof(long)); /* msg type plus queue id */
-        struct message request;
-        request.message_type = 1; /* TODO: change this */
-        request.message_text.psa_type = type;
-        vector_sizes_t vec_sizes;
-
-        /* If the client is non-secure then set the NS bit */
-        if (__psa_ff_client_security_state != 0) {
-            request.message_type |= NON_SECURE;
-        }
-
-        assert(request.message_type >= 0);
-
-        INFO("SEND: Sending message of type %ld with psa_type %d", request.message_type, type);
-        INFO("     internal_server_qid = %i", *internal_server_qid);
-
-        request.message_text.qid = rx_qid;
-
-        if (type == PSA_IPC_CONNECT) {
-            memcpy(request.message_text.buf, &minor_version, sizeof(minor_version));
-            request_msg_size = request_msg_size + sizeof(minor_version);
-            INFO("   Request msg size is %lu", request_msg_size);
-        } else {
-            assert(internal_server_qid > 0);
-        }
-
-        if (vecs != NULL && type >= PSA_IPC_CALL) {
-
-            memset(&vec_sizes, 0, sizeof(vec_sizes));
-
-            /* Copy invec sizes */
-            for (size_t i = 0; i < (vecs->in_len); i++) {
-                vec_sizes.invec_sizes[i] = vecs->in_vec[i].len;
-                INFO("   Client sending vector %lu: %lu", i, vec_sizes.invec_sizes[i]);
-            }
-
-            /* Copy outvec sizes */
-            for (size_t i = 0; i < (vecs->out_len); i++) {
-                vec_sizes.outvec_sizes[i] = vecs->out_vec[i].len;
-
-                /* Reset to 0 since we need to eventually fill in with bytes written */
-                vecs->out_vec[i].len = 0;
-            }
-
-            memcpy(request.message_text.buf, &vec_sizes, sizeof(vec_sizes));
-            request_msg_size = request_msg_size + sizeof(vec_sizes);
-        }
-
-        INFO("   Sending and then waiting");
-
-        // send message to server
-        if (msgsnd(server_qid, &request, request_msg_size, 0) == -1) {
-            ERROR("   msgsnd failed");
-            return ret;
-        }
-
-        return process_response(rx_qid, vecs, type, internal_server_qid);
-    }
-}
-
-
-uint32_t psa_framework_version(void)
-{
-    return PSA_FRAMEWORK_VERSION;
-}
-
-psa_handle_t psa_connect(uint32_t sid, uint32_t minor_version)
-{
-
-    int idx;
-    psa_status_t ret;
-    char pathname[PATHNAMESIZE] = { 0 };
-
-    idx = get_next_free_handle();
-
-    /* if there's a free handle available */
-    if (idx >= 0) {
-        snprintf(pathname, PATHNAMESIZE - 1, TMP_FILE_BASE_PATH "psa_service_%u", sid);
-        INFO("Attempting to contact RoT service at %s", pathname);
-
-        /* if communication is possible */
-        if (get_queue_info(pathname, &handles[idx].client_qid, &handles[idx].server_qid) >= 0) {
-
-            ret = send(handles[idx].client_qid,
-                       handles[idx].server_qid,
-                       &handles[idx].internal_server_qid,
-                       PSA_IPC_CONNECT,
-                       minor_version,
-                       NULL);
-
-            /* if connection accepted by RoT service */
-            if (ret >= 0) {
-                handles[idx].valid = 1;
-                return idx;
-            } else {
-                INFO("Server didn't like you");
-            }
-        } else {
-            INFO("Couldn't contact RoT service. Does it exist?");
-
-            if (__psa_ff_client_security_state == 0) {
-                ERROR("Invalid SID");
-            }
-        }
-    }
-
-    INFO("Couldn't obtain a free handle");
-    return PSA_ERROR_CONNECTION_REFUSED;
-}
-
-uint32_t psa_version(uint32_t sid)
-{
-    int idx;
-    psa_status_t ret;
-    char pathname[PATHNAMESIZE] = { 0 };
-
-    idx = get_next_free_handle();
-
-    if (idx >= 0) {
-        snprintf(pathname, PATHNAMESIZE, TMP_FILE_BASE_PATH "psa_service_%u", sid);
-        if (get_queue_info(pathname, &handles[idx].client_qid, &handles[idx].server_qid) >= 0) {
-            ret = send(handles[idx].client_qid,
-                       handles[idx].server_qid,
-                       &handles[idx].internal_server_qid,
-                       VERSION_REQUEST,
-                       0,
-                       NULL);
-            INFO("psa_version: Recieved from server %d", ret);
-            if (ret > 0) {
-                return ret;
-            }
-        }
-    }
-    INFO("psa_version failed: does the service exist?");
-    return PSA_VERSION_NONE;
-}
-
-psa_status_t psa_call(psa_handle_t handle,
-                      int32_t type,
-                      const psa_invec *in_vec,
-                      size_t in_len,
-                      psa_outvec *out_vec,
-                      size_t out_len)
-{
-
-    handle_is_valid(handle);
-
-    if ((in_len + out_len) > PSA_MAX_IOVEC) {
-        ERROR("Too many iovecs: %lu + %lu", in_len, out_len);
-    }
-
-    vectors_t vecs = { 0 };
-    vecs.in_vec = in_vec;
-    vecs.in_len = in_len;
-    vecs.out_vec = out_vec;
-    vecs.out_len = out_len;
-
-    return send(handles[handle].client_qid,
-                handles[handle].server_qid,
-                &handles[handle].internal_server_qid,
-                type,
-                0,
-                &vecs);
-}
-
-void psa_close(psa_handle_t handle)
-{
-    handle_is_valid(handle);
-    if (send(handles[handle].client_qid, handles[handle].server_qid,
-             &handles[handle].internal_server_qid, PSA_IPC_DISCONNECT, 0, NULL)) {
-        ERROR("ERROR: Couldn't send disconnect msg");
-    } else {
-        if (msgctl(handles[handle].client_qid, IPC_RMID, NULL) != 0) {
-            ERROR("ERROR: Failed to delete msg queue");
-        }
-    }
-    INFO("Closing handle %u", handle);
-    handles[handle].valid = 0;
-}
diff --git a/tests/psa-client-server/psasim/src/common.c b/tests/psa-client-server/psasim/src/common.c
deleted file mode 100644
index 287bb50..0000000
--- a/tests/psa-client-server/psasim/src/common.c
+++ /dev/null
@@ -1,8 +0,0 @@
-/* Common code between clients and services */
-
-/*
- *  Copyright The Mbed TLS Contributors
- *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
- */
-
-#include "psa/common.h"
diff --git a/tests/psa-client-server/psasim/test/manifest.json b/tests/psa-client-server/psasim/src/manifest.json
similarity index 83%
rename from tests/psa-client-server/psasim/test/manifest.json
rename to tests/psa-client-server/psasim/src/manifest.json
index 0ab83ef..e67b636 100644
--- a/tests/psa-client-server/psasim/test/manifest.json
+++ b/tests/psa-client-server/psasim/src/manifest.json
@@ -3,14 +3,14 @@
    "name":"TEST_PARTITION",
    "type":"PSA-ROT",
    "priority":"LOW",
-   "entry_point":"psa_sha256_main",
+   "entry_point":"psa_server_main",
    "stack_size":"0x400",
    "heap_size":"0x100",
    "services":[
       {
-         "name":"PSA_SID_SHA256",
+         "name":"PSA_SID_CRYPTO",
          "sid":"0x0000F000",
-         "signal":"PSA_SHA256",
+         "signal":"PSA_CRYPTO",
          "non_secure_clients": "true",
          "minor_version":1,
          "minor_policy":"STRICT"
diff --git a/tests/psa-client-server/psasim/src/psa_ff_client.c b/tests/psa-client-server/psasim/src/psa_ff_client.c
new file mode 100644
index 0000000..21a43b3
--- /dev/null
+++ b/tests/psa-client-server/psasim/src/psa_ff_client.c
@@ -0,0 +1,386 @@
+/* PSA firmware framework client API */
+
+/*
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
+ */
+
+#include <stdint.h>
+#include <stdlib.h>
+#include <stddef.h>
+#include <assert.h>
+#include <stdio.h>
+#include <string.h>
+#include <strings.h>
+#include <inttypes.h>
+#include <sys/types.h>
+#include <sys/ipc.h>
+#include <sys/msg.h>
+
+#include "client.h"
+#include "common.h"
+#include "error_ext.h"
+#include "util.h"
+
+typedef struct internal_handle {
+    int server_qid;
+    int client_qid;
+    int internal_server_qid;
+    int valid;
+} internal_handle_t;
+
+typedef struct vectors {
+    const psa_invec *in_vec;
+    size_t in_len;
+    psa_outvec *out_vec;
+    size_t out_len;
+} vectors_t;
+
+/* Note that this implementation is functional and not secure */
+int __psa_ff_client_security_state = NON_SECURE;
+
+/* Access to this global is not thread safe */
+#define MAX_HANDLES 32
+static internal_handle_t handles[MAX_HANDLES] = { { 0 } };
+
+static int get_next_free_handle()
+{
+    /* Never return handle 0 as it's a special null handle */
+    for (int i = 1; i < MAX_HANDLES; i++) {
+        if (handles[i].valid == 0) {
+            return i;
+        }
+    }
+    return -1;
+}
+
+static int handle_is_valid(psa_handle_t handle)
+{
+    if (handle > 0 && handle < MAX_HANDLES) {
+        if (handles[handle].valid == 1) {
+            return 1;
+        }
+    }
+    ERROR("ERROR: Invalid handle");
+    return 0;
+}
+
+static int get_queue_info(char *path, int *cqid, int *sqid)
+{
+    key_t server_queue_key;
+    int rx_qid, server_qid;
+
+    INFO("Attempting to contact a RoT service queue");
+
+    if ((rx_qid = msgget(IPC_PRIVATE, 0660)) == -1) {
+        ERROR("msgget: rx_qid");
+        return -1;
+    }
+
+    if ((server_queue_key = ftok(path, PROJECT_ID)) == -1) {
+        ERROR("ftok");
+        return -2;
+    }
+
+    if ((server_qid = msgget(server_queue_key, 0)) == -1) {
+        ERROR("msgget: server_qid");
+        return -3;
+    }
+
+    *cqid = rx_qid;
+    *sqid = server_qid;
+
+    return 0;
+}
+
+static psa_status_t process_response(int rx_qid, vectors_t *vecs, int type,
+                                     int *internal_server_qid)
+{
+    struct message response, request;
+    psa_status_t ret = PSA_ERROR_CONNECTION_REFUSED;
+    size_t invec_seek[4] = { 0 };
+    size_t data_size;
+    psa_status_t invec, outvec; /* TODO: Should these be size_t ? */
+
+    assert(internal_server_qid > 0);
+
+    while (1) {
+        data_size = 0;
+        invec = 0;
+        outvec = 0;
+
+        /* read response from server */
+        if (msgrcv(rx_qid, &response, sizeof(struct message_text), 0, 0) == -1) {
+            ERROR("   msgrcv failed");
+            return ret;
+        }
+
+        /* process return message from server */
+        switch (response.message_type) {
+            case PSA_REPLY:
+                memcpy(&ret, response.message_text.buf, sizeof(psa_status_t));
+                INFO("   Message received from server: %d", ret);
+                if (type == PSA_IPC_CONNECT && ret > 0) {
+                    *internal_server_qid = ret;
+                    INFO("   ASSSIGNED q ID %d", *internal_server_qid);
+                    ret = PSA_SUCCESS;
+                }
+                return ret;
+                break;
+            case READ_REQUEST:
+                /* read data request */
+                request.message_type = READ_RESPONSE;
+
+                assert(vecs != 0);
+
+                memcpy(&invec, response.message_text.buf, sizeof(psa_status_t));
+                memcpy(&data_size, response.message_text.buf+sizeof(size_t), sizeof(size_t));
+                INFO("   Partition asked for %lu bytes from invec %d", data_size, invec);
+
+                /* need to add more checks here */
+                assert(invec >= 0 && invec < PSA_MAX_IOVEC);
+
+                if (data_size > MAX_FRAGMENT_SIZE) {
+                    data_size = MAX_FRAGMENT_SIZE;
+                }
+
+                /* send response */
+                INFO("   invec_seek[invec] is %lu", invec_seek[invec]);
+                INFO("   Reading from offset %p", vecs->in_vec[invec].base + invec_seek[invec]);
+                memcpy(request.message_text.buf,
+                       (vecs->in_vec[invec].base + invec_seek[invec]),
+                       data_size);
+
+                /* update invec base TODO: check me */
+                invec_seek[invec] = invec_seek[invec] + data_size;
+
+                INFO("   Sending message of type %li", request.message_type);
+                INFO("       with content %s", request.message_text.buf);
+
+                if (msgsnd(*internal_server_qid, &request,
+                           sizeof(int) + sizeof(uint32_t) + data_size, 0) == -1) {
+                    ERROR("Internal error: failed to respond to read request");
+                }
+                break;
+            case WRITE_REQUEST:
+                assert(vecs != 0);
+
+                request.message_type = WRITE_RESPONSE;
+
+                memcpy(&outvec, response.message_text.buf, sizeof(psa_status_t));
+                memcpy(&data_size, response.message_text.buf + sizeof(size_t), sizeof(size_t));
+                INFO("   Partition wants to write %lu bytes to outvec %d", data_size, outvec);
+
+                assert(outvec >= 0 && outvec < PSA_MAX_IOVEC);
+
+                /* copy memory into message and send back amount written */
+                size_t sofar = vecs->out_vec[outvec].len;
+                memcpy(vecs->out_vec[outvec].base + sofar,
+                       response.message_text.buf+(sizeof(size_t)*2), data_size);
+                INFO("   Data size is %lu", data_size);
+                vecs->out_vec[outvec].len += data_size;
+
+                INFO("   Sending message of type %li", request.message_type);
+
+                /* send response */
+                if (msgsnd(*internal_server_qid, &request, sizeof(int) + data_size, 0) == -1) {
+                    ERROR("Internal error: failed to respond to write request");
+                }
+                break;
+            case SKIP_REQUEST:
+                memcpy(&invec, response.message_text.buf, sizeof(psa_status_t));
+                memcpy(&data_size, response.message_text.buf+sizeof(size_t), sizeof(size_t));
+                INFO("   Partition asked to skip %lu bytes in invec %d", data_size, invec);
+                assert(invec >= 0 && invec < PSA_MAX_IOVEC);
+                /* update invec base TODO: check me */
+                invec_seek[invec] = invec_seek[invec] + data_size;
+                break;
+
+            default:
+                FATAL("   ERROR: unknown internal message type: %ld",
+                      response.message_type);
+                return ret;
+        }
+    }
+}
+
+static psa_status_t send(int rx_qid, int server_qid, int *internal_server_qid,
+                         int32_t type, uint32_t minor_version, vectors_t *vecs)
+{
+    psa_status_t ret = PSA_ERROR_CONNECTION_REFUSED;
+    size_t request_msg_size = (sizeof(int) + sizeof(long)); /* msg type plus queue id */
+    struct message request;
+    request.message_type = 1; /* TODO: change this */
+    request.message_text.psa_type = type;
+    vector_sizes_t vec_sizes;
+
+    /* If the client is non-secure then set the NS bit */
+    if (__psa_ff_client_security_state != 0) {
+        request.message_type |= NON_SECURE;
+    }
+
+    assert(request.message_type >= 0);
+
+    INFO("SEND: Sending message of type %ld with psa_type %d", request.message_type, type);
+    INFO("     internal_server_qid = %i", *internal_server_qid);
+
+    request.message_text.qid = rx_qid;
+
+    if (type == PSA_IPC_CONNECT) {
+        memcpy(request.message_text.buf, &minor_version, sizeof(minor_version));
+        request_msg_size = request_msg_size + sizeof(minor_version);
+        INFO("   Request msg size is %lu", request_msg_size);
+    } else {
+        assert(internal_server_qid > 0);
+    }
+
+    if (vecs != NULL && type >= PSA_IPC_CALL) {
+
+        memset(&vec_sizes, 0, sizeof(vec_sizes));
+
+        /* Copy invec sizes */
+        for (size_t i = 0; i < (vecs->in_len); i++) {
+            vec_sizes.invec_sizes[i] = vecs->in_vec[i].len;
+            INFO("   Client sending vector %lu: %lu", i, vec_sizes.invec_sizes[i]);
+        }
+
+        /* Copy outvec sizes */
+        for (size_t i = 0; i < (vecs->out_len); i++) {
+            vec_sizes.outvec_sizes[i] = vecs->out_vec[i].len;
+
+            /* Reset to 0 since we need to eventually fill in with bytes written */
+            vecs->out_vec[i].len = 0;
+        }
+
+        memcpy(request.message_text.buf, &vec_sizes, sizeof(vec_sizes));
+        request_msg_size = request_msg_size + sizeof(vec_sizes);
+    }
+
+    INFO("   Sending and then waiting");
+
+    /* send message to server */
+    if (msgsnd(server_qid, &request, request_msg_size, 0) == -1) {
+        ERROR("   msgsnd failed");
+        return ret;
+    }
+
+    return process_response(rx_qid, vecs, type, internal_server_qid);
+}
+
+
+uint32_t psa_framework_version(void)
+{
+    return PSA_FRAMEWORK_VERSION;
+}
+
+psa_handle_t psa_connect(uint32_t sid, uint32_t minor_version)
+{
+    int idx;
+    psa_status_t ret;
+    char pathname[PATHNAMESIZE] = { 0 };
+
+    idx = get_next_free_handle();
+
+    /* if there's a free handle available */
+    if (idx >= 0) {
+        snprintf(pathname, PATHNAMESIZE - 1, TMP_FILE_BASE_PATH "psa_service_%u", sid);
+        INFO("Attempting to contact RoT service at %s", pathname);
+
+        /* if communication is possible */
+        if (get_queue_info(pathname, &handles[idx].client_qid, &handles[idx].server_qid) >= 0) {
+
+            ret = send(handles[idx].client_qid,
+                       handles[idx].server_qid,
+                       &handles[idx].internal_server_qid,
+                       PSA_IPC_CONNECT,
+                       minor_version,
+                       NULL);
+
+            /* if connection accepted by RoT service */
+            if (ret >= 0) {
+                handles[idx].valid = 1;
+                return idx;
+            } else {
+                INFO("Server didn't like you");
+            }
+        } else {
+            INFO("Couldn't contact RoT service. Does it exist?");
+
+            if (__psa_ff_client_security_state == 0) {
+                ERROR("Invalid SID");
+            }
+        }
+    }
+
+    INFO("Couldn't obtain a free handle");
+    return PSA_ERROR_CONNECTION_REFUSED;
+}
+
+uint32_t psa_version(uint32_t sid)
+{
+    int idx;
+    psa_status_t ret;
+    char pathname[PATHNAMESIZE] = { 0 };
+
+    idx = get_next_free_handle();
+
+    if (idx >= 0) {
+        snprintf(pathname, PATHNAMESIZE, TMP_FILE_BASE_PATH "psa_service_%u", sid);
+        if (get_queue_info(pathname, &handles[idx].client_qid, &handles[idx].server_qid) >= 0) {
+            ret = send(handles[idx].client_qid,
+                       handles[idx].server_qid,
+                       &handles[idx].internal_server_qid,
+                       VERSION_REQUEST,
+                       0,
+                       NULL);
+            INFO("psa_version: Recieved from server %d", ret);
+            if (ret > 0) {
+                return ret;
+            }
+        }
+    }
+    INFO("psa_version failed: does the service exist?");
+    return PSA_VERSION_NONE;
+}
+
+psa_status_t psa_call(psa_handle_t handle,
+                      int32_t type,
+                      const psa_invec *in_vec,
+                      size_t in_len,
+                      psa_outvec *out_vec,
+                      size_t out_len)
+{
+    handle_is_valid(handle);
+
+    if ((in_len + out_len) > PSA_MAX_IOVEC) {
+        ERROR("Too many iovecs: %lu + %lu", in_len, out_len);
+    }
+
+    vectors_t vecs = { 0 };
+    vecs.in_vec = in_vec;
+    vecs.in_len = in_len;
+    vecs.out_vec = out_vec;
+    vecs.out_len = out_len;
+
+    return send(handles[handle].client_qid,
+                handles[handle].server_qid,
+                &handles[handle].internal_server_qid,
+                type,
+                0,
+                &vecs);
+}
+
+void psa_close(psa_handle_t handle)
+{
+    handle_is_valid(handle);
+    if (send(handles[handle].client_qid, handles[handle].server_qid,
+             &handles[handle].internal_server_qid, PSA_IPC_DISCONNECT, 0, NULL)) {
+        ERROR("ERROR: Couldn't send disconnect msg");
+    } else {
+        if (msgctl(handles[handle].client_qid, IPC_RMID, NULL) != 0) {
+            ERROR("ERROR: Failed to delete msg queue");
+        }
+    }
+    INFO("Closing handle %u", handle);
+    handles[handle].valid = 0;
+}
diff --git a/tests/psa-client-server/psasim/src/service.c b/tests/psa-client-server/psasim/src/psa_ff_server.c
similarity index 98%
rename from tests/psa-client-server/psasim/src/service.c
rename to tests/psa-client-server/psasim/src/psa_ff_server.c
index b2b6a08..ea797d8 100644
--- a/tests/psa-client-server/psasim/src/service.c
+++ b/tests/psa-client-server/psasim/src/psa_ff_server.c
@@ -16,11 +16,11 @@
 #include <time.h>
 #include <assert.h>
 
-#include "psa/service.h"
-#include "psasim/init.h"
-#include "psa/error.h"
-#include "psa/common.h"
-#include "psa/util.h"
+#include "service.h"
+#include "init.h"
+#include "error_ext.h"
+#include "common.h"
+#include "util.h"
 
 #define MAX_CLIENTS 128
 #define MAX_MESSAGES 32
@@ -34,7 +34,7 @@
 };
 
 /* Note that this implementation is functional and not secure. */
-extern int __psa_ff_client_security_state;
+int __psa_ff_client_security_state = NON_SECURE;
 
 static psa_msg_t messages[MAX_MESSAGES]; /* Message slots */
 static uint8_t pending_message[MAX_MESSAGES] = { 0 }; /* Booleans indicating active message slots */
diff --git a/tests/psa-client-server/psasim/src/psa_functions_codes.h b/tests/psa-client-server/psasim/src/psa_functions_codes.h
new file mode 100644
index 0000000..34897b9
--- /dev/null
+++ b/tests/psa-client-server/psasim/src/psa_functions_codes.h
@@ -0,0 +1,9 @@
+#ifndef _PSA_FUNCTIONS_CODES_H_
+#define  _PSA_FUNCTIONS_CODES_H_
+
+enum {
+    PSA_CRYPTO_INIT = 0x00,
+    /* Add other PSA functions here */
+};
+
+#endif /*  _PSA_FUNCTIONS_CODES_H_ */
diff --git a/tests/psa-client-server/psasim/test/server.c b/tests/psa-client-server/psasim/src/server.c
similarity index 68%
rename from tests/psa-client-server/psasim/test/server.c
rename to tests/psa-client-server/psasim/src/server.c
index c4b6d9c..630bd73 100644
--- a/tests/psa-client-server/psasim/test/server.c
+++ b/tests/psa-client-server/psasim/src/server.c
@@ -8,10 +8,16 @@
 #include <unistd.h>
 #include <stdio.h>
 
-#include "psa/service.h"
-#include "psa/error.h"
-#include "psa/util.h"
+/* Includes from psasim */
+#include "service.h"
+#include "error_ext.h"
+#include "util.h"
 #include "psa_manifest/manifest.h"
+#include "psa_functions_codes.h"
+
+/* Includes from mbedtls */
+#include "mbedtls/version.h"
+#include "psa/crypto.h"
 
 #define SERVER_PRINT(fmt, ...) \
     PRINT("Server: " fmt, ##__VA_ARGS__)
@@ -36,13 +42,16 @@
     }
 }
 
-int psa_sha256_main(int argc, char *argv[])
+int psa_server_main(int argc, char *argv[])
 {
     psa_status_t ret = PSA_ERROR_PROGRAMMER_ERROR;
     psa_msg_t msg = { -1 };
-    char foo[BUF_SIZE] = { 0 };
     const int magic_num = 66;
     int client_disconnected = 0;
+    char mbedtls_version[18];
+
+    mbedtls_version_get_string_full(mbedtls_version);
+    SERVER_PRINT("%s", mbedtls_version);
 
     parse_input_args(argc, argv);
     SERVER_PRINT("Starting");
@@ -54,10 +63,9 @@
             SERVER_PRINT("Signals: 0x%08x", signals);
         }
 
-        if (signals & PSA_SHA256_SIGNAL) {
-            if (PSA_SUCCESS == psa_get(PSA_SHA256_SIGNAL, &msg)) {
-                SERVER_PRINT("My handle is %d", msg.handle);
-                SERVER_PRINT("My rhandle is %p", (int *) msg.rhandle);
+        if (signals & PSA_CRYPTO_SIGNAL) {
+            if (PSA_SUCCESS == psa_get(PSA_CRYPTO_SIGNAL, &msg)) {
+                SERVER_PRINT("handle: %d - rhandle: %p", msg.handle, (int *) msg.rhandle);
                 switch (msg.type) {
                     case PSA_IPC_CONNECT:
                         SERVER_PRINT("Got a connection message");
@@ -69,34 +77,23 @@
                         ret = PSA_SUCCESS;
                         client_disconnected = 1;
                         break;
-
                     default:
                         SERVER_PRINT("Got an IPC call of type %d", msg.type);
-                        ret = 42;
-                        size_t size = msg.in_size[0];
-
-                        if ((size > 0) && (size <= sizeof(foo))) {
-                            psa_read(msg.handle, 0, foo, 6);
-                            foo[(BUF_SIZE-1)] = '\0';
-                            SERVER_PRINT("Reading payload: %s", foo);
-                            psa_read(msg.handle, 0, foo+6, 6);
-                            foo[(BUF_SIZE-1)] = '\0';
-                            SERVER_PRINT("Reading payload: %s", foo);
+                        switch (msg.type) {
+                            case PSA_CRYPTO_INIT:
+                                ret = psa_crypto_init();
+                                break;
+                            default:
+                                SERVER_PRINT("Unknown PSA function code");
+                                break;
                         }
-
-                        size = msg.out_size[0];
-                        if ((size > 0)) {
-                            SERVER_PRINT("Writing response");
-                            psa_write(msg.handle, 0, "RESP", 4);
-                            psa_write(msg.handle, 0, "ONSE", 4);
-                        }
+                        SERVER_PRINT("Internal function call returned %d", ret);
 
                         if (msg.client_id > 0) {
                             psa_notify(msg.client_id);
                         } else {
                             SERVER_PRINT("Client is non-secure, so won't notify");
                         }
-
                 }
 
                 psa_reply(msg.handle, ret);
diff --git a/tests/psa-client-server/psasim/test/Makefile b/tests/psa-client-server/psasim/test/Makefile
deleted file mode 100644
index 34b86b6..0000000
--- a/tests/psa-client-server/psasim/test/Makefile
+++ /dev/null
@@ -1,29 +0,0 @@
-INCLUDE := -I../include/ -I./psa_manifest
-LIB := -L../src -lpsaff
-
-TEST_BIN = 	psa_client \
-			psa_partition
-
-GENERATED_H_FILES =	psa_manifest/manifest.h \
-					psa_manifest/pid.h \
-					psa_manifest/sid.h
-
-PARTITION_SERVER_BOOTSTRAP = psa_ff_bootstrap_TEST_PARTITION.c
-
-.PHONY: all clean
-
-all: $(TEST_BIN)
-
-psa_client: client.c $(GENERATED_H_FILES)
-	$(CC) $(INCLUDE) $(CFLAGS) $< $(LIB) -o $@
-
-psa_partition: $(PARTITION_SERVER_BOOTSTRAP) $(GENERATED_H_FILES)
-	$(CC) $(INCLUDE) $(CFLAGS) $< $(LIB) -o $@
-
-$(PARTITION_SERVER_BOOTSTRAP) $(GENERATED_H_FILES): manifest.json server.c
-	../tools/psa_autogen.py $<
-
-clean:
-	rm -f $(TEST_BIN) psa_ff_bootstrap_*.c
-	rm -f psa_notify_* psa_service_*
-	rm -f psa_manifest/*
diff --git a/tests/psa-client-server/psasim/test/client.c b/tests/psa-client-server/psasim/test/client.c
deleted file mode 100644
index 5bde82f..0000000
--- a/tests/psa-client-server/psasim/test/client.c
+++ /dev/null
@@ -1,48 +0,0 @@
-/* psasim test client */
-
-/*
- *  Copyright The Mbed TLS Contributors
- *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
- */
-
-#include <psa/client.h>
-#include <psa/util.h>
-#include "psa_manifest/sid.h"
-#include <stdio.h>
-#include <unistd.h>
-
-#define CLIENT_PRINT(fmt, ...) \
-    PRINT("Client: " fmt, ##__VA_ARGS__)
-
-int main()
-{
-    const char *text = "FOOBARCOOL!!";
-    char output[100] = { 0 };
-    CLIENT_PRINT("My PID: %d", getpid());
-
-    CLIENT_PRINT("PSA version: %u", psa_version(PSA_SID_SHA256_SID));
-    psa_handle_t h = psa_connect(PSA_SID_SHA256_SID, 1);
-
-    if (h < 0) {
-        CLIENT_PRINT("Couldn't connect %d", h);
-        return 1;
-    } else {
-        int type = 2;
-        CLIENT_PRINT("psa_call() w/o invec returned: %d", psa_call(h, type, NULL, 0, NULL, 0));
-        psa_invec invecs[1];
-        psa_outvec outvecs[1];
-        invecs[0].base = text;
-        invecs[0].len = sizeof(text);
-        outvecs[0].base = output;
-        outvecs[0].len = sizeof(output);
-
-        CLIENT_PRINT("invec len: %lu", invecs[0].len);
-        CLIENT_PRINT("psa_call() w/ invec returned: %d", psa_call(h, type, invecs, 1, outvecs, 1));
-        CLIENT_PRINT("Received payload len: %ld", outvecs[0].len);
-        CLIENT_PRINT("Received payload content: %s", output);
-        CLIENT_PRINT("Closing handle");
-        psa_close(h);
-    }
-
-    return 0;
-}
diff --git a/tests/psa-client-server/psasim/test/run_test.sh b/tests/psa-client-server/psasim/test/run_test.sh
index f0e7a62..6a5605f 100755
--- a/tests/psa-client-server/psasim/test/run_test.sh
+++ b/tests/psa-client-server/psasim/test/run_test.sh
@@ -11,7 +11,10 @@
 
 set -e
 
+cd "$(dirname "$0")"
+
 function clean_run() {
+    rm -f psa_notify_*
     pkill psa_partition || true
     pkill psa_client || true
     ipcs | grep q | awk '{ printf " -q " $$2 }' | xargs ipcrm > /dev/null 2>&1 || true
@@ -21,14 +24,14 @@
 # event as signal that the server is ready so that we can start client(s).
 function wait_for_server_startup() {
     while [ ! -f ./psa_notify_* ]; do
-    sleep 0.1
+        sleep 0.1
     done
 }
 
 clean_run
 
-./psa_partition -k &
+./psa_partition -k > psa_partition.log 2>&1 &
 SERV_PID=$!
 wait_for_server_startup
-./psa_client
+./psa_client > psa_client.log 2>&1
 wait $SERV_PID
diff --git a/tests/psa-client-server/psasim/tools/psa_autogen.py b/tests/psa-client-server/psasim/tools/psa_autogen.py
index 53b1fea..fbc9806 100755
--- a/tests/psa-client-server/psasim/tools/psa_autogen.py
+++ b/tests/psa-client-server/psasim/tools/psa_autogen.py
@@ -15,6 +15,13 @@
 
 FILENAME = str(sys.argv[1])
 
+SCRIPT_PATH = os.path.dirname(__file__)
+GENERATED_H_PATH = os.path.join(SCRIPT_PATH, "..", "include", "psa_manifest")
+GENERATED_C_PATH = os.path.join(SCRIPT_PATH, "..", "src")
+
+MANIFEST_FILE = os.path.join(GENERATED_H_PATH, "manifest.h")
+PID_FILE = os.path.join(GENERATED_H_PATH, "pid.h")
+SID_FILE = os.path.join(GENERATED_H_PATH, "sid.h")
 
 with open(str(FILENAME), "r") as read_file:
     data = json.load(read_file)
@@ -32,14 +39,14 @@
             irqs = []
 
         try:
-            os.mkdir("psa_manifest")
+            os.mkdir(GENERATED_H_PATH)
             print("Generating psa_manifest directory")
         except OSError:
-            print ("PSA manifest directory already exists")
+            print("PSA manifest directory already exists")
 
-        man  = open(str("psa_manifest/" + FILENAME + ".h"), "w")
-        pids = open("psa_manifest/pid.h", "a")
-        sids = open("psa_manifest/sid.h", "a")
+        manifest_content = []
+        pids_content = []
+        sids_content = []
 
         if len(services) > 28:
            print ("Unsupported number of services")
@@ -60,8 +67,8 @@
 
         # Go through all the services to make sid.h and pid.h
         for svc in services:
-            man.write("#define {}_SIGNAL    0x{:08x}\n".format(svc['signal'], 2**count))
-            sids.write("#define {}_SID    {}\n".format(svc['name'], svc['sid']))
+            manifest_content.append("#define {}_SIGNAL    0x{:08x}".format(svc['signal'], 2**count))
+            sids_content.append("#define {}_SID    {}".format(svc['name'], svc['sid']))
             qcode = qcode + "\"" + queue_path + str(int(svc['sid'], 16)) + "\","
             ns_clients = svc['non_secure_clients']
             print(str(svc))
@@ -91,7 +98,7 @@
         handlercode = "void __sig_handler(int signo) {\n"
         irqcount = count
         for irq in irqs:
-            man.write("#define {}    0x{:08x}\n".format(irq['signal'], 2**irqcount))
+            manifest_content.append("#define {}    0x{:08x}".format(irq['signal'], 2**irqcount))
             sigcode = sigcode + "    signal({}, __sig_handler);\n".format(irq['source'])
             handlercode = handlercode + \
                           "    if (signo == {}) {{ raise_signal(0x{:08x}); }};\n".format(irq['source'], 2**irqcount)
@@ -111,28 +118,28 @@
         versions = versions + "};\n"
         policy = policy + "};\n"
 
-        pids.close()
-        sids.close()
-        man.close()
+        with open(MANIFEST_FILE, "wt") as output:
+            output.write("\n".join(manifest_content))
+        with open(SID_FILE, "wt") as output:
+            output.write("\n".join(sids_content))
+        with open(PID_FILE, "wt") as output:
+            output.write("\n".join(pids_content))
 
         symbols = []
-        # Go through all the files in the current directory and look for the entrypoint
 
-        for root, directories, filenames in os.walk('.'):
+        # Go through source files and look for the entrypoint
+        for root, directories, filenames in os.walk(GENERATED_C_PATH):
             for filename in filenames:
-
                 if "psa_ff_bootstrap" in filename or filename == "psa_manifest":
                     continue
-
                 try:
                     fullpath = os.path.join(root,filename)
                     with open(fullpath, encoding='utf-8') as currentFile:
                         text = currentFile.read()
                         if str(entry_point + "(") in text:
-                            symbols.append(fullpath)
+                            symbols.append(filename)
                 except IOError:
                     print("Couldn't open " + filename)
-
                 except UnicodeDecodeError:
                     pass
 
@@ -144,22 +151,24 @@
             print("Duplicate entrypoint symbol detected: " + str(symbols))
             sys.exit(2)
         else:
-            bs = open(str("psa_ff_bootstrap_" + str(partition_name) + ".c"), "w")
-            bs.write("#include <psasim/init.h>\n")
-            bs.write("#include \"" + symbols[0] + "\"\n")
-            bs.write("#include <signal.h>\n\n")
-            bs.write(qcode)
-            bs.write(nsacl)
-            bs.write(policy)
-            bs.write(versions)
-            bs.write("\n")
-            bs.write(handlercode)
-            bs.write("\n")
-            bs.write("int main(int argc, char *argv[]) {\n")
-            bs.write("    (void) argc;\n")
-            bs.write(sigcode)
-            bs.write("    __init_psasim(psa_queues, 32, ns_allowed, versions, strict_policy);\n")
-            bs.write("    " + entry_point + "(argc, argv);\n}\n")
-            bs.close()
+            C_FILENAME = os.path.join(GENERATED_C_PATH, "psa_ff_bootstrap_" + partition_name + ".c")
+            c_content = []
+            c_content.append("#include <init.h>")
+            c_content.append("#include \"" + symbols[0] + "\"")
+            c_content.append("#include <signal.h>")
+            c_content.append(qcode)
+            c_content.append(nsacl)
+            c_content.append(policy)
+            c_content.append(versions)
+            c_content.append(handlercode)
+            c_content.append("int main(int argc, char *argv[]) {")
+            c_content.append("    (void) argc;")
+            c_content.append(sigcode)
+            c_content.append("    __init_psasim(psa_queues, 32, ns_allowed, versions,"
+                             "strict_policy);")
+            c_content.append("    " + entry_point + "(argc, argv);")
+            c_content.append("}")
+            with open(C_FILENAME, "wt") as output:
+                output.write("\n".join(c_content))
 
             print("Success")
diff --git a/tests/scripts/all.sh b/tests/scripts/all.sh
index 8158c8d..0bdfb6f 100755
--- a/tests/scripts/all.sh
+++ b/tests/scripts/all.sh
@@ -945,6 +945,39 @@
     make CC=$ASAN_CC CFLAGS="$ASAN_CFLAGS -I../tests/include -I../tests -I../../tests -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_TEST_LIBTESTDRIVER1 $loc_accel_flags" LDFLAGS="-ltestdriver1 $ASAN_CFLAGS" "$@"
 }
 
+# $1: target which can be "client" or "server"
+helper_crypto_client_build() {
+    TARGET=$1
+    shift
+    TARGET_LIB=libpsa$TARGET
+
+    if [ "$TARGET" == "client" ]; then
+        scripts/config.py full
+        scripts/config.py unset MBEDTLS_PSA_CRYPTO_C
+        scripts/config.py unset MBEDTLS_PSA_CRYPTO_STORAGE_C
+        # Dynamic secure element support is a deprecated feature and it is not
+        # available when CRYPTO_C and PSA_CRYPTO_STORAGE_C are disabled.
+        scripts/config.py unset MBEDTLS_PSA_CRYPTO_SE_C
+        # Disable potentially problematic features
+        scripts/config.py unset MBEDTLS_X509_RSASSA_PSS_SUPPORT
+        scripts/config.py unset MBEDTLS_KEY_EXCHANGE_DHE_RSA_ENABLED
+        scripts/config.py unset MBEDTLS_KEY_EXCHANGE_ECDHE_RSA_ENABLED
+        scripts/config.py unset MBEDTLS_KEY_EXCHANGE_ECDH_RSA_ENABLED
+        scripts/config.py unset MBEDTLS_ECP_RESTARTABLE
+    else
+        scripts/config.py crypto_full
+        scripts/config.py unset MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS
+        scripts/config.py set MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER
+    fi
+
+    make -C tests/psa-client-server/psasim/ CFLAGS="$ASAN_CFLAGS" LDFLAGS="$ASAN_CFLAGS" $TARGET_LIB "$@"
+
+    # cleanup() will restore some backed-up files which include $CONFIG_H and
+    # $CRYPTO_CONFIG_H. Built libraries were already copied to psasim at this
+    # point.
+    cleanup
+}
+
 ################################################################
 #### Configuration helpers
 ################################################################
@@ -6175,11 +6208,21 @@
 }
 
 component_test_psasim() {
+    msg "build library for client"
+
+    helper_crypto_client_build client
+
+    msg "build library for server"
+
+    scripts/config.py crypto
+
+    helper_crypto_client_build server
+
     msg "build psasim"
-    make -C tests/psa-client-server/psasim
+    make -C tests/psa-client-server/psasim CFLAGS="$ASAN_CFLAGS" LDFLAGS="$ASAN_CFLAGS"
 
     msg "test psasim"
-    make -C tests/psa-client-server/psasim run
+    tests/psa-client-server/psasim/test/run_test.sh
 
     msg "clean psasim"
     make -C tests/psa-client-server/psasim clean