token.c

   1 /*
   2  * Routines used by the file-transfer code.
   3  *
   4  * Copyright (C) 1996 Andrew Tridgell
   5  * Copyright (C) 1996 Paul Mackerras
   6  * Copyright (C) 2003-2022 Wayne Davison
   7  *
   8  * This program is free software; you can redistribute it and/or modify
   9  * it under the terms of the GNU General Public License as published by
  10  * the Free Software Foundation; either version 3 of the License, or
  11  * (at your option) any later version.
  12  *
  13  * This program is distributed in the hope that it will be useful,
  14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  16  * GNU General Public License for more details.
  17  *
  18  * You should have received a copy of the GNU General Public License along
  19  * with this program; if not, visit the http://fsf.org website.
  20  */
  21
  22 #include "rsync.h"
  23 #include "itypes.h"
  24 #include <zlib.h>
  25 #ifdef SUPPORT_ZSTD
  26 #include <zstd.h>
  27 #endif
  28 #ifdef SUPPORT_LZ4
  29 #include <lz4.h>
  30 #endif
  31
  32 extern int do_compression;
  33 extern int protocol_version;
  34 extern int module_id;
  35 extern int do_compression_level;
  36 extern char *skip_compress;
  37
  38 #ifndef Z_INSERT_ONLY
  39 #define Z_INSERT_ONLY Z_SYNC_FLUSH
  40 #endif
  41
  42 static int skip_compression_level; /* The least possible compressing for handling skip-compress files. */
  43 static int per_file_default_level; /* The default level that each new file gets prior to checking its suffix. */
  44
  45 struct suffix_tree {
  46         struct suffix_tree *sibling;
  47         struct suffix_tree *child;
  48         char letter, word_end;
  49 };
  50
  51 static char *match_list;
  52 static struct suffix_tree *suftree;
  53
  54 void init_compression_level(void)
  55 {
  56         int min_level, max_level, def_level, off_level;
  57
  58         switch (do_compression) {
  59         case CPRES_NONE:
  60                 return;
  61         case CPRES_ZLIB:
  62         case CPRES_ZLIBX:
  63                 min_level = 1;
  64                 max_level = Z_BEST_COMPRESSION;
  65                 def_level = 6; /* Z_DEFAULT_COMPRESSION is -1, so set it to the real default */
  66                 off_level = skip_compression_level = Z_NO_COMPRESSION;
  67                 if (do_compression_level == Z_DEFAULT_COMPRESSION)
  68                         do_compression_level = def_level;
  69                 break;
  70 #ifdef SUPPORT_ZSTD
  71         case CPRES_ZSTD:
  72                 min_level = skip_compression_level = ZSTD_minCLevel();
  73                 max_level = ZSTD_maxCLevel();
  74                 def_level = ZSTD_CLEVEL_DEFAULT;
  75                 off_level = CLVL_NOT_SPECIFIED;
  76                 if (do_compression_level == 0)
  77                         do_compression_level = def_level;
  78                 break;
  79 #endif
  80 #ifdef SUPPORT_LZ4
  81         case CPRES_LZ4:
  82                 min_level = skip_compression_level = 0;
  83                 max_level = 0;
  84                 def_level = 0;
  85                 off_level = CLVL_NOT_SPECIFIED;
  86                 break;
  87 #endif
  88         default: /* paranoia to prevent missing case values */
  89                 NOISY_DEATH("Unknown do_compression value");
  90         }
  91
  92         if (do_compression_level == CLVL_NOT_SPECIFIED)
  93                 do_compression_level = def_level;
  94         else if (do_compression_level == off_level) {
  95                 do_compression = CPRES_NONE;
  96                 return;
  97         }
  98
  99         /* We don't bother with any errors or warnings -- just make sure that the values are valid. */
 100         if (do_compression_level < min_level)
 101                 do_compression_level = min_level;
 102         else if (do_compression_level > max_level)
 103                 do_compression_level = max_level;
 104 }
 105
 106 static void add_suffix(struct suffix_tree **prior, char ltr, const char *str)
 107 {
 108         struct suffix_tree *node, *newnode;
 109
 110         if (ltr == '[') {
 111                 const char *after = strchr(str, ']');
 112                 /* Treat "[foo" and "[]" as having a literal '['. */
 113                 if (after && after++ != str+1) {
 114                         while ((ltr = *str++) != ']')
 115                                 add_suffix(prior, ltr, after);
 116                         return;
 117                 }
 118         }
 119
 120         for (node = *prior; node; prior = &node->sibling, node = node->sibling) {
 121                 if (node->letter == ltr) {
 122                         if (*str)
 123                                 add_suffix(&node->child, *str, str+1);
 124                         else
 125                                 node->word_end = 1;
 126                         return;
 127                 }
 128                 if (node->letter > ltr)
 129                         break;
 130         }
 131         newnode = new(struct suffix_tree);
 132         newnode->sibling = node;
 133         newnode->child = NULL;
 134         newnode->letter = ltr;
 135         *prior = newnode;
 136         if (*str) {
 137                 add_suffix(&newnode->child, *str, str+1);
 138                 newnode->word_end = 0;
 139         } else
 140                 newnode->word_end = 1;
 141 }
 142
 143 static void add_nocompress_suffixes(const char *str)
 144 {
 145         char *buf, *t;
 146         const char *f = str;
 147
 148         buf = new_array(char, strlen(f) + 1);
 149
 150         while (*f) {
 151                 if (*f == '/') {
 152                         f++;
 153                         continue;
 154                 }
 155
 156                 t = buf;
 157                 do {
 158                         if (isUpper(f))
 159                                 *t++ = toLower(f);
 160                         else
 161                                 *t++ = *f;
 162                 } while (*++f != '/' && *f);
 163                 *t++ = '\0';
 164
 165                 add_suffix(&suftree, *buf, buf+1);
 166         }
 167
 168         free(buf);
 169 }
 170
 171 static void init_set_compression(void)
 172 {
 173         const char *f;
 174         char *t, *start;
 175
 176         if (skip_compress)
 177                 add_nocompress_suffixes(skip_compress);
 178
 179         /* A non-daemon transfer skips the default suffix list if the
 180          * user specified --skip-compress. */
 181         if (skip_compress && module_id < 0)
 182                 f = "";
 183         else
 184                 f = lp_dont_compress(module_id);
 185
 186         match_list = t = new_array(char, strlen(f) + 2);
 187
 188         per_file_default_level = do_compression_level;
 189
 190         while (*f) {
 191                 if (*f == ' ') {
 192                         f++;
 193                         continue;
 194                 }
 195
 196                 start = t;
 197                 do {
 198                         if (isUpper(f))
 199                                 *t++ = toLower(f);
 200                         else
 201                                 *t++ = *f;
 202                 } while (*++f != ' ' && *f);
 203                 *t++ = '\0';
 204
 205                 if (t - start == 1+1 && *start == '*') {
 206                         /* Optimize a match-string of "*". */
 207                         *match_list = '\0';
 208                         suftree = NULL;
 209                         per_file_default_level = skip_compression_level;
 210                         break;
 211                 }
 212
 213                 /* Move *.foo items into the stuffix tree. */
 214                 if (*start == '*' && start[1] == '.' && start[2]
 215                  && !strpbrk(start+2, ".?*")) {
 216                         add_suffix(&suftree, start[2], start+3);
 217                         t = start;
 218                 }
 219         }
 220         *t++ = '\0';
 221 }
 222
 223 /* determine the compression level based on a wildcard filename list */
 224 void set_compression(const char *fname)
 225 {
 226 #if 0 /* No compression algorithms currently allow mid-stream changing of the level. */
 227         const struct suffix_tree *node;
 228         const char *s;
 229         char ltr;
 230 #endif
 231
 232         if (!do_compression)
 233                 return;
 234
 235         if (!match_list)
 236                 init_set_compression();
 237
 238 #if 0
 239         compression_level = per_file_default_level;
 240
 241         if (!*match_list && !suftree)
 242                 return;
 243
 244         if ((s = strrchr(fname, '/')) != NULL)
 245                 fname = s + 1;
 246
 247         for (s = match_list; *s; s += strlen(s) + 1) {
 248                 if (iwildmatch(s, fname)) {
 249                         compression_level = skip_compression_level;
 250                         return;
 251                 }
 252         }
 253
 254         if (!(node = suftree) || !(s = strrchr(fname, '.'))
 255          || s == fname || !(ltr = *++s))
 256                 return;
 257
 258         while (1) {
 259                 if (isUpper(&ltr))
 260                         ltr = toLower(&ltr);
 261                 while (node->letter != ltr) {
 262                         if (node->letter > ltr)
 263                                 return;
 264                         if (!(node = node->sibling))
 265                                 return;
 266                 }
 267                 if ((ltr = *++s) == '\0') {
 268                         if (node->word_end)
 269                                 compression_level = skip_compression_level;
 270                         return;
 271                 }
 272                 if (!(node = node->child))
 273                         return;
 274         }
 275 #else
 276         (void)fname;
 277 #endif
 278 }
 279
 280 /* non-compressing recv token */
 281 static int32 simple_recv_token(int f, char **data)
 282 {
 283         static int32 residue;
 284         static char *buf;
 285         int32 n;
 286
 287         if (!buf)
 288                 buf = new_array(char, CHUNK_SIZE);
 289
 290         if (residue == 0) {
 291                 int32 i = read_int(f);
 292                 if (i <= 0)
 293                         return i;
 294                 residue = i;
 295         }
 296
 297         *data = buf;
 298         n = MIN(CHUNK_SIZE,residue);
 299         residue -= n;
 300         read_buf(f,buf,n);
 301         return n;
 302 }
 303
 304 /* non-compressing send token */
 305 static void simple_send_token(int f, int32 token, struct map_struct *buf, OFF_T offset, int32 n)
 306 {
 307         if (n > 0) {
 308                 int32 len = 0;
 309                 while (len < n) {
 310                         int32 n1 = MIN(CHUNK_SIZE, n-len);
 311                         write_int(f, n1);
 312                         write_buf(f, map_ptr(buf, offset+len, n1), n1);
 313                         len += n1;
 314                 }
 315         }
 316         /* a -2 token means to send data only and no token */
 317         if (token != -2)
 318                 write_int(f, -(token+1));
 319 }
 320
 321 /* Flag bytes in compressed stream are encoded as follows: */
 322 #define END_FLAG        0       /* that's all folks */
 323 #define TOKEN_LONG      0x20    /* followed by 32-bit token number */
 324 #define TOKENRUN_LONG   0x21    /* ditto with 16-bit run count */
 325 #define DEFLATED_DATA   0x40    /* + 6-bit high len, then low len byte */
 326 #define TOKEN_REL       0x80    /* + 6-bit relative token number */
 327 #define TOKENRUN_REL    0xc0    /* ditto with 16-bit run count */
 328
 329 #define MAX_DATA_COUNT  16383   /* fit 14 bit count into 2 bytes with flags */
 330
 331 /* zlib.h says that if we want to be able to compress something in a single
 332  * call, avail_out must be at least 0.1% larger than avail_in plus 12 bytes.
 333  * We'll add in 0.1%+16, just to be safe (and we'll avoid floating point,
 334  * to ensure that this is a compile-time value). */
 335 #define AVAIL_OUT_SIZE(avail_in_size) ((avail_in_size)*1001/1000+16)
 336
 337 /* For coding runs of tokens */
 338 static int32 last_token = -1;
 339 static int32 run_start;
 340 static int32 last_run_end;
 341
 342 /* Deflation state */
 343 static z_stream tx_strm;
 344
 345 /* Output buffer */
 346 static char *obuf;
 347
 348 /* We want obuf to be able to hold both MAX_DATA_COUNT+2 bytes as well as
 349  * AVAIL_OUT_SIZE(CHUNK_SIZE) bytes, so make sure that it's large enough. */
 350 #if MAX_DATA_COUNT+2 > AVAIL_OUT_SIZE(CHUNK_SIZE)
 351 #define OBUF_SIZE       (MAX_DATA_COUNT+2)
 352 #else
 353 #define OBUF_SIZE       AVAIL_OUT_SIZE(CHUNK_SIZE)
 354 #endif
 355
 356 /* Send a deflated token */
 357 static void
 358 send_deflated_token(int f, int32 token, struct map_struct *buf, OFF_T offset, int32 nb, int32 toklen)
 359 {
 360         static int init_done, flush_pending;
 361         int32 n, r;
 362
 363         if (last_token == -1) {
 364                 /* initialization */
 365                 if (!init_done) {
 366                         tx_strm.next_in = NULL;
 367                         tx_strm.zalloc = NULL;
 368                         tx_strm.zfree = NULL;
 369                         if (deflateInit2(&tx_strm, per_file_default_level,
 370                                          Z_DEFLATED, -15, 8,
 371                                          Z_DEFAULT_STRATEGY) != Z_OK) {
 372                                 rprintf(FERROR, "compression init failed\n");
 373                                 exit_cleanup(RERR_PROTOCOL);
 374                         }
 375                         obuf = new_array(char, OBUF_SIZE);
 376                         init_done = 1;
 377                 } else
 378                         deflateReset(&tx_strm);
 379                 last_run_end = 0;
 380                 run_start = token;
 381                 flush_pending = 0;
 382         } else if (last_token == -2) {
 383                 run_start = token;
 384         } else if (nb != 0 || token != last_token + 1 || token >= run_start + 65536) {
 385                 /* output previous run */
 386                 r = run_start - last_run_end;
 387                 n = last_token - run_start;
 388                 if (r >= 0 && r <= 63) {
 389                         write_byte(f, (n==0? TOKEN_REL: TOKENRUN_REL) + r);
 390                 } else {
 391                         write_byte(f, (n==0? TOKEN_LONG: TOKENRUN_LONG));
 392                         write_int(f, run_start);
 393                 }
 394                 if (n != 0) {
 395                         write_byte(f, n);
 396                         write_byte(f, n >> 8);
 397                 }
 398                 last_run_end = last_token;
 399                 run_start = token;
 400         }
 401
 402         last_token = token;
 403
 404         if (nb != 0 || flush_pending) {
 405                 /* deflate the data starting at offset */
 406                 int flush = Z_NO_FLUSH;
 407                 tx_strm.avail_in = 0;
 408                 tx_strm.avail_out = 0;
 409                 do {
 410                         if (tx_strm.avail_in == 0 && nb != 0) {
 411                                 /* give it some more input */
 412                                 n = MIN(nb, CHUNK_SIZE);
 413                                 tx_strm.next_in = (Bytef *)
 414                                         map_ptr(buf, offset, n);
 415                                 tx_strm.avail_in = n;
 416                                 nb -= n;
 417                                 offset += n;
 418                         }
 419                         if (tx_strm.avail_out == 0) {
 420                                 tx_strm.next_out = (Bytef *)(obuf + 2);
 421                                 tx_strm.avail_out = MAX_DATA_COUNT;
 422                                 if (flush != Z_NO_FLUSH) {
 423                                         /*
 424                                          * We left the last 4 bytes in the
 425                                          * buffer, in case they are the
 426                                          * last 4.  Move them to the front.
 427                                          */
 428                                         memcpy(tx_strm.next_out, obuf+MAX_DATA_COUNT-2, 4);
 429                                         tx_strm.next_out += 4;
 430                                         tx_strm.avail_out -= 4;
 431                                 }
 432                         }
 433                         if (nb == 0 && token != -2)
 434                                 flush = Z_SYNC_FLUSH;
 435                         r = deflate(&tx_strm, flush);
 436                         if (r != Z_OK) {
 437                                 rprintf(FERROR, "deflate returned %d\n", r);
 438                                 exit_cleanup(RERR_STREAMIO);
 439                         }
 440                         if (nb == 0 || tx_strm.avail_out == 0) {
 441                                 n = MAX_DATA_COUNT - tx_strm.avail_out;
 442                                 if (flush != Z_NO_FLUSH) {
 443                                         /*
 444                                          * We have to trim off the last 4
 445                                          * bytes of output when flushing
 446                                          * (they are just 0, 0, ff, ff).
 447                                          */
 448                                         n -= 4;
 449                                 }
 450                                 if (n > 0) {
 451                                         obuf[0] = DEFLATED_DATA + (n >> 8);
 452                                         obuf[1] = n;
 453                                         write_buf(f, obuf, n+2);
 454                                 }
 455                         }
 456                 } while (nb != 0 || tx_strm.avail_out == 0);
 457                 flush_pending = token == -2;
 458         }
 459
 460         if (token == -1) {
 461                 /* end of file - clean up */
 462                 write_byte(f, END_FLAG);
 463         } else if (token != -2 && do_compression == CPRES_ZLIB) {
 464                 /* Add the data in the current block to the compressor's
 465                  * history and hash table. */
 466                 do {
 467                         /* Break up long sections in the same way that
 468                          * see_deflate_token() does. */
 469                         int32 n1 = toklen > 0xffff ? 0xffff : toklen;
 470                         toklen -= n1;
 471                         tx_strm.next_in = (Bytef *)map_ptr(buf, offset, n1);
 472                         tx_strm.avail_in = n1;
 473                         if (protocol_version >= 31) /* Newer protocols avoid a data-duplicating bug */
 474                                 offset += n1;
 475                         tx_strm.next_out = (Bytef *) obuf;
 476                         tx_strm.avail_out = AVAIL_OUT_SIZE(CHUNK_SIZE);
 477                         r = deflate(&tx_strm, Z_INSERT_ONLY);
 478                         if (r != Z_OK || tx_strm.avail_in != 0) {
 479                                 rprintf(FERROR, "deflate on token returned %d (%d bytes left)\n",
 480                                         r, tx_strm.avail_in);
 481                                 exit_cleanup(RERR_STREAMIO);
 482                         }
 483                 } while (toklen > 0);
 484         }
 485 }
 486
 487 /* tells us what the receiver is in the middle of doing */
 488 static enum { r_init, r_idle, r_running, r_inflating, r_inflated } recv_state;
 489
 490 /* for inflating stuff */
 491 static z_stream rx_strm;
 492 static char *cbuf;
 493 static char *dbuf;
 494
 495 /* for decoding runs of tokens */
 496 static int32 rx_token;
 497 static int32 rx_run;
 498
 499 /* Receive a deflated token and inflate it */
 500 static int32 recv_deflated_token(int f, char **data)
 501 {
 502         static int init_done;
 503         static int32 saved_flag;
 504         int32 n, flag;
 505         int r;
 506
 507         for (;;) {
 508                 switch (recv_state) {
 509                 case r_init:
 510                         if (!init_done) {
 511                                 rx_strm.next_out = NULL;
 512                                 rx_strm.zalloc = NULL;
 513                                 rx_strm.zfree = NULL;
 514                                 if (inflateInit2(&rx_strm, -15) != Z_OK) {
 515                                         rprintf(FERROR, "inflate init failed\n");
 516                                         exit_cleanup(RERR_PROTOCOL);
 517                                 }
 518                                 cbuf = new_array(char, MAX_DATA_COUNT);
 519                                 dbuf = new_array(char, AVAIL_OUT_SIZE(CHUNK_SIZE));
 520                                 init_done = 1;
 521                         } else {
 522                                 inflateReset(&rx_strm);
 523                         }
 524                         recv_state = r_idle;
 525                         rx_token = 0;
 526                         break;
 527
 528                 case r_idle:
 529                 case r_inflated:
 530                         if (saved_flag) {
 531                                 flag = saved_flag & 0xff;
 532                                 saved_flag = 0;
 533                         } else
 534                                 flag = read_byte(f);
 535                         if ((flag & 0xC0) == DEFLATED_DATA) {
 536                                 n = ((flag & 0x3f) << 8) + read_byte(f);
 537                                 read_buf(f, cbuf, n);
 538                                 rx_strm.next_in = (Bytef *)cbuf;
 539                                 rx_strm.avail_in = n;
 540                                 recv_state = r_inflating;
 541                                 break;
 542                         }
 543                         if (recv_state == r_inflated) {
 544                                 /* check previous inflated stuff ended correctly */
 545                                 rx_strm.avail_in = 0;
 546                                 rx_strm.next_out = (Bytef *)dbuf;
 547                                 rx_strm.avail_out = AVAIL_OUT_SIZE(CHUNK_SIZE);
 548                                 r = inflate(&rx_strm, Z_SYNC_FLUSH);
 549                                 n = AVAIL_OUT_SIZE(CHUNK_SIZE) - rx_strm.avail_out;
 550                                 /*
 551                                  * Z_BUF_ERROR just means no progress was
 552                                  * made, i.e. the decompressor didn't have
 553                                  * any pending output for us.
 554                                  */
 555                                 if (r != Z_OK && r != Z_BUF_ERROR) {
 556                                         rprintf(FERROR, "inflate flush returned %d (%d bytes)\n",
 557                                                 r, n);
 558                                         exit_cleanup(RERR_STREAMIO);
 559                                 }
 560                                 if (n != 0 && r != Z_BUF_ERROR) {
 561                                         /* have to return some more data and
 562                                            save the flag for later. */
 563                                         saved_flag = flag + 0x10000;
 564                                         *data = dbuf;
 565                                         return n;
 566                                 }
 567                                 /*
 568                                  * At this point the decompressor should
 569                                  * be expecting to see the 0, 0, ff, ff bytes.
 570                                  */
 571                                 if (!inflateSyncPoint(&rx_strm)) {
 572                                         rprintf(FERROR, "decompressor lost sync!\n");
 573                                         exit_cleanup(RERR_STREAMIO);
 574                                 }
 575                                 rx_strm.avail_in = 4;
 576                                 rx_strm.next_in = (Bytef *)cbuf;
 577                                 cbuf[0] = cbuf[1] = 0;
 578                                 cbuf[2] = cbuf[3] = (char)0xff;
 579                                 inflate(&rx_strm, Z_SYNC_FLUSH);
 580                                 recv_state = r_idle;
 581                         }
 582                         if (flag == END_FLAG) {
 583                                 /* that's all folks */
 584                                 recv_state = r_init;
 585                                 return 0;
 586                         }
 587
 588                         /* here we have a token of some kind */
 589                         if (flag & TOKEN_REL) {
 590                                 rx_token += flag & 0x3f;
 591                                 flag >>= 6;
 592                         } else
 593                                 rx_token = read_int(f);
 594                         if (flag & 1) {
 595                                 rx_run = read_byte(f);
 596                                 rx_run += read_byte(f) << 8;
 597                                 recv_state = r_running;
 598                         }
 599                         return -1 - rx_token;
 600
 601                 case r_inflating:
 602                         rx_strm.next_out = (Bytef *)dbuf;
 603                         rx_strm.avail_out = AVAIL_OUT_SIZE(CHUNK_SIZE);
 604                         r = inflate(&rx_strm, Z_NO_FLUSH);
 605                         n = AVAIL_OUT_SIZE(CHUNK_SIZE) - rx_strm.avail_out;
 606                         if (r != Z_OK) {
 607                                 rprintf(FERROR, "inflate returned %d (%d bytes)\n", r, n);
 608                                 exit_cleanup(RERR_STREAMIO);
 609                         }
 610                         if (rx_strm.avail_in == 0)
 611                                 recv_state = r_inflated;
 612                         if (n != 0) {
 613                                 *data = dbuf;
 614                                 return n;
 615                         }
 616                         break;
 617
 618                 case r_running:
 619                         ++rx_token;
 620                         if (--rx_run == 0)
 621                                 recv_state = r_idle;
 622                         return -1 - rx_token;
 623                 }
 624         }
 625 }
 626
 627 /*
 628  * put the data corresponding to a token that we've just returned
 629  * from recv_deflated_token into the decompressor's history buffer.
 630  */
 631 static void see_deflate_token(char *buf, int32 len)
 632 {
 633         int r;
 634         int32 blklen;
 635         unsigned char hdr[5];
 636
 637         rx_strm.avail_in = 0;
 638         blklen = 0;
 639         hdr[0] = 0;
 640         do {
 641                 if (rx_strm.avail_in == 0 && len != 0) {
 642                         if (blklen == 0) {
 643                                 /* Give it a fake stored-block header. */
 644                                 rx_strm.next_in = (Bytef *)hdr;
 645                                 rx_strm.avail_in = 5;
 646                                 blklen = len;
 647                                 if (blklen > 0xffff)
 648                                         blklen = 0xffff;
 649                                 hdr[1] = blklen;
 650                                 hdr[2] = blklen >> 8;
 651                                 hdr[3] = ~hdr[1];
 652                                 hdr[4] = ~hdr[2];
 653                         } else {
 654                                 rx_strm.next_in = (Bytef *)buf;
 655                                 rx_strm.avail_in = blklen;
 656                                 if (protocol_version >= 31) /* Newer protocols avoid a data-duplicating bug */
 657                                         buf += blklen;
 658                                 len -= blklen;
 659                                 blklen = 0;
 660                         }
 661                 }
 662                 rx_strm.next_out = (Bytef *)dbuf;
 663                 rx_strm.avail_out = AVAIL_OUT_SIZE(CHUNK_SIZE);
 664                 r = inflate(&rx_strm, Z_SYNC_FLUSH);
 665                 if (r != Z_OK && r != Z_BUF_ERROR) {
 666                         rprintf(FERROR, "inflate (token) returned %d\n", r);
 667                         exit_cleanup(RERR_STREAMIO);
 668                 }
 669         } while (len || rx_strm.avail_out == 0);
 670 }
 671
 672 #ifdef SUPPORT_ZSTD
 673
 674 static ZSTD_inBuffer zstd_in_buff;
 675 static ZSTD_outBuffer zstd_out_buff;
 676 static ZSTD_CCtx *zstd_cctx;
 677
 678 static void send_zstd_token(int f, int32 token, struct map_struct *buf, OFF_T offset, int32 nb)
 679 {
 680         static int comp_init_done, flush_pending;
 681         ZSTD_EndDirective flush = ZSTD_e_continue;
 682         int32 n, r;
 683
 684         /* initialization */
 685         if (!comp_init_done) {
 686                 zstd_cctx = ZSTD_createCCtx();
 687                 if (!zstd_cctx) {
 688                         rprintf(FERROR, "compression init failed\n");
 689                         exit_cleanup(RERR_PROTOCOL);
 690                 }
 691
 692                 obuf = new_array(char, OBUF_SIZE);
 693
 694                 ZSTD_CCtx_setParameter(zstd_cctx, ZSTD_c_compressionLevel, do_compression_level);
 695                 zstd_out_buff.dst = obuf + 2;
 696
 697                 comp_init_done = 1;
 698         }
 699
 700         if (last_token == -1) {
 701                 last_run_end = 0;
 702                 run_start = token;
 703                 flush_pending = 0;
 704         } else if (last_token == -2) {
 705                 run_start = token;
 706         } else if (nb != 0 || token != last_token + 1 || token >= run_start + 65536) {
 707                 /* output previous run */
 708                 r = run_start - last_run_end;
 709                 n = last_token - run_start;
 710
 711                 if (r >= 0 && r <= 63) {
 712                         write_byte(f, (n==0? TOKEN_REL: TOKENRUN_REL) + r);
 713                 } else {
 714                         write_byte(f, (n==0? TOKEN_LONG: TOKENRUN_LONG));
 715                         write_int(f, run_start);
 716                 }
 717                 if (n != 0) {
 718                         write_byte(f, n);
 719                         write_byte(f, n >> 8);
 720                 }
 721                 last_run_end = last_token;
 722                 run_start = token;
 723         }
 724
 725         last_token = token;
 726
 727         if (nb || flush_pending) {
 728
 729                 zstd_in_buff.src = map_ptr(buf, offset, nb);
 730                 zstd_in_buff.size = nb;
 731                 zstd_in_buff.pos = 0;
 732
 733                 do {
 734                         if (zstd_out_buff.size == 0) {
 735                                 zstd_out_buff.size = MAX_DATA_COUNT;
 736                                 zstd_out_buff.pos = 0;
 737                         }
 738
 739                         /* File ended, flush */
 740                         if (token != -2)
 741                                 flush = ZSTD_e_flush;
 742
 743                         r = ZSTD_compressStream2(zstd_cctx, &zstd_out_buff, &zstd_in_buff, flush);
 744                         if (ZSTD_isError(r)) {
 745                                 rprintf(FERROR, "ZSTD_compressStream returned %d\n", r);
 746                                 exit_cleanup(RERR_STREAMIO);
 747                         }
 748
 749                         /*
 750                          * Nothing is sent if the buffer isn't full so avoid smaller
 751                          * transfers. If a file is finished then we flush the internal
 752                          * state and send a smaller buffer so that the remote side can
 753                          * finish the file.
 754                          */
 755                         if (zstd_out_buff.pos == zstd_out_buff.size || flush == ZSTD_e_flush) {
 756                                 n = zstd_out_buff.pos;
 757
 758                                 obuf[0] = DEFLATED_DATA + (n >> 8);
 759                                 obuf[1] = n;
 760                                 write_buf(f, obuf, n+2);
 761
 762                                 zstd_out_buff.size = 0;
 763                         }
 764                         /*
 765                          * Loop while the input buffer isn't full consumed or the
 766                          * internal state isn't fully flushed.
 767                          */
 768                 } while (zstd_in_buff.pos < zstd_in_buff.size || r > 0);
 769                 flush_pending = token == -2;
 770         }
 771
 772         if (token == -1) {
 773                 /* end of file - clean up */
 774                 write_byte(f, END_FLAG);
 775         }
 776 }
 777
 778 static ZSTD_DCtx *zstd_dctx;
 779
 780 static int32 recv_zstd_token(int f, char **data)
 781 {
 782         static int decomp_init_done;
 783         static int out_buffer_size;
 784         int32 n, flag;
 785         int r;
 786
 787         if (!decomp_init_done) {
 788                 zstd_dctx = ZSTD_createDCtx();
 789                 if (!zstd_dctx) {
 790                         rprintf(FERROR, "ZSTD_createDStream failed\n");
 791                         exit_cleanup(RERR_PROTOCOL);
 792                 }
 793
 794                 /* Output buffer fits two decompressed blocks */
 795                 out_buffer_size = ZSTD_DStreamOutSize() * 2;
 796                 cbuf = new_array(char, MAX_DATA_COUNT);
 797                 dbuf = new_array(char, out_buffer_size);
 798
 799                 zstd_in_buff.src = cbuf;
 800                 zstd_out_buff.dst = dbuf;
 801
 802                 decomp_init_done = 1;
 803         }
 804
 805         for (;;) {
 806                 switch (recv_state) {
 807                 case r_init:
 808                         recv_state = r_idle;
 809                         rx_token = 0;
 810                         break;
 811
 812                 case r_idle:
 813                         flag = read_byte(f);
 814                         if ((flag & 0xC0) == DEFLATED_DATA) {
 815                                 n = ((flag & 0x3f) << 8) + read_byte(f);
 816                                 read_buf(f, cbuf, n);
 817
 818                                 zstd_in_buff.size = n;
 819                                 zstd_in_buff.pos = 0;
 820
 821                                 recv_state = r_inflating;
 822                                 break;
 823                         }
 824
 825                         if (flag == END_FLAG) {
 826                                 /* that's all folks */
 827                                 recv_state = r_init;
 828                                 return 0;
 829                         }
 830                         /* here we have a token of some kind */
 831                         if (flag & TOKEN_REL) {
 832                                 rx_token += flag & 0x3f;
 833                                 flag >>= 6;
 834                         } else
 835                                 rx_token = read_int(f);
 836                         if (flag & 1) {
 837                                 rx_run = read_byte(f);
 838                                 rx_run += read_byte(f) << 8;
 839                                 recv_state = r_running;
 840                         }
 841                         return -1 - rx_token;
 842
 843                 case r_inflated: /* zstd doesn't get into this state */
 844                         break;
 845
 846                 case r_inflating:
 847                         zstd_out_buff.size = out_buffer_size;
 848                         zstd_out_buff.pos = 0;
 849
 850                         r = ZSTD_decompressStream(zstd_dctx, &zstd_out_buff, &zstd_in_buff);
 851                         n = zstd_out_buff.pos;
 852                         if (ZSTD_isError(r)) {
 853                                 rprintf(FERROR, "ZSTD decomp returned %d (%d bytes)\n", r, n);
 854                                 exit_cleanup(RERR_STREAMIO);
 855                         }
 856
 857                         /*
 858                          * If the input buffer is fully consumed and the output
 859                          * buffer is not full then next step is to read more
 860                          * data.
 861                          */
 862                         if (zstd_in_buff.size == zstd_in_buff.pos && n < out_buffer_size)
 863                                 recv_state = r_idle;
 864
 865                         if (n != 0) {
 866                                 *data = dbuf;
 867                                 return n;
 868                         }
 869                         break;
 870
 871                 case r_running:
 872                         ++rx_token;
 873                         if (--rx_run == 0)
 874                                 recv_state = r_idle;
 875                         return -1 - rx_token;
 876                 }
 877         }
 878 }
 879 #endif /* SUPPORT_ZSTD */
 880
 881 #ifdef SUPPORT_LZ4
 882 static void
 883 send_compressed_token(int f, int32 token, struct map_struct *buf, OFF_T offset, int32 nb)
 884 {
 885         static int init_done, flush_pending;
 886         int size = MAX(LZ4_compressBound(CHUNK_SIZE), MAX_DATA_COUNT+2);
 887         int32 n, r;
 888
 889         if (last_token == -1) {
 890                 if (!init_done) {
 891                         obuf = new_array(char, size);
 892                         init_done = 1;
 893                 }
 894                 last_run_end = 0;
 895                 run_start = token;
 896                 flush_pending = 0;
 897         } else if (last_token == -2) {
 898                 run_start = token;
 899         } else if (nb != 0 || token != last_token + 1 || token >= run_start + 65536) {
 900                 /* output previous run */
 901                 r = run_start - last_run_end;
 902                 n = last_token - run_start;
 903                 if (r >= 0 && r <= 63) {
 904                         write_byte(f, (n==0? TOKEN_REL: TOKENRUN_REL) + r);
 905                 } else {
 906                         write_byte(f, (n==0? TOKEN_LONG: TOKENRUN_LONG));
 907                         write_int(f, run_start);
 908                 }
 909                 if (n != 0) {
 910                         write_byte(f, n);
 911                         write_byte(f, n >> 8);
 912                 }
 913                 last_run_end = last_token;
 914                 run_start = token;
 915         }
 916
 917         last_token = token;
 918
 919         if (nb != 0 || flush_pending) {
 920                 int available_in, available_out = 0;
 921                 const char *next_in;
 922
 923                 do {
 924                         char *next_out = obuf + 2;
 925
 926                         if (available_out == 0) {
 927                                 available_in = MIN(nb, MAX_DATA_COUNT);
 928                                 next_in = map_ptr(buf, offset, available_in);
 929                         } else
 930                                 available_in /= 2;
 931
 932                         available_out = LZ4_compress_default(next_in, next_out, available_in, size - 2);
 933                         if (!available_out) {
 934                                 rprintf(FERROR, "compress returned %d\n", available_out);
 935                                 exit_cleanup(RERR_STREAMIO);
 936                         }
 937                         if (available_out <= MAX_DATA_COUNT) {
 938                                 obuf[0] = DEFLATED_DATA + (available_out >> 8);
 939                                 obuf[1] = available_out;
 940
 941                                 write_buf(f, obuf, available_out + 2);
 942
 943                                 available_out = 0;
 944                                 nb -= available_in;
 945                                 offset += available_in;
 946                         }
 947                 } while (nb != 0);
 948                 flush_pending = token == -2;
 949         }
 950         if (token == -1) {
 951                 /* end of file - clean up */
 952                 write_byte(f, END_FLAG);
 953         }
 954 }
 955
 956 static int32 recv_compressed_token(int f, char **data)
 957 {
 958         static int init_done;
 959         int32 n, flag;
 960         int size = MAX(LZ4_compressBound(CHUNK_SIZE), MAX_DATA_COUNT+2);
 961         static const char *next_in;
 962         static int avail_in;
 963         int avail_out;
 964
 965         for (;;) {
 966                 switch (recv_state) {
 967                 case r_init:
 968                         if (!init_done) {
 969                                 cbuf = new_array(char, MAX_DATA_COUNT);
 970                                 dbuf = new_array(char, size);
 971                                 init_done = 1;
 972                         }
 973                         recv_state = r_idle;
 974                         rx_token = 0;
 975                         break;
 976
 977                 case r_idle:
 978                         flag = read_byte(f);
 979                         if ((flag & 0xC0) == DEFLATED_DATA) {
 980                                 n = ((flag & 0x3f) << 8) + read_byte(f);
 981                                 read_buf(f, cbuf, n);
 982                                 next_in = (char *)cbuf;
 983                                 avail_in = n;
 984                                 recv_state = r_inflating;
 985                                 break;
 986                         }
 987
 988                         if (flag == END_FLAG) {
 989                                 /* that's all folks */
 990                                 recv_state = r_init;
 991                                 return 0;
 992                         }
 993
 994                         /* here we have a token of some kind */
 995                         if (flag & TOKEN_REL) {
 996                                 rx_token += flag & 0x3f;
 997                                 flag >>= 6;
 998                         } else
 999                                 rx_token = read_int(f);
1000                         if (flag & 1) {
1001                                 rx_run = read_byte(f);
1002                                 rx_run += read_byte(f) << 8;
1003                                 recv_state = r_running;
1004                         }
1005                         return -1 - rx_token;
1006
1007                 case r_inflating:
1008                         avail_out = LZ4_decompress_safe(next_in, dbuf, avail_in, size);
1009                         if (avail_out < 0) {
1010                                 rprintf(FERROR, "uncompress failed: %d\n", avail_out);
1011                                 exit_cleanup(RERR_STREAMIO);
1012                         }
1013                         recv_state = r_idle;
1014                         *data = dbuf;
1015                         return avail_out;
1016
1017                 case r_inflated: /* lz4 doesn't get into this state */
1018                         break;
1019
1020                 case r_running:
1021                         ++rx_token;
1022                         if (--rx_run == 0)
1023                                 recv_state = r_idle;
1024                         return -1 - rx_token;
1025                 }
1026         }
1027 }
1028 #endif /* SUPPORT_LZ4 */
1029
1030 /**
1031  * Transmit a verbatim buffer of length @p n followed by a token.
1032  * If token == -1 then we have reached EOF
1033  * If n == 0 then don't send a buffer
1034  */
1035 void send_token(int f, int32 token, struct map_struct *buf, OFF_T offset,
1036                 int32 n, int32 toklen)
1037 {
1038         switch (do_compression) {
1039         case CPRES_NONE:
1040                 simple_send_token(f, token, buf, offset, n);
1041                 break;
1042         case CPRES_ZLIB:
1043         case CPRES_ZLIBX:
1044                 send_deflated_token(f, token, buf, offset, n, toklen);
1045                 break;
1046 #ifdef SUPPORT_ZSTD
1047         case CPRES_ZSTD:
1048                 send_zstd_token(f, token, buf, offset, n);
1049                 break;
1050 #endif
1051 #ifdef SUPPORT_LZ4
1052         case CPRES_LZ4:
1053                 send_compressed_token(f, token, buf, offset, n);
1054                 break;
1055 #endif
1056         default:
1057                 NOISY_DEATH("Unknown do_compression value");
1058         }
1059 }
1060
1061 /*
1062  * receive a token or buffer from the other end. If the return value is >0 then
1063  * it is a data buffer of that length, and *data will point at the data.
1064  * if the return value is -i then it represents token i-1
1065  * if the return value is 0 then the end has been reached
1066  */
1067 int32 recv_token(int f, char **data)
1068 {
1069         switch (do_compression) {
1070         case CPRES_NONE:
1071                 return simple_recv_token(f,data);
1072         case CPRES_ZLIB:
1073         case CPRES_ZLIBX:
1074                 return recv_deflated_token(f, data);
1075 #ifdef SUPPORT_ZSTD
1076         case CPRES_ZSTD:
1077                 return recv_zstd_token(f, data);
1078 #endif
1079 #ifdef SUPPORT_LZ4
1080         case CPRES_LZ4:
1081                 return recv_compressed_token(f, data);
1082 #endif
1083         default:
1084                 NOISY_DEATH("Unknown do_compression value");
1085         }
1086 }
1087
1088 /*
1089  * look at the data corresponding to a token, if necessary
1090  */
1091 void see_token(char *data, int32 toklen)
1092 {
1093         switch (do_compression) {
1094         case CPRES_NONE:
1095                 break;
1096         case CPRES_ZLIB:
1097                 see_deflate_token(data, toklen);
1098                 break;
1099         case CPRES_ZLIBX:
1100                 break;
1101 #ifdef SUPPORT_ZSTD
1102         case CPRES_ZSTD:
1103                 break;
1104 #endif
1105 #ifdef SUPPORT_LZ4
1106         case CPRES_LZ4:
1107                 /*see_uncompressed_token(data, toklen);*/
1108                 break;
1109 #endif
1110         default:
1111                 NOISY_DEATH("Unknown do_compression value");
1112         }
1113 }