| 1 | // SPDX-License-Identifier: GPL-3.0-or-later |
| 2 | |
| 3 | #include "websocket-internal.h" |
| 4 | |
| 5 | /** |
| 6 | * @brief Compresses a message into the client's c_payload buffer |
| 7 | * |
| 8 | * This function compresses the given data using the client's deflate stream |
| 9 | * and stores the result in the client's c_payload buffer. |
| 10 | * |
| 11 | * @param wsc WebSocket client |
| 12 | * @param data Data to compress |
| 13 | * @param length Length of the data |
| 14 | * @return true if compression was successful, false otherwise |
| 15 | */ |
| 16 | bool websocket_client_compress_message(WS_CLIENT *wsc, const char *data, size_t length) { |
| 17 | if (!wsc || !data || !length || !wsc->compression.enabled || !wsc->compression.deflate_stream) |
| 18 | return false; |
| 19 | |
| 20 | if (length < WS_COMPRESS_MIN_SIZE) |
| 21 | return false; // Too small to benefit from compression |
| 22 | |
| 23 | // Clear and prepare the compression buffer |
| 24 | wsb_reset(&wsc->c_payload); |
| 25 | |
| 26 | z_stream *zstrm = wsc->compression.deflate_stream; |
| 27 | |
| 28 | // Calculate maximum possible compressed size |
| 29 | uLong max_compressed_size = deflateBound(zstrm, length) + 4; // +4 for Z_SYNC_FLUSH trailer |
| 30 | |
| 31 | // Ensure the buffer has enough capacity |
| 32 | wsb_need_bytes(&wsc->c_payload, max_compressed_size); |
| 33 | |
| 34 | // Set up the deflate stream |
| 35 | zstrm->next_in = (Bytef *)data; |
| 36 | zstrm->avail_in = length; |
| 37 | zstrm->next_out = (Bytef *)wsb_data(&wsc->c_payload); |
| 38 | zstrm->avail_out = wsb_size(&wsc->c_payload); |
| 39 | zstrm->total_in = 0; |
| 40 | zstrm->total_out = 0; |
| 41 | |
| 42 | // Compress with sync flush |
| 43 | int ret = deflate(zstrm, Z_SYNC_FLUSH); |
| 44 | |
| 45 | bool success = false; |
| 46 | if (ret == Z_STREAM_END || (ret == Z_OK && zstrm->avail_in == 0 && zstrm->avail_out > 0)) |
| 47 | success = true; |
| 48 | else if (ret == Z_OK && zstrm->avail_in == 0 && zstrm->avail_out == 0) { |
| 49 | unsigned pending = 0; |
| 50 | int bits = 0; |
| 51 | if(deflatePending(zstrm, &pending, &bits) == Z_OK && |
| 52 | (pending == 0 && bits == 0)) |
| 53 | success = true; |
| 54 | } |
| 55 | |
| 56 | uLong total_out = zstrm->total_out; |
| 57 | |
| 58 | // Reset the stream for future use |
| 59 | if (deflateReset(zstrm) != Z_OK) { |
| 60 | websocket_error(wsc, "Deflate reset failed"); |
| 61 | |
| 62 | // Clear pointers for safety |
| 63 | zstrm->next_in = NULL; |
| 64 | zstrm->avail_in = 0; |
| 65 | zstrm->next_out = NULL; |
| 66 | zstrm->avail_out = 0; |
| 67 | zstrm->total_in = 0; |
| 68 | zstrm->total_out = 0; |
| 69 | |
| 70 | return false; |
| 71 | } |
| 72 | |
| 73 | // Clear all pointers for safety |
| 74 | zstrm->next_in = NULL; |
| 75 | zstrm->avail_in = 0; |
| 76 | zstrm->next_out = NULL; |
| 77 | zstrm->avail_out = 0; |
| 78 | zstrm->total_in = 0; |
| 79 | zstrm->total_out = 0; |
| 80 | |
| 81 | if (!success || total_out <= 4) { |
| 82 | // Compression failed or didn't save space |
| 83 | websocket_debug(wsc, "Compression not beneficial (in=%zu, out=%lu) - not using compression", |
| 84 | length, total_out); |
| 85 | return false; |
| 86 | } |
| 87 | |
| 88 | // As per RFC 7692, remove trailing 4 bytes (00 00 FF FF) from Z_SYNC_FLUSH |
| 89 | size_t compressed_size = total_out - 4; |
| 90 | |
| 91 | // Update the buffer's length |
| 92 | wsb_set_length(&wsc->c_payload, compressed_size); |
| 93 | |
| 94 | websocket_debug(wsc, "Compressed message from %zu to %zu bytes (%.1f%%)", |
| 95 | length, compressed_size, (double)compressed_size * 100.0 / (double)length); |
| 96 | |
| 97 | return true; |
| 98 | } |
| 99 | |
| 100 | // Initialize compression resources using the parsed options |
| 101 | bool websocket_compression_init(WS_CLIENT *wsc) { |
| 102 | internal_fatal(wsc->wth->tid != gettid_cached(), "Function %s() should only be used by the websocket thread", __FUNCTION__ ); |
| 103 | |
| 104 | if (!wsc->compression.enabled) { |
| 105 | websocket_debug(wsc, "Compression is disabled"); |
| 106 | return false; |
| 107 | } |
| 108 | |
| 109 | // Initialize deflate (compression) context for server-to-client messages |
| 110 | wsc->compression.deflate_stream = mallocz(sizeof(z_stream)); |
| 111 | wsc->compression.deflate_stream->zalloc = Z_NULL; |
| 112 | wsc->compression.deflate_stream->zfree = Z_NULL; |
| 113 | wsc->compression.deflate_stream->opaque = Z_NULL; |
| 114 | |
| 115 | // Initialize with negative window bits for raw deflate (no zlib/gzip header) |
| 116 | // Use server_max_window_bits for outgoing (server-to-client) messages |
| 117 | int ret = deflateInit2( |
| 118 | wsc->compression.deflate_stream, |
| 119 | wsc->compression.compression_level, |
| 120 | Z_DEFLATED, |
| 121 | -wsc->compression.server_max_window_bits, |
| 122 | WS_COMPRESS_MEMLEVEL, |
| 123 | Z_DEFAULT_STRATEGY |
| 124 | ); |
| 125 | |
| 126 | if (ret != Z_OK) { |
| 127 | websocket_error(wsc, "Failed to initialize deflate context: %s (%d)", |
| 128 | zError(ret), ret); |
| 129 | freez(wsc->compression.deflate_stream); |
| 130 | wsc->compression.deflate_stream = NULL; |
| 131 | return false; |
| 132 | } |
| 133 | |
| 134 | websocket_debug(wsc, "Compression initialized (server window bits: %d)", |
| 135 | wsc->compression.server_max_window_bits); |
| 136 | |
| 137 | return true; |
| 138 | } |
| 139 | |
| 140 | // Initialize decompression resources for a client |
| 141 | bool websocket_decompression_init(WS_CLIENT *wsc) { |
| 142 | internal_fatal(wsc->wth->tid != gettid_cached(), "Function %s() should only be used by the websocket thread", __FUNCTION__ ); |
| 143 | |
| 144 | if (!wsc->compression.enabled) { |
| 145 | websocket_debug(wsc, "Decompression is disabled"); |
| 146 | return false; |
| 147 | } |
| 148 | |
| 149 | // Create a new inflate stream |
| 150 | wsc->compression.inflate_stream = mallocz(sizeof(z_stream)); |
| 151 | wsc->compression.inflate_stream->zalloc = Z_NULL; |
| 152 | wsc->compression.inflate_stream->zfree = Z_NULL; |
| 153 | wsc->compression.inflate_stream->opaque = Z_NULL; |
| 154 | |
| 155 | // Initialize with negative window bits for raw deflate (no zlib/gzip header) |
| 156 | // Use client_max_window_bits for incoming (client-to-server) messages |
| 157 | int init_ret = inflateInit2(wsc->compression.inflate_stream, -wsc->compression.client_max_window_bits); |
| 158 | |
| 159 | if (init_ret != Z_OK) { |
| 160 | websocket_error(wsc, "Failed to initialize inflate stream: %s (%d)", |
| 161 | zError(init_ret), init_ret); |
| 162 | freez(wsc->compression.inflate_stream); |
| 163 | wsc->compression.inflate_stream = NULL; |
| 164 | return false; |
| 165 | } |
| 166 | |
| 167 | websocket_debug(wsc, "Decompression initialized (client window bits: %d)", |
| 168 | wsc->compression.client_max_window_bits); |
| 169 | |
| 170 | return true; |
| 171 | } |
| 172 | |
| 173 | // Clean up compression resources for a WebSocket client |
| 174 | void websocket_compression_cleanup(WS_CLIENT *wsc) { |
| 175 | // Clean up deflate context |
| 176 | if (!wsc->compression.deflate_stream) |
| 177 | return; |
| 178 | |
| 179 | internal_fatal(wsc->wth->tid != gettid_cached(), "Function %s() should only be used by the websocket thread", __FUNCTION__ ); |
| 180 | |
| 181 | // Set up dummy I/O pointers to ensure clean state |
| 182 | unsigned char dummy_buffer[16] = {0}; |
| 183 | wsc->compression.deflate_stream->next_in = dummy_buffer; |
| 184 | wsc->compression.deflate_stream->avail_in = 0; |
| 185 | wsc->compression.deflate_stream->next_out = dummy_buffer; |
| 186 | wsc->compression.deflate_stream->avail_out = sizeof(dummy_buffer); |
| 187 | |
| 188 | // Always call deflateEnd to release internal zlib resources |
| 189 | // Don't bother with deflateReset as deflateEnd will clean up properly |
| 190 | int ret = deflateEnd(wsc->compression.deflate_stream); |
| 191 | |
| 192 | if (ret != Z_OK && ret != Z_DATA_ERROR) { |
| 193 | // Z_DATA_ERROR can happen in some edge cases, it's not critical here |
| 194 | // as we're cleaning up anyway |
| 195 | websocket_debug(wsc, "deflateEnd returned %d: %s", ret, zError(ret)); |
| 196 | } |
| 197 | |
| 198 | // Free the stream structure |
| 199 | freez(wsc->compression.deflate_stream); |
| 200 | wsc->compression.deflate_stream = NULL; |
| 201 | |
| 202 | websocket_debug(wsc, "Compression resources cleaned up"); |
| 203 | } |
| 204 | |
| 205 | // Clean up decompression resources for a client's inflate stream |
| 206 | void websocket_decompression_cleanup(WS_CLIENT *wsc) { |
| 207 | if (!wsc->compression.inflate_stream) |
| 208 | return; |
| 209 | |
| 210 | internal_fatal(wsc->wth->tid != gettid_cached(), "Function %s() should only be used by the websocket thread", __FUNCTION__ ); |
| 211 | |
| 212 | // End the current inflate stream and free its resources |
| 213 | inflateEnd(wsc->compression.inflate_stream); |
| 214 | freez(wsc->compression.inflate_stream); |
| 215 | wsc->compression.inflate_stream = NULL; |
| 216 | |
| 217 | websocket_debug(wsc, "Decompression resources cleaned up"); |
| 218 | } |
| 219 | |
| 220 | // Reset compression resources for a client - calls cleanup and init |
| 221 | ALWAYS_INLINE |
| 222 | bool websocket_compression_reset(WS_CLIENT *wsc) { |
| 223 | websocket_compression_cleanup(wsc); |
| 224 | return websocket_compression_init(wsc); |
| 225 | } |
| 226 | |
| 227 | // Reset decompression resources for a client - calls cleanup and init |
| 228 | ALWAYS_INLINE |
| 229 | bool websocket_decompression_reset(WS_CLIENT *wsc) { |
| 230 | websocket_decompression_cleanup(wsc); |
| 231 | return websocket_decompression_init(wsc); |
| 232 | } |
| 233 | |
| 234 | // Decompress a client's message from payload to u_payload |
| 235 | bool websocket_client_decompress_message(WS_CLIENT *wsc) { |
| 236 | internal_fatal(wsc->wth->tid != gettid_cached(), "Function %s() should only be used by the websocket thread", __FUNCTION__ ); |
| 237 | |
| 238 | if (!wsc->is_compressed || !wsc->compression.enabled || !wsc->compression.inflate_stream) |
| 239 | return false; |
| 240 | |
| 241 | if (wsb_is_empty(&wsc->payload)) { |
| 242 | websocket_debug(wsc, "Empty compressed message"); |
| 243 | wsb_reset(&wsc->u_payload); |
| 244 | wsb_null_terminate(&wsc->u_payload); |
| 245 | return true; |
| 246 | } |
| 247 | |
| 248 | websocket_debug(wsc, "Decompressing message (%zu bytes)", wsb_length(&wsc->payload)); |
| 249 | |
| 250 | z_stream *zstrm = wsc->compression.inflate_stream; |
| 251 | wsb_reset(&wsc->u_payload); |
| 252 | |
| 253 | // Per RFC 7692, we need to append 4 bytes (00 00 FF FF) to the compressed data |
| 254 | // to ensure the inflate operation completes |
| 255 | static const unsigned char trailer[4] = {0x00, 0x00, 0xFF, 0xFF}; |
| 256 | wsb_append_padding(&wsc->payload, trailer, 4); |
| 257 | |
| 258 | zstrm->next_in = (Bytef *)wsb_data(&wsc->payload); |
| 259 | zstrm->avail_in = wsb_length(&wsc->payload) + 4; |
| 260 | zstrm->next_out = (Bytef *)wsb_data(&wsc->u_payload); |
| 261 | zstrm->avail_out = wsb_size(&wsc->u_payload); |
| 262 | zstrm->total_in = 0; |
| 263 | zstrm->total_out = 0; |
| 264 | |
| 265 | // Decompress with loop for multiple buffer expansions if needed |
| 266 | int ret = Z_MEM_ERROR; |
| 267 | bool success = false; |
| 268 | int retries = 24; |
| 269 | size_t wanted_size = MAX(wsb_size(&wsc->u_payload), wsb_length(&wsc->payload) * 2); |
| 270 | do { |
| 271 | wsb_resize(&wsc->u_payload, wanted_size); |
| 272 | |
| 273 | // Position next_out to point to the end of the currently decompressed data |
| 274 | zstrm->next_out = (Bytef *)wsb_data(&wsc->u_payload) + wsb_length(&wsc->u_payload); |
| 275 | |
| 276 | // Only make the newly available space available to zlib |
| 277 | zstrm->avail_out = wsb_size(&wsc->u_payload) - wsb_length(&wsc->u_payload); |
| 278 | |
| 279 | // Try to decompress |
| 280 | ret = inflate(zstrm, Z_SYNC_FLUSH); |
| 281 | |
| 282 | websocket_debug(wsc, "inflate() returned %d (%s), " |
| 283 | "avail_in=%u, avail_out=%u, total_in=%lu, total_out=%lu", |
| 284 | ret, zError(ret), |
| 285 | zstrm->avail_in, zstrm->avail_out, zstrm->total_in, zstrm->total_out); |
| 286 | |
| 287 | // Handle different return codes from inflate() |
| 288 | // Z_STREAM_END - Complete decompression success |
| 289 | // Z_OK - Partial success, all input processed or output buffer full |
| 290 | // Z_BUF_ERROR - Need more output space |
| 291 | |
| 292 | success = ret == Z_STREAM_END || |
| 293 | (zstrm->avail_in == 0 && zstrm->avail_out > 0 && (ret == Z_OK || ret == Z_BUF_ERROR)); |
| 294 | |
| 295 | // Update the buffer's length to include the newly written data |
| 296 | wsb_set_length(&wsc->u_payload, wsb_size(&wsc->u_payload) - zstrm->avail_out); |
| 297 | |
| 298 | // Check if we need more output space |
| 299 | if (!success && (ret == Z_BUF_ERROR || ret == Z_OK)) { |
| 300 | wanted_size = MIN(wanted_size * 2, WS_MAX_DECOMPRESSED_SIZE); |
| 301 | if (wanted_size == WS_MAX_DECOMPRESSED_SIZE && wanted_size == wsb_size(&wsc->u_payload)) |
| 302 | break; // we cannot resize more |
| 303 | } |
| 304 | } while (!success && retries-- > 0); |
| 305 | |
| 306 | if(!success) { |
| 307 | // Decompression failed |
| 308 | websocket_error(wsc, "Decompression failed: %s (ret = %d, avail_in = %u)", zError(ret), ret, zstrm->avail_in); |
| 309 | wsb_reset(&wsc->u_payload); |
| 310 | websocket_decompression_reset(wsc); |
| 311 | return false; |
| 312 | } |
| 313 | |
| 314 | // Log successful decompression with detailed information |
| 315 | websocket_debug(wsc, "Successfully decompressed %zu bytes to %zu bytes (ratio: %.2fx)", |
| 316 | wsb_length(&wsc->payload), wsb_length(&wsc->u_payload), |
| 317 | (double)wsb_length(&wsc->u_payload) / (double)wsb_length(&wsc->payload)); |
| 318 | |
| 319 | // Show a preview of the decompressed data |
| 320 | websocket_dump_debug(wsc, wsb_data(&wsc->u_payload), wsb_length(&wsc->u_payload), "RX UNCOMPRESSED PAYLOAD"); |
| 321 | |
| 322 | // when client context takeover is disabled, reset the decompressor |
| 323 | if (!wsc->compression.client_context_takeover) { |
| 324 | websocket_debug(wsc, "resetting compression"); |
| 325 | if(inflateReset2(zstrm, -wsc->compression.client_max_window_bits) != Z_OK) { |
| 326 | websocket_debug(wsc, "reset failed, re-initializing compression"); |
| 327 | if (!websocket_decompression_reset(wsc)) { |
| 328 | websocket_debug(wsc, "re-initializing failed, reporting failure"); |
| 329 | return false; |
| 330 | } |
| 331 | zstrm = wsc->compression.inflate_stream; |
| 332 | } |
| 333 | } |
| 334 | |
| 335 | zstrm->next_in = NULL; |
| 336 | zstrm->next_out = NULL; |
| 337 | zstrm->avail_in = 0; |
| 338 | zstrm->avail_out = 0; |
| 339 | zstrm->total_in = 0; |
| 340 | zstrm->total_out = 0; |
| 341 | |
| 342 | return true; |
| 343 | } |