diff --git a/src/main.h b/src/main.h
index 1216283..b5ef1b9 100644
--- a/src/main.h
+++ b/src/main.h
@@ -37,6 +37,7 @@ enum state {
   S_WAIT_SERVER,
   S_PRINT_SERVER_MESSAGE,
   S_FILE_LIST,
+  S_WAIT_PASS,
   /* S_FILE_SELECT, */
   S_FILE_DOWNLOAD,
   S_UPLOAD_PARAMS,
@@ -88,7 +89,6 @@ typedef struct query_args {
   char *buf;
   bool from_server;
   params_t *params;
-  char *next_server_command;
   struct {
     char *text;
     uint64_t capacity;
diff --git a/src/query.c b/src/query.c
index 186210c..e561708 100644
--- a/src/query.c
+++ b/src/query.c
@@ -74,13 +74,7 @@ void query_loop(app_t *app) {
           maxfd = query_args->file->fd;
       }
     }
-    if (query_args->state == S_NEXT_ACTION &&
-        query_args->next_server_command != NULL) {
-      query_args->buf_used = strlen(query_args->next_server_command);
-      memcpy(query_args->buf, query_args->next_server_command,
-             query_args->buf_used);
-      free(query_args->next_server_command);
-      query_args->next_server_command = NULL;
+    if (query_args->state == S_NEXT_ACTION) {
       query_args->from_server = TRUE;
       if (ERR == process_query(app)) {
         EXIT_APP("*** error while processing query ***", app, 1)
@@ -133,7 +127,7 @@ void wait_register(query_args_t *q_args) {
                      email);
 }
 
-int upload_confirm_cb(app_t *app, char *query) {
+int upload_confirm_cb(app_t *app, char *query, int q_len) {
   if (!strncmp("finished\n", query, sizeof("finished\n") - 1)) {
     ui_file_list_t *fui = (ui_file_list_t *)app->query_args->file_list_ui;
     notification("File upload", dc_normal, "File %s is uploaded to the server!",
@@ -157,12 +151,16 @@ int process_query(app_t *app) {
   case WAIT_SERVER:
   case WAIT_SERVER_INIT:
   case WAIT_CLIENT:
-  case S_NEXT_ACTION:
-  case S_WAIT_SERVER:
   case WAIT_REGISTER_CONFIRMATION:
-  case S_WAIT_REGISTER_CONFIRMATION:
+  case S_WAIT_PASS:
     wait_side(app, NULL);
     break;
+  case S_NEXT_ACTION:
+    wait_side(app, NULL);
+    break;
+  case S_WAIT_REGISTER_CONFIRMATION:
+    wait_side(app, server_wait_reg_confirm_cb);
+    break;
   case S_ERR:
     break;
   case S_FILE_LIST:
@@ -196,6 +194,9 @@ int process_query(app_t *app) {
   case WAIT_REGISTER:
     wait_register(query_args);
     break;
+  case S_WAIT_SERVER:
+    wait_side(app, server_print_message_cb);
+    break;
   default:
     break;
   }
@@ -220,7 +221,7 @@ int query_extract_from_buf(char *buf, int *buf_used, char **output_line) {
     line = malloc(b_used + 1);
     strncpy(line, buf, b_used);
 
-    // line[b_used] = 0;
+    line[b_used] = 0;
 
     *output_line = line;
     return b_used;
@@ -240,21 +241,32 @@ int query_extract_from_buf(char *buf, int *buf_used, char **output_line) {
   return pos + 1;
 }
 
+void query_return_to_buf(char *buf, int *buf_used, char *input_line) {
+  ssize_t l_len = strlen(input_line);
+  memmove(buf + l_len, buf, l_len);
+  memcpy(buf, input_line, l_len);
+  *buf_used = l_len;
+  buf[*buf_used] = 0;
+}
+
 static int wait_side(app_t *app, wait_server_cb *callback) {
   query_args_t *q_args = app->query_args;
   int qlen;
-  int buf_used = q_args->buf_used;
   char *query = NULL;
   int res = 0;
 
-  while ((qlen = query_extract_from_buf(q_args->buf, &buf_used, &query))) {
+  while ((qlen = query_extract_from_buf(q_args->buf, &(q_args->buf_used),
+                                        &query))) {
     if (callback == NULL) {
       process_server_command(query, qlen, app);
     } else {
-      res = callback(app, query);
+      res = callback(app, query, qlen);
     }
     free(query);
     query = NULL;
+    if (res == 1) {
+      break;
+    }
   }
   return res;
 }
@@ -270,7 +282,6 @@ void init_query_args(query_args_t *q_args, params_t *params) {
   q_args->server_message.text = NULL;
   q_args->server_message.capacity = 0;
   q_args->server_message.size = 0;
-  q_args->next_server_command = NULL;
   q_args->progress_bar = NULL;
   q_args->active_dialogue = NULL;
   q_args->notification = NULL;
diff --git a/src/query.h b/src/query.h
index 7d533a5..c3c1a68 100644
--- a/src/query.h
+++ b/src/query.h
@@ -13,6 +13,6 @@ void wait_register(query_args_t *q_args);
 int process_query(app_t *app);
 int query_extract_from_buf(char *buf, int *buf_used, char **output_line);
 void init_query_args(query_args_t *q_args, params_t *params);
-typedef int wait_server_cb(app_t *app, char *query);
-
+typedef int wait_server_cb(app_t *app, char *query, int q_len);
+void query_return_to_buf(char *buf, int *buf_used, char *input_line);
 #endif
\ No newline at end of file
diff --git a/src/server.c b/src/server.c
index ff753e6..ec0ebf5 100644
--- a/src/server.c
+++ b/src/server.c
@@ -4,6 +4,7 @@
 #include "app.h"
 #include "dialogue.h"
 #include "main.h"
+#include "query.h"
 #include <arpa/inet.h>
 #include <fcntl.h>
 #include <modals.h>
@@ -21,15 +22,6 @@
 void ask_uname_and_password(params_t *params);
 void ask_register(params_t *params, char *email);
 
-/* TODO: questionable solution. Do better  */
-#define PRINT_SRV_MESSAGE(q_args, l_len, line)                                 \
-  if (q_args->server_message.size > 0) {                                       \
-    q_args->state = S_PRINT_SERVER_MESSAGE;                                    \
-    q_args->next_server_command = malloc(l_len + 1);                           \
-    strcpy(q_args->next_server_command, line);                                 \
-    return 0;                                                                  \
-  }
-
 void server_send_string(query_args_t *q, const char *fmt, ...) {
   char buf[INBUFSIZE];
   size_t b_len = 0;
@@ -40,14 +32,72 @@ void server_send_string(query_args_t *q, const char *fmt, ...) {
   va_end(args);
 }
 
+int server_wait_reg_confirm_cb(app_t *app, char *line, int l_len) {
+  query_args_t *q_args = app->query_args;
+  char query[INBUFSIZE];
+  ui_file_list_t *fui = (ui_file_list_t *)q_args->file_list_ui;
+
+  if (!strcmp(line, "ok\n")) {
+    server_send_string(q_args, "file list %u %u\n", fui->max_lines,
+                       fui->current_page);
+    q_args->state = S_FILE_LIST;
+    destroy_dialogue(&(app->modal), app);
+    sprintf(query,
+            "You've been successfully registered.\n"
+            "Welcome, %s!",
+            app->params->uname);
+    notification("Registration", dc_normal, query);
+    print_bars(app);
+    return 0;
+  } else {
+    q_args->state = S_ASK_REGISTER;
+    alert(line);
+    return 1;
+  }
+}
+
+int server_print_message_cb(app_t *app, char *line, int l_len) {
+  int64_t new_capacity;
+  query_args_t *q_args = app->query_args;
+  if (q_args->server_message.text == NULL) {
+    q_args->server_message.text = malloc(l_len + 1);
+    memcpy(q_args->server_message.text, line, l_len);
+    q_args->server_message.text[l_len] = 0;
+    q_args->server_message.size = l_len + 1;
+    q_args->server_message.capacity = l_len + 1;
+  } else {
+    if (q_args->server_message.size + l_len > q_args->server_message.capacity) {
+      new_capacity = (l_len + q_args->server_message.capacity) * 2;
+      q_args->server_message.text =
+          realloc(q_args->server_message.text, new_capacity);
+      q_args->server_message.capacity = new_capacity;
+    }
+    q_args->server_message.size += l_len;
+    strncat(q_args->server_message.text, line, l_len);
+    q_args->server_message.text[q_args->server_message.size + l_len] = 0;
+  }
+  char *end_char = strchr(q_args->server_message.text,
+                          '\04'); /* Search for the end of the notification */
+  if (end_char != NULL) {
+    query_return_to_buf(app->query_args->buf, &(app->query_args->buf_used),
+                        end_char + 1);
+    *end_char = '\0'; /* Don't show this symbol */
+    notification("Server response", dc_normal, q_args->server_message.text);
+    /* q_args->state = S_PRINT_SERVER_MESSAGE; */
+    q_args->server_message.size = 0;
+    q_args->server_message.text[0] = 0;
+    q_args->state = S_NEXT_ACTION;
+    return 1;
+  }
+  return 0;
+}
+
 /* TODO: refactor to process callbacks */
 int process_server_command(char *line, int l_len, app_t *app) {
   int ws_pos = l_len;
   query_args_t *q_args = app->query_args;
   params_t *params = q_args->params;
-  uint64_t new_capacity;
   ui_file_list_t *fui = (ui_file_list_t *)q_args->file_list_ui;
-  char query[INBUFSIZE];
 
   char *cptr = strchr(line, ' ');
   if (cptr != NULL && cptr > line)
@@ -55,9 +105,6 @@ int process_server_command(char *line, int l_len, app_t *app) {
 
   /* LOGIN */
   if (!strncmp(line, "login>", ws_pos)) {
-    if (q_args->server_message.size > 0) {
-      PRINT_SRV_MESSAGE(q_args, l_len, line);
-    }
     if (q_args->params->uname != NULL && q_args->params->pass != NULL) {
       server_send_string(q_args, q_args->params->uname);
     } else {
@@ -86,27 +133,6 @@ int process_server_command(char *line, int l_len, app_t *app) {
     return 0;
   }
 
-  /* REGISTER CONFIRMATION */
-  if (q_args->state == S_WAIT_REGISTER_CONFIRMATION) {
-    if (!strcmp(line, "ok\n")) {
-      server_send_string(q_args, "file list %u %u\n", fui->max_lines,
-                         fui->current_page);
-      q_args->state = S_FILE_LIST;
-      destroy_dialogue(&(app->modal), app);
-      sprintf(query,
-              "You've been successfully registered.\n"
-              "Welcome, %s!",
-              app->params->uname);
-      notification("Registration", dc_normal, query);
-      print_bars(app);
-      return 0;
-    } else {
-      q_args->state = S_ASK_REGISTER;
-      alert(line);
-      return 1;
-    }
-  }
-
   /* DOWNLOAD START */
   if (!strncmp(line, "download_start", ws_pos)) {
     return 0;
@@ -120,36 +146,7 @@ int process_server_command(char *line, int l_len, app_t *app) {
     return 0;
   }
 
-  /* TODO: Hide details */
-  if (q_args->state == S_WAIT_SERVER) {
-    if (q_args->server_message.text == NULL) {
-      q_args->server_message.text = malloc(l_len + 1);
-      memcpy(q_args->server_message.text, line, l_len);
-      q_args->server_message.text[l_len] = 0;
-      q_args->server_message.size = l_len + 1;
-      q_args->server_message.capacity = l_len + 1;
-    } else {
-      if (q_args->server_message.size + l_len >
-          q_args->server_message.capacity) {
-        new_capacity = (l_len + q_args->server_message.capacity) * 2;
-        q_args->server_message.text =
-            realloc(q_args->server_message.text, new_capacity);
-        q_args->server_message.capacity = new_capacity;
-      }
-      q_args->server_message.size += l_len;
-      strncat(q_args->server_message.text, line, l_len);
-      q_args->server_message.text[q_args->server_message.size + l_len] = 0;
-    }
-    char *end_char = strchr(q_args->server_message.text,
-                            '\04'); /* Search for the end of the notification */
-    if (end_char != NULL) {
-      *end_char = '\0'; /* Don't show this symbol */
-      q_args->state = S_PRINT_SERVER_MESSAGE;
-      q_args->server_message.size--;
-    }
-  }
-
-  return 0;
+  return -1;
 }
 
 void ask_uname_and_password(params_t *params) {
diff --git a/src/server.h b/src/server.h
index 6e3cdf8..139ef92 100644
--- a/src/server.h
+++ b/src/server.h
@@ -6,7 +6,9 @@
 
 #include "app.h"
 #include "main.h"
-int process_server_command(char *line, int l_len, app_t *app);
+int  process_server_command(char *line, int l_len, app_t *app);
 void ask_register(params_t *params, char *email);
 void server_send_string(query_args_t *q, const char *fmt, ...);
+int  server_print_message_cb(app_t *app, char *line, int l_len);
+int  server_wait_reg_confirm_cb(app_t *app, char *line, int l_len);
 #endif
\ No newline at end of file