diff --git a/library/ssl_cli.c b/library/ssl_cli.c
index be80de7..0c587e9 100644
--- a/library/ssl_cli.c
+++ b/library/ssl_cli.c
@@ -1762,12 +1762,7 @@
 
 #if defined(MBEDTLS_ZLIB_SUPPORT)
     /* See comments in ssl_write_client_hello() */
-#if defined(MBEDTLS_SSL_PROTO_DTLS)
-    if( ssl->conf->transport == MBEDTLS_SSL_TRANSPORT_DATAGRAM )
-        accept_comp = 0;
-    else
-#endif
-        accept_comp = 1;
+    accept_comp = MBEDTLS_SSL_TRANSPORT_IS_TLS( ssl->conf->transport );
 
     if( comp != MBEDTLS_SSL_COMPRESS_NULL &&
         ( comp != MBEDTLS_SSL_COMPRESS_DEFLATE || accept_comp == 0 ) )