2 files changed, 42 insertions, 22 deletions

diff --git a/src/tools/echo/echo_client.c b/src/tools/echo/echo_client.c
index 8d3fc322..9cf56cee 100644
--- a/src/tools/echo/echo_client.c
+++ b/src/tools/echo/echo_client.c
@@ -23,19 +23,25 @@
 #define CLIENT_AP_NAME "echo-client"
 
 #include <ouroboros/dev.h>
+#include <stdlib.h>
 
 int client_main()
 {
         int fd = 0;
         int result = 0;
-        uint8_t buf[BUF_SIZE];
+        char buf[BUF_SIZE];
         char * message  = "Client says hi!";
         ssize_t count = 0;
 
-        fd = flow_alloc(SERVER_AP_NAME, CLIENT_AP_NAME,
-                        NULL, NULL, 0);
+        if(ap_init(CLIENT_AP_NAME)) {
+                printf("Failed to init AP.");
+                return -1;
+        }
+
+        fd = flow_alloc(SERVER_AP_NAME, NULL, NULL);
         if (fd < 0) {
                 printf("Failed to allocate flow\n");
+                ap_fini();
                 return -1;
         }
 
@@ -43,12 +49,14 @@ int client_main()
         if (result < 0) {
                 printf("Flow allocation refused\n");
                 flow_dealloc(fd);
+                ap_fini();
                 return -1;
         }
 
         if (flow_write(fd, message, strlen(message) + 1) == -1) {
                 printf("Failed to write SDU\n");
                 flow_dealloc(fd);
+                ap_fini();
                 return -1;
         }
 
@@ -56,6 +64,7 @@ int client_main()
         if (count < 0) {
                 printf("Failed to read SDU\n");
                 flow_dealloc(fd);
+                ap_fini();
                 return -1;
         }
 
@@ -63,5 +72,7 @@ int client_main()
 
         flow_dealloc(fd);
 
+        ap_fini();
+
         return 0;
 }
diff --git a/src/tools/echo/echo_server.c b/src/tools/echo/echo_server.c
index e457e22b..d9af1c1f 100644
--- a/src/tools/echo/echo_server.c
+++ b/src/tools/echo/echo_server.c
@@ -33,65 +33,72 @@ void shutdown_server(int signo)
 {
         char * dif = DIF_NAME;
 
-        if (ap_unreg(SERVER_AP_NAME, &dif, 1)) {
-                printf("Failed to unregister application\n");
+        if (ap_unreg(&dif, 1)) {
+                printf("Failed to unregister application.\n");
+                ap_fini();
                 exit(EXIT_FAILURE);
         }
 
+        ap_fini();
         exit(EXIT_SUCCESS);
 }
 
 int server_main()
 {
-        int server_fd = 0;
-        int client_fd = 0;
+        int    server_fd = 0;
+        int    client_fd = 0;
         char * dif = DIF_NAME;
         char * client_name = NULL;
-        uint8_t buf[BUF_SIZE];
+        char   buf[BUF_SIZE];
         ssize_t count = 0;
 
-        printf("Starting the server\n");
+        printf("Starting the server.\n");
 
         /* Manual cleanup is required for now */
         if (signal(SIGINT, shutdown_server) == SIG_ERR) {
-                printf("Can't install signal handler\n");
+                printf("Can't install signal handler.\n");
                 return -1;
         }
 
-        server_fd = ap_reg(SERVER_AP_NAME, &dif, 1);
-        if (server_fd < 0) {
-                printf("Failed to register application\n");
+        if(ap_init(SERVER_AP_NAME)) {
+                printf("Failed to init AP.");
                 return -1;
         }
 
-        printf("Echo server started...\n");
+        server_fd = ap_reg(&dif, 1);
+        if (server_fd < 0) {
+                printf("Failed to register application.\n");
+                ap_fini();
+                return -1;
+        }
 
         while (true) {
                 client_fd = flow_accept(server_fd,
-                                        client_name, NULL);
+                                        &client_name, NULL);
                 if (client_fd < 0) {
-                        continue;
+                        printf("Failed to accept flow.\n");
+                        break;
                 }
 
-                printf("New flow from %s\n", client_name);
+                printf("New flow from %s.\n", client_name);
 
                 if (flow_alloc_resp(client_fd, 0)) {
-                        printf("Failed to give an allocate response\n");
+                        printf("Failed to give an allocate response.\n");
                         flow_dealloc(client_fd);
                         continue;
                 }
 
-                count = flow_read(client_fd, buf, BUF_SIZE);
+                count = flow_read(client_fd, (void **) &buf, BUF_SIZE);
                 if (count < 0) {
-                        printf("Failed to read SDU\n");
+                        printf("Failed to read SDU.\n");
                         flow_dealloc(client_fd);
                         continue;
                 }
 
-                printf("Message from client is %.*s\n", (int) count, buf);
+                printf("Message from client is %.*s.\n", (int) count, buf);
 
                 if (flow_write(client_fd, buf, count) == -1) {
-                        printf("Failed to write SDU\n");
+                        printf("Failed to write SDU.\n");
                         flow_dealloc(client_fd);
                         continue;
                 }
@@ -99,5 +106,7 @@ int server_main()
                 flow_dealloc(client_fd);
         }
 
+        ap_fini();
+
         return 0;
 }