diff --git a/src/modbus.c b/src/modbus.c
index 8d0aafa..0e64c01 100644
--- a/src/modbus.c
+++ b/src/modbus.c
@@ -1294,6 +1294,11 @@ void modbus_set_timeout_end(modbus_t *ctx, const struct timeval *timeout)
     ctx->timeout_end = *timeout;
 }
 
+int modbus_get_header_length(modbus_t *ctx)
+{
+    return ctx->backend->header_length;
+}
+
 int modbus_connect(modbus_t *ctx)
 {
     return ctx->backend->connect(ctx);
diff --git a/src/modbus.h b/src/modbus.h
index d9a2674..e2b70d5 100644
--- a/src/modbus.h
+++ b/src/modbus.h
@@ -139,6 +139,8 @@ void modbus_set_timeout_begin(modbus_t *ctx, const struct timeval *timeout);
 void modbus_get_timeout_end(modbus_t *ctx, struct timeval *timeout);
 void modbus_set_timeout_end(modbus_t *ctx, const struct timeval *timeout);
 
+int modbus_get_header_length(modbus_t *ctx);
+
 int modbus_connect(modbus_t *ctx);
 void modbus_close(modbus_t *ctx);
 
diff --git a/tests/unit-test-server.c b/tests/unit-test-server.c
index cf9693b..3b3d6f9 100644
--- a/tests/unit-test-server.c
+++ b/tests/unit-test-server.c
@@ -40,6 +40,8 @@ int main(int argc, char*argv[])
     int rc;
     int i;
     int use_backend;
+    uint8_t *query;
+    int header_length;
 
     if (argc > 1) {
         if (strcmp(argv[1], "tcp") == 0) {
@@ -57,10 +59,13 @@ int main(int argc, char*argv[])
 
     if (use_backend == TCP) {
         ctx = modbus_new_tcp("127.0.0.1", 1502);
+        query = malloc(MODBUS_TCP_MAX_ADU_LENGTH);
     } else {
         ctx = modbus_new_rtu("/dev/ttyUSB0", 115200, 'N', 8, 1);
         modbus_set_slave(ctx, SERVER_ID);
+        query = malloc(MODBUS_RTU_MAX_ADU_LENGTH);
     }
+    header_length = modbus_get_header_length(ctx);
 
     modbus_set_debug(ctx, TRUE);
     modbus_set_error_recovery(ctx, TRUE);
@@ -104,17 +109,15 @@ int main(int argc, char*argv[])
     }
 
     for (;;) {
-        uint8_t query[MODBUS_TCP_MAX_ADU_LENGTH];
-
         rc = modbus_receive(ctx, -1, query);
         if (rc > 0) {
-            if (((query[HEADER_LENGTH_TCP + 3] << 8) +
-                 query[HEADER_LENGTH_TCP + 4])
+            if (((query[header_length + 3] << 8) +
+                 query[header_length + 4])
                 == UT_REGISTERS_NB_POINTS_SPECIAL) {
                 /* Change the number of values (offset
                    TCP = 6) */
-                query[HEADER_LENGTH_TCP + 3] = 0;
-                query[HEADER_LENGTH_TCP + 4] = UT_REGISTERS_NB_POINTS;
+                query[header_length + 3] = 0;
+                query[header_length + 4] = UT_REGISTERS_NB_POINTS;
             }
 
             rc = modbus_reply(ctx, query, rc, mb_mapping);
@@ -133,6 +136,7 @@ int main(int argc, char*argv[])
         close(socket);
     }
     modbus_mapping_free(mb_mapping);
+    free(query);
     modbus_free(ctx);
 
     return 0;