Python/marshal.c

   1
   2 /* Write Python objects to files and read them back.
   3    This is intended for writing and reading compiled Python code only;
   4    a true persistent storage facility would be much harder, since
   5    it would have to take circular links and sharing into account. */
   6
   7 #define PY_SSIZE_T_CLEAN
   8
   9 #include "Python.h"
  10 #include "longintrepr.h"
  11 #include "code.h"
  12 #include "marshal.h"
  13
  14 /* High water mark to determine when the marshalled object is dangerously deep
  15  * and risks coring the interpreter.  When the object stack gets this deep,
  16  * raise an exception instead of continuing.
  17  */
  18 #define MAX_MARSHAL_STACK_DEPTH 5000
  19
  20 #define TYPE_NULL               '0'
  21 #define TYPE_NONE               'N'
  22 #define TYPE_FALSE              'F'
  23 #define TYPE_TRUE               'T'
  24 #define TYPE_STOPITER           'S'
  25 #define TYPE_ELLIPSIS           '.'
  26 #define TYPE_INT                'i'
  27 #define TYPE_INT64              'I'
  28 #define TYPE_FLOAT              'f'
  29 #define TYPE_BINARY_FLOAT       'g'
  30 #define TYPE_COMPLEX            'x'
  31 #define TYPE_BINARY_COMPLEX     'y'
  32 #define TYPE_LONG               'l'
  33 #define TYPE_STRING             's'
  34 #define TYPE_INTERNED           't'
  35 #define TYPE_STRINGREF          'R'
  36 #define TYPE_TUPLE              '('
  37 #define TYPE_LIST               '['
  38 #define TYPE_DICT               '{'
  39 #define TYPE_CODE               'c'
  40 #define TYPE_UNICODE            'u'
  41 #define TYPE_UNKNOWN            '?'
  42 #define TYPE_SET                '<'
  43 #define TYPE_FROZENSET          '>'
  44
  45 typedef struct {
  46         FILE *fp;
  47         int error;
  48         int depth;
  49         /* If fp == NULL, the following are valid: */
  50         PyObject *str;
  51         char *ptr;
  52         char *end;
  53         PyObject *strings; /* dict on marshal, list on unmarshal */
  54         int version;
  55 } WFILE;
  56
  57 #define w_byte(c, p) if (((p)->fp)) putc((c), (p)->fp); \
  58                       else if ((p)->ptr != (p)->end) *(p)->ptr++ = (c); \
  59                            else w_more(c, p)
  60
  61 static void
  62 w_more(int c, WFILE *p)
  63 {
  64         Py_ssize_t size, newsize;
  65         if (p->str == NULL)
  66                 return; /* An error already occurred */
  67         size = PyString_Size(p->str);
  68         newsize = size + 1024;
  69         if (_PyString_Resize(&p->str, newsize) != 0) {
  70                 p->ptr = p->end = NULL;
  71         }
  72         else {
  73                 p->ptr = PyString_AS_STRING((PyStringObject *)p->str) + size;
  74                 p->end =
  75                         PyString_AS_STRING((PyStringObject *)p->str) + newsize;
  76                 *p->ptr++ = Py_SAFE_DOWNCAST(c, int, char);
  77         }
  78 }
  79
  80 static void
  81 w_string(char *s, int n, WFILE *p)
  82 {
  83         if (p->fp != NULL) {
  84                 fwrite(s, 1, n, p->fp);
  85         }
  86         else {
  87                 while (--n >= 0) {
  88                         w_byte(*s, p);
  89                         s++;
  90                 }
  91         }
  92 }
  93
  94 static void
  95 w_short(int x, WFILE *p)
  96 {
  97         w_byte((char)( x      & 0xff), p);
  98         w_byte((char)((x>> 8) & 0xff), p);
  99 }
 100
 101 static void
 102 w_long(long x, WFILE *p)
 103 {
 104         w_byte((char)( x      & 0xff), p);
 105         w_byte((char)((x>> 8) & 0xff), p);
 106         w_byte((char)((x>>16) & 0xff), p);
 107         w_byte((char)((x>>24) & 0xff), p);
 108 }
 109
 110 #if SIZEOF_LONG > 4
 111 static void
 112 w_long64(long x, WFILE *p)
 113 {
 114         w_long(x, p);
 115         w_long(x>>32, p);
 116 }
 117 #endif
 118
 119 static void
 120 w_object(PyObject *v, WFILE *p)
 121 {
 122         Py_ssize_t i, n;
 123
 124         p->depth++;
 125
 126         if (p->depth > MAX_MARSHAL_STACK_DEPTH) {
 127                 p->error = 2;
 128         }
 129         else if (v == NULL) {
 130                 w_byte(TYPE_NULL, p);
 131         }
 132         else if (v == Py_None) {
 133                 w_byte(TYPE_NONE, p);
 134         }
 135         else if (v == PyExc_StopIteration) {
 136                 w_byte(TYPE_STOPITER, p);
 137         }
 138         else if (v == Py_Ellipsis) {
 139                 w_byte(TYPE_ELLIPSIS, p);
 140         }
 141         else if (v == Py_False) {
 142                 w_byte(TYPE_FALSE, p);
 143         }
 144         else if (v == Py_True) {
 145                 w_byte(TYPE_TRUE, p);
 146         }
 147         else if (PyInt_Check(v)) {
 148                 long x = PyInt_AS_LONG((PyIntObject *)v);
 149 #if SIZEOF_LONG > 4
 150                 long y = Py_ARITHMETIC_RIGHT_SHIFT(long, x, 31);
 151                 if (y && y != -1) {
 152                         w_byte(TYPE_INT64, p);
 153                         w_long64(x, p);
 154                 }
 155                 else
 156 #endif
 157                         {
 158                         w_byte(TYPE_INT, p);
 159                         w_long(x, p);
 160                 }
 161         }
 162         else if (PyLong_Check(v)) {
 163                 PyLongObject *ob = (PyLongObject *)v;
 164                 w_byte(TYPE_LONG, p);
 165                 n = ob->ob_size;
 166                 w_long((long)n, p);
 167                 if (n < 0)
 168                         n = -n;
 169                 for (i = 0; i < n; i++)
 170                         w_short(ob->ob_digit[i], p);
 171         }
 172         else if (PyFloat_Check(v)) {
 173                 if (p->version > 1) {
 174                         unsigned char buf[8];
 175                         if (_PyFloat_Pack8(PyFloat_AsDouble(v),
 176                                            buf, 1) < 0) {
 177                                 p->error = 1;
 178                                 return;
 179                         }
 180                         w_byte(TYPE_BINARY_FLOAT, p);
 181                         w_string((char*)buf, 8, p);
 182                 }
 183                 else {
 184                         char buf[256]; /* Plenty to format any double */
 185                         PyFloat_AsReprString(buf, (PyFloatObject *)v);
 186                         n = strlen(buf);
 187                         w_byte(TYPE_FLOAT, p);
 188                         w_byte((int)n, p);
 189                         w_string(buf, (int)n, p);
 190                 }
 191         }
 192 #ifndef WITHOUT_COMPLEX
 193         else if (PyComplex_Check(v)) {
 194                 if (p->version > 1) {
 195                         unsigned char buf[8];
 196                         if (_PyFloat_Pack8(PyComplex_RealAsDouble(v),
 197                                            buf, 1) < 0) {
 198                                 p->error = 1;
 199                                 return;
 200                         }
 201                         w_byte(TYPE_BINARY_COMPLEX, p);
 202                         w_string((char*)buf, 8, p);
 203                         if (_PyFloat_Pack8(PyComplex_ImagAsDouble(v),
 204                                            buf, 1) < 0) {
 205                                 p->error = 1;
 206                                 return;
 207                         }
 208                         w_string((char*)buf, 8, p);
 209                 }
 210                 else {
 211                         char buf[256]; /* Plenty to format any double */
 212                         PyFloatObject *temp;
 213                         w_byte(TYPE_COMPLEX, p);
 214                         temp = (PyFloatObject*)PyFloat_FromDouble(
 215                                 PyComplex_RealAsDouble(v));
 216                         if (!temp) {
 217                                 p->error = 1;
 218                                 return;
 219                         }
 220                         PyFloat_AsReprString(buf, temp);
 221                         Py_DECREF(temp);
 222                         n = strlen(buf);
 223                         w_byte((int)n, p);
 224                         w_string(buf, (int)n, p);
 225                         temp = (PyFloatObject*)PyFloat_FromDouble(
 226                                 PyComplex_ImagAsDouble(v));
 227                         if (!temp) {
 228                                 p->error = 1;
 229                                 return;
 230                         }
 231                         PyFloat_AsReprString(buf, temp);
 232                         Py_DECREF(temp);
 233                         n = strlen(buf);
 234                         w_byte((int)n, p);
 235                         w_string(buf, (int)n, p);
 236                 }
 237         }
 238 #endif
 239         else if (PyString_Check(v)) {
 240                 if (p->strings && PyString_CHECK_INTERNED(v)) {
 241                         PyObject *o = PyDict_GetItem(p->strings, v);
 242                         if (o) {
 243                                 long w = PyInt_AsLong(o);
 244                                 w_byte(TYPE_STRINGREF, p);
 245                                 w_long(w, p);
 246                                 goto exit;
 247                         }
 248                         else {
 249                                 o = PyInt_FromSsize_t(PyDict_Size(p->strings));
 250                                 PyDict_SetItem(p->strings, v, o);
 251                                 Py_DECREF(o);
 252                                 w_byte(TYPE_INTERNED, p);
 253                         }
 254                 }
 255                 else {
 256                         w_byte(TYPE_STRING, p);
 257                 }
 258                 n = PyString_GET_SIZE(v);
 259                 if (n > INT_MAX) {
 260                         /* huge strings are not supported */
 261                         p->depth--;
 262                         p->error = 1;
 263                         return;
 264                 }
 265                 w_long((long)n, p);
 266                 w_string(PyString_AS_STRING(v), (int)n, p);
 267         }
 268 #ifdef Py_USING_UNICODE
 269         else if (PyUnicode_Check(v)) {
 270                 PyObject *utf8;
 271                 utf8 = PyUnicode_AsUTF8String(v);
 272                 if (utf8 == NULL) {
 273                         p->depth--;
 274                         p->error = 1;
 275                         return;
 276                 }
 277                 w_byte(TYPE_UNICODE, p);
 278                 n = PyString_GET_SIZE(utf8);
 279                 if (n > INT_MAX) {
 280                         p->depth--;
 281                         p->error = 1;
 282                         return;
 283                 }
 284                 w_long((long)n, p);
 285                 w_string(PyString_AS_STRING(utf8), (int)n, p);
 286                 Py_DECREF(utf8);
 287         }
 288 #endif
 289         else if (PyTuple_Check(v)) {
 290                 w_byte(TYPE_TUPLE, p);
 291                 n = PyTuple_Size(v);
 292                 w_long((long)n, p);
 293                 for (i = 0; i < n; i++) {
 294                         w_object(PyTuple_GET_ITEM(v, i), p);
 295                 }
 296         }
 297         else if (PyList_Check(v)) {
 298                 w_byte(TYPE_LIST, p);
 299                 n = PyList_GET_SIZE(v);
 300                 w_long((long)n, p);
 301                 for (i = 0; i < n; i++) {
 302                         w_object(PyList_GET_ITEM(v, i), p);
 303                 }
 304         }
 305         else if (PyDict_Check(v)) {
 306                 Py_ssize_t pos;
 307                 PyObject *key, *value;
 308                 w_byte(TYPE_DICT, p);
 309                 /* This one is NULL object terminated! */
 310                 pos = 0;
 311                 while (PyDict_Next(v, &pos, &key, &value)) {
 312                         w_object(key, p);
 313                         w_object(value, p);
 314                 }
 315                 w_object((PyObject *)NULL, p);
 316         }
 317         else if (PyAnySet_Check(v)) {
 318                 PyObject *value, *it;
 319
 320                 if (PyObject_TypeCheck(v, &PySet_Type))
 321                         w_byte(TYPE_SET, p);
 322                 else
 323                         w_byte(TYPE_FROZENSET, p);
 324                 n = PyObject_Size(v);
 325                 if (n == -1) {
 326                         p->depth--;
 327                         p->error = 1;
 328                         return;
 329                 }
 330                 w_long((long)n, p);
 331                 it = PyObject_GetIter(v);
 332                 if (it == NULL) {
 333                         p->depth--;
 334                         p->error = 1;
 335                         return;
 336                 }
 337                 while ((value = PyIter_Next(it)) != NULL) {
 338                         w_object(value, p);
 339                         Py_DECREF(value);
 340                 }
 341                 Py_DECREF(it);
 342                 if (PyErr_Occurred()) {
 343                         p->depth--;
 344                         p->error = 1;
 345                         return;
 346                 }
 347         }
 348         else if (PyCode_Check(v)) {
 349                 PyCodeObject *co = (PyCodeObject *)v;
 350                 w_byte(TYPE_CODE, p);
 351                 w_long(co->co_argcount, p);
 352                 w_long(co->co_nlocals, p);
 353                 w_long(co->co_stacksize, p);
 354                 w_long(co->co_flags, p);
 355                 w_object(co->co_code, p);
 356                 w_object(co->co_consts, p);
 357                 w_object(co->co_names, p);
 358                 w_object(co->co_varnames, p);
 359                 w_object(co->co_freevars, p);
 360                 w_object(co->co_cellvars, p);
 361                 w_object(co->co_filename, p);
 362                 w_object(co->co_name, p);
 363                 w_long(co->co_firstlineno, p);
 364                 w_object(co->co_lnotab, p);
 365         }
 366         else if (PyObject_CheckReadBuffer(v)) {
 367                 /* Write unknown buffer-style objects as a string */
 368                 char *s;
 369                 PyBufferProcs *pb = v->ob_type->tp_as_buffer;
 370                 w_byte(TYPE_STRING, p);
 371                 n = (*pb->bf_getreadbuffer)(v, 0, (void **)&s);
 372                 if (n > INT_MAX) {
 373                         p->depth--;
 374                         p->error = 1;
 375                         return;
 376                 }
 377                 w_long((long)n, p);
 378                 w_string(s, (int)n, p);
 379         }
 380         else {
 381                 w_byte(TYPE_UNKNOWN, p);
 382                 p->error = 1;
 383         }
 384    exit:
 385         p->depth--;
 386 }
 387
 388 /* version currently has no effect for writing longs. */
 389 void
 390 PyMarshal_WriteLongToFile(long x, FILE *fp, int version)
 391 {
 392         WFILE wf;
 393         wf.fp = fp;
 394         wf.error = 0;
 395         wf.depth = 0;
 396         wf.strings = NULL;
 397         wf.version = version;
 398         w_long(x, &wf);
 399 }
 400
 401 void
 402 PyMarshal_WriteObjectToFile(PyObject *x, FILE *fp, int version)
 403 {
 404         WFILE wf;
 405         wf.fp = fp;
 406         wf.error = 0;
 407         wf.depth = 0;
 408         wf.strings = (version > 0) ? PyDict_New() : NULL;
 409         wf.version = version;
 410         w_object(x, &wf);
 411         Py_XDECREF(wf.strings);
 412 }
 413
 414 typedef WFILE RFILE; /* Same struct with different invariants */
 415
 416 #define rs_byte(p) (((p)->ptr != (p)->end) ? (unsigned char)*(p)->ptr++ : EOF)
 417
 418 #define r_byte(p) ((p)->fp ? getc((p)->fp) : rs_byte(p))
 419
 420 static int
 421 r_string(char *s, int n, RFILE *p)
 422 {
 423         if (p->fp != NULL)
 424                 /* The result fits into int because it must be <=n. */
 425                 return (int)fread(s, 1, n, p->fp);
 426         if (p->end - p->ptr < n)
 427                 n = (int)(p->end - p->ptr);
 428         memcpy(s, p->ptr, n);
 429         p->ptr += n;
 430         return n;
 431 }
 432
 433 static int
 434 r_short(RFILE *p)
 435 {
 436         register short x;
 437         x = r_byte(p);
 438         x |= r_byte(p) << 8;
 439         /* Sign-extension, in case short greater than 16 bits */
 440         x |= -(x & 0x8000);
 441         return x;
 442 }
 443
 444 static long
 445 r_long(RFILE *p)
 446 {
 447         register long x;
 448         register FILE *fp = p->fp;
 449         if (fp) {
 450                 x = getc(fp);
 451                 x |= (long)getc(fp) << 8;
 452                 x |= (long)getc(fp) << 16;
 453                 x |= (long)getc(fp) << 24;
 454         }
 455         else {
 456                 x = rs_byte(p);
 457                 x |= (long)rs_byte(p) << 8;
 458                 x |= (long)rs_byte(p) << 16;
 459                 x |= (long)rs_byte(p) << 24;
 460         }
 461 #if SIZEOF_LONG > 4
 462         /* Sign extension for 64-bit machines */
 463         x |= -(x & 0x80000000L);
 464 #endif
 465         return x;
 466 }
 467
 468 /* r_long64 deals with the TYPE_INT64 code.  On a machine with
 469    sizeof(long) > 4, it returns a Python int object, else a Python long
 470    object.  Note that w_long64 writes out TYPE_INT if 32 bits is enough,
 471    so there's no inefficiency here in returning a PyLong on 32-bit boxes
 472    for everything written via TYPE_INT64 (i.e., if an int is written via
 473    TYPE_INT64, it *needs* more than 32 bits).
 474 */
 475 static PyObject *
 476 r_long64(RFILE *p)
 477 {
 478         long lo4 = r_long(p);
 479         long hi4 = r_long(p);
 480 #if SIZEOF_LONG > 4
 481         long x = (hi4 << 32) | (lo4 & 0xFFFFFFFFL);
 482         return PyInt_FromLong(x);
 483 #else
 484         unsigned char buf[8];
 485         int one = 1;
 486         int is_little_endian = (int)*(char*)&one;
 487         if (is_little_endian) {
 488                 memcpy(buf, &lo4, 4);
 489                 memcpy(buf+4, &hi4, 4);
 490         }
 491         else {
 492                 memcpy(buf, &hi4, 4);
 493                 memcpy(buf+4, &lo4, 4);
 494         }
 495         return _PyLong_FromByteArray(buf, 8, is_little_endian, 1);
 496 #endif
 497 }
 498
 499 static PyObject *
 500 r_object(RFILE *p)
 501 {
 502         /* NULL is a valid return value, it does not necessarily means that
 503            an exception is set. */
 504         PyObject *v, *v2, *v3;
 505         long i, n;
 506         int type = r_byte(p);
 507
 508         switch (type) {
 509
 510         case EOF:
 511                 PyErr_SetString(PyExc_EOFError,
 512                                 "EOF read where object expected");
 513                 return NULL;
 514
 515         case TYPE_NULL:
 516                 return NULL;
 517
 518         case TYPE_NONE:
 519                 Py_INCREF(Py_None);
 520                 return Py_None;
 521
 522         case TYPE_STOPITER:
 523                 Py_INCREF(PyExc_StopIteration);
 524                 return PyExc_StopIteration;
 525
 526         case TYPE_ELLIPSIS:
 527                 Py_INCREF(Py_Ellipsis);
 528                 return Py_Ellipsis;
 529
 530         case TYPE_FALSE:
 531                 Py_INCREF(Py_False);
 532                 return Py_False;
 533
 534         case TYPE_TRUE:
 535                 Py_INCREF(Py_True);
 536                 return Py_True;
 537
 538         case TYPE_INT:
 539                 return PyInt_FromLong(r_long(p));
 540
 541         case TYPE_INT64:
 542                 return r_long64(p);
 543
 544         case TYPE_LONG:
 545                 {
 546                         int size;
 547                         PyLongObject *ob;
 548                         n = r_long(p);
 549                         if (n < -INT_MAX || n > INT_MAX) {
 550                                 PyErr_SetString(PyExc_ValueError,
 551                                                 "bad marshal data");
 552                                 return NULL;
 553                         }
 554                         size = n<0 ? -n : n;
 555                         ob = _PyLong_New(size);
 556                         if (ob == NULL)
 557                                 return NULL;
 558                         ob->ob_size = n;
 559                         for (i = 0; i < size; i++) {
 560                                 int digit = r_short(p);
 561                                 if (digit < 0) {
 562                                         Py_DECREF(ob);
 563                                         PyErr_SetString(PyExc_ValueError,
 564                                                         "bad marshal data");
 565                                         return NULL;
 566                                 }
 567                                 ob->ob_digit[i] = digit;
 568                         }
 569                         return (PyObject *)ob;
 570                 }
 571
 572         case TYPE_FLOAT:
 573                 {
 574                         char buf[256];
 575                         double dx;
 576                         n = r_byte(p);
 577                         if (n == EOF || r_string(buf, (int)n, p) != n) {
 578                                 PyErr_SetString(PyExc_EOFError,
 579                                         "EOF read where object expected");
 580                                 return NULL;
 581                         }
 582                         buf[n] = '\0';
 583                         PyFPE_START_PROTECT("atof", return 0)
 584                         dx = PyOS_ascii_atof(buf);
 585                         PyFPE_END_PROTECT(dx)
 586                         return PyFloat_FromDouble(dx);
 587                 }
 588
 589         case TYPE_BINARY_FLOAT:
 590                 {
 591                         unsigned char buf[8];
 592                         double x;
 593                         if (r_string((char*)buf, 8, p) != 8) {
 594                                 PyErr_SetString(PyExc_EOFError,
 595                                         "EOF read where object expected");
 596                                 return NULL;
 597                         }
 598                         x = _PyFloat_Unpack8(buf, 1);
 599                         if (x == -1.0 && PyErr_Occurred()) {
 600                                 return NULL;
 601                         }
 602                         return PyFloat_FromDouble(x);
 603                 }
 604
 605 #ifndef WITHOUT_COMPLEX
 606         case TYPE_COMPLEX:
 607                 {
 608                         char buf[256];
 609                         Py_complex c;
 610                         n = r_byte(p);
 611                         if (n == EOF || r_string(buf, (int)n, p) != n) {
 612                                 PyErr_SetString(PyExc_EOFError,
 613                                         "EOF read where object expected");
 614                                 return NULL;
 615                         }
 616                         buf[n] = '\0';
 617                         PyFPE_START_PROTECT("atof", return 0)
 618                         c.real = PyOS_ascii_atof(buf);
 619                         PyFPE_END_PROTECT(c)
 620                         n = r_byte(p);
 621                         if (n == EOF || r_string(buf, (int)n, p) != n) {
 622                                 PyErr_SetString(PyExc_EOFError,
 623                                         "EOF read where object expected");
 624                                 return NULL;
 625                         }
 626                         buf[n] = '\0';
 627                         PyFPE_START_PROTECT("atof", return 0)
 628                         c.imag = PyOS_ascii_atof(buf);
 629                         PyFPE_END_PROTECT(c)
 630                         return PyComplex_FromCComplex(c);
 631                 }
 632
 633         case TYPE_BINARY_COMPLEX:
 634                 {
 635                         unsigned char buf[8];
 636                         Py_complex c;
 637                         if (r_string((char*)buf, 8, p) != 8) {
 638                                 PyErr_SetString(PyExc_EOFError,
 639                                         "EOF read where object expected");
 640                                 return NULL;
 641                         }
 642                         c.real = _PyFloat_Unpack8(buf, 1);
 643                         if (c.real == -1.0 && PyErr_Occurred()) {
 644                                 return NULL;
 645                         }
 646                         if (r_string((char*)buf, 8, p) != 8) {
 647                                 PyErr_SetString(PyExc_EOFError,
 648                                         "EOF read where object expected");
 649                                 return NULL;
 650                         }
 651                         c.imag = _PyFloat_Unpack8(buf, 1);
 652                         if (c.imag == -1.0 && PyErr_Occurred()) {
 653                                 return NULL;
 654                         }
 655                         return PyComplex_FromCComplex(c);
 656                 }
 657 #endif
 658
 659         case TYPE_INTERNED:
 660         case TYPE_STRING:
 661                 n = r_long(p);
 662                 if (n < 0 || n > INT_MAX) {
 663                         PyErr_SetString(PyExc_ValueError, "bad marshal data");
 664                         return NULL;
 665                 }
 666                 v = PyString_FromStringAndSize((char *)NULL, n);
 667                 if (v == NULL)
 668                         return v;
 669                 if (r_string(PyString_AS_STRING(v), (int)n, p) != n) {
 670                         Py_DECREF(v);
 671                         PyErr_SetString(PyExc_EOFError,
 672                                         "EOF read where object expected");
 673                         return NULL;
 674                 }
 675                 if (type == TYPE_INTERNED) {
 676                         PyString_InternInPlace(&v);
 677                         PyList_Append(p->strings, v);
 678                 }
 679                 return v;
 680
 681         case TYPE_STRINGREF:
 682                 n = r_long(p);
 683                 if (n < 0 || n >= PyList_GET_SIZE(p->strings)) {
 684                         PyErr_SetString(PyExc_ValueError, "bad marshal data");
 685                         return NULL;
 686                 }
 687                 v = PyList_GET_ITEM(p->strings, n);
 688                 Py_INCREF(v);
 689                 return v;
 690
 691 #ifdef Py_USING_UNICODE
 692         case TYPE_UNICODE:
 693             {
 694                 char *buffer;
 695
 696                 n = r_long(p);
 697                 if (n < 0 || n > INT_MAX) {
 698                         PyErr_SetString(PyExc_ValueError, "bad marshal data");
 699                         return NULL;
 700                 }
 701                 buffer = PyMem_NEW(char, n);
 702                 if (buffer == NULL)
 703                         return PyErr_NoMemory();
 704                 if (r_string(buffer, (int)n, p) != n) {
 705                         PyMem_DEL(buffer);
 706                         PyErr_SetString(PyExc_EOFError,
 707                                 "EOF read where object expected");
 708                         return NULL;
 709                 }
 710                 v = PyUnicode_DecodeUTF8(buffer, n, NULL);
 711                 PyMem_DEL(buffer);
 712                 return v;
 713             }
 714 #endif
 715
 716         case TYPE_TUPLE:
 717                 n = r_long(p);
 718                 if (n < 0 || n > INT_MAX) {
 719                         PyErr_SetString(PyExc_ValueError, "bad marshal data");
 720                         return NULL;
 721                 }
 722                 v = PyTuple_New((int)n);
 723                 if (v == NULL)
 724                         return v;
 725                 for (i = 0; i < n; i++) {
 726                         v2 = r_object(p);
 727                         if ( v2 == NULL ) {
 728                                 if (!PyErr_Occurred())
 729                                         PyErr_SetString(PyExc_TypeError,
 730                                                 "NULL object in marshal data");
 731                                 Py_DECREF(v);
 732                                 v = NULL;
 733                                 break;
 734                         }
 735                         PyTuple_SET_ITEM(v, (int)i, v2);
 736                 }
 737                 return v;
 738
 739         case TYPE_LIST:
 740                 n = r_long(p);
 741                 if (n < 0 || n > INT_MAX) {
 742                         PyErr_SetString(PyExc_ValueError, "bad marshal data");
 743                         return NULL;
 744                 }
 745                 v = PyList_New((int)n);
 746                 if (v == NULL)
 747                         return v;
 748                 for (i = 0; i < n; i++) {
 749                         v2 = r_object(p);
 750                         if ( v2 == NULL ) {
 751                                 if (!PyErr_Occurred())
 752                                         PyErr_SetString(PyExc_TypeError,
 753                                                 "NULL object in marshal data");
 754                                 Py_DECREF(v);
 755                                 v = NULL;
 756                                 break;
 757                         }
 758                         PyList_SetItem(v, (int)i, v2);
 759                 }
 760                 return v;
 761
 762         case TYPE_DICT:
 763                 v = PyDict_New();
 764                 if (v == NULL)
 765                         return NULL;
 766                 for (;;) {
 767                         PyObject *key, *val;
 768                         key = r_object(p);
 769                         if (key == NULL)
 770                                 break;
 771                         val = r_object(p);
 772                         if (val != NULL)
 773                                 PyDict_SetItem(v, key, val);
 774                         Py_DECREF(key);
 775                         Py_XDECREF(val);
 776                 }
 777                 if (PyErr_Occurred()) {
 778                         Py_DECREF(v);
 779                         v = NULL;
 780                 }
 781                 return v;
 782
 783         case TYPE_SET:
 784         case TYPE_FROZENSET:
 785                 n = r_long(p);
 786                 if (n < 0) {
 787                         PyErr_SetString(PyExc_ValueError, "bad marshal data");
 788                         return NULL;
 789                 }
 790                 v = PyTuple_New((int)n);
 791                 if (v == NULL)
 792                         return v;
 793                 for (i = 0; i < n; i++) {
 794                         v2 = r_object(p);
 795                         if ( v2 == NULL ) {
 796                                 if (!PyErr_Occurred())
 797                                         PyErr_SetString(PyExc_TypeError,
 798                                                 "NULL object in marshal data");
 799                                 Py_DECREF(v);
 800                                 v = NULL;
 801                                 break;
 802                         }
 803                         PyTuple_SET_ITEM(v, (int)i, v2);
 804                 }
 805                 if (v == NULL)
 806                         return v;
 807                 if (type == TYPE_SET)
 808                         v3 = PySet_New(v);
 809                 else
 810                         v3 = PyFrozenSet_New(v);
 811                 Py_DECREF(v);
 812                 return v3;
 813
 814         case TYPE_CODE:
 815                 if (PyEval_GetRestricted()) {
 816                         PyErr_SetString(PyExc_RuntimeError,
 817                                 "cannot unmarshal code objects in "
 818                                 "restricted execution mode");
 819                         return NULL;
 820                 }
 821                 else {
 822                         int argcount;
 823                         int nlocals;
 824                         int stacksize;
 825                         int flags;
 826                         PyObject *code = NULL;
 827                         PyObject *consts = NULL;
 828                         PyObject *names = NULL;
 829                         PyObject *varnames = NULL;
 830                         PyObject *freevars = NULL;
 831                         PyObject *cellvars = NULL;
 832                         PyObject *filename = NULL;
 833                         PyObject *name = NULL;
 834                         int firstlineno;
 835                         PyObject *lnotab = NULL;
 836
 837                         v = NULL;
 838
 839                         /* XXX ignore long->int overflows for now */
 840                         argcount = (int)r_long(p);
 841                         nlocals = (int)r_long(p);
 842                         stacksize = (int)r_long(p);
 843                         flags = (int)r_long(p);
 844                         code = r_object(p);
 845                         if (code == NULL)
 846                                 goto code_error;
 847                         consts = r_object(p);
 848                         if (consts == NULL)
 849                                 goto code_error;
 850                         names = r_object(p);
 851                         if (names == NULL)
 852                                 goto code_error;
 853                         varnames = r_object(p);
 854                         if (varnames == NULL)
 855                                 goto code_error;
 856                         freevars = r_object(p);
 857                         if (freevars == NULL)
 858                                 goto code_error;
 859                         cellvars = r_object(p);
 860                         if (cellvars == NULL)
 861                                 goto code_error;
 862                         filename = r_object(p);
 863                         if (filename == NULL)
 864                                 goto code_error;
 865                         name = r_object(p);
 866                         if (name == NULL)
 867                                 goto code_error;
 868                         firstlineno = (int)r_long(p);
 869                         lnotab = r_object(p);
 870                         if (lnotab == NULL)
 871                                 goto code_error;
 872
 873                         v = (PyObject *) PyCode_New(
 874                                         argcount, nlocals, stacksize, flags,
 875                                         code, consts, names, varnames,
 876                                         freevars, cellvars, filename, name,
 877                                         firstlineno, lnotab);
 878
 879                   code_error:
 880                         Py_XDECREF(code);
 881                         Py_XDECREF(consts);
 882                         Py_XDECREF(names);
 883                         Py_XDECREF(varnames);
 884                         Py_XDECREF(freevars);
 885                         Py_XDECREF(cellvars);
 886                         Py_XDECREF(filename);
 887                         Py_XDECREF(name);
 888                         Py_XDECREF(lnotab);
 889
 890                 }
 891                 return v;
 892
 893         default:
 894                 /* Bogus data got written, which isn't ideal.
 895                    This will let you keep working and recover. */
 896                 PyErr_SetString(PyExc_ValueError, "bad marshal data");
 897                 return NULL;
 898
 899         }
 900 }
 901
 902 static PyObject *
 903 read_object(RFILE *p)
 904 {
 905         PyObject *v;
 906         if (PyErr_Occurred()) {
 907                 fprintf(stderr, "XXX readobject called with exception set\n");
 908                 return NULL;
 909         }
 910         v = r_object(p);
 911         if (v == NULL && !PyErr_Occurred())
 912                 PyErr_SetString(PyExc_TypeError, "NULL object in marshal data");
 913         return v;
 914 }
 915
 916 int
 917 PyMarshal_ReadShortFromFile(FILE *fp)
 918 {
 919         RFILE rf;
 920         assert(fp);
 921         rf.fp = fp;
 922         rf.strings = NULL;
 923         rf.end = rf.ptr = NULL;
 924         return r_short(&rf);
 925 }
 926
 927 long
 928 PyMarshal_ReadLongFromFile(FILE *fp)
 929 {
 930         RFILE rf;
 931         rf.fp = fp;
 932         rf.strings = NULL;
 933         return r_long(&rf);
 934 }
 935
 936 #ifdef HAVE_FSTAT
 937 /* Return size of file in bytes; < 0 if unknown. */
 938 static off_t
 939 getfilesize(FILE *fp)
 940 {
 941         struct stat st;
 942         if (fstat(fileno(fp), &st) != 0)
 943                 return -1;
 944         else
 945                 return st.st_size;
 946 }
 947 #endif
 948
 949 /* If we can get the size of the file up-front, and it's reasonably small,
 950  * read it in one gulp and delegate to ...FromString() instead.  Much quicker
 951  * than reading a byte at a time from file; speeds .pyc imports.
 952  * CAUTION:  since this may read the entire remainder of the file, don't
 953  * call it unless you know you're done with the file.
 954  */
 955 PyObject *
 956 PyMarshal_ReadLastObjectFromFile(FILE *fp)
 957 {
 958 /* 75% of 2.1's .pyc files can exploit SMALL_FILE_LIMIT.
 959  * REASONABLE_FILE_LIMIT is by defn something big enough for Tkinter.pyc.
 960  */
 961 #define SMALL_FILE_LIMIT (1L << 14)
 962 #define REASONABLE_FILE_LIMIT (1L << 18)
 963 #ifdef HAVE_FSTAT
 964         off_t filesize;
 965 #endif
 966 #ifdef HAVE_FSTAT
 967         filesize = getfilesize(fp);
 968         if (filesize > 0) {
 969                 char buf[SMALL_FILE_LIMIT];
 970                 char* pBuf = NULL;
 971                 if (filesize <= SMALL_FILE_LIMIT)
 972                         pBuf = buf;
 973                 else if (filesize <= REASONABLE_FILE_LIMIT)
 974                         pBuf = (char *)PyMem_MALLOC(filesize);
 975                 if (pBuf != NULL) {
 976                         PyObject* v;
 977                         size_t n;
 978                         /* filesize must fit into an int, because it
 979                            is smaller than REASONABLE_FILE_LIMIT */
 980                         n = fread(pBuf, 1, (int)filesize, fp);
 981                         v = PyMarshal_ReadObjectFromString(pBuf, n);
 982                         if (pBuf != buf)
 983                                 PyMem_FREE(pBuf);
 984                         return v;
 985                 }
 986
 987         }
 988 #endif
 989         /* We don't have fstat, or we do but the file is larger than
 990          * REASONABLE_FILE_LIMIT or malloc failed -- read a byte at a time.
 991          */
 992         return PyMarshal_ReadObjectFromFile(fp);
 993
 994 #undef SMALL_FILE_LIMIT
 995 #undef REASONABLE_FILE_LIMIT
 996 }
 997
 998 PyObject *
 999 PyMarshal_ReadObjectFromFile(FILE *fp)
1000 {
1001         RFILE rf;
1002         PyObject *result;
1003         rf.fp = fp;
1004         rf.strings = PyList_New(0);
1005         result = r_object(&rf);
1006         Py_DECREF(rf.strings);
1007         return result;
1008 }
1009
1010 PyObject *
1011 PyMarshal_ReadObjectFromString(char *str, Py_ssize_t len)
1012 {
1013         RFILE rf;
1014         PyObject *result;
1015         rf.fp = NULL;
1016         rf.ptr = str;
1017         rf.end = str + len;
1018         rf.strings = PyList_New(0);
1019         result = r_object(&rf);
1020         Py_DECREF(rf.strings);
1021         return result;
1022 }
1023
1024 PyObject *
1025 PyMarshal_WriteObjectToString(PyObject *x, int version)
1026 {
1027         WFILE wf;
1028         wf.fp = NULL;
1029         wf.str = PyString_FromStringAndSize((char *)NULL, 50);
1030         if (wf.str == NULL)
1031                 return NULL;
1032         wf.ptr = PyString_AS_STRING((PyStringObject *)wf.str);
1033         wf.end = wf.ptr + PyString_Size(wf.str);
1034         wf.error = 0;
1035         wf.depth = 0;
1036         wf.version = version;
1037         wf.strings = (version > 0) ? PyDict_New() : NULL;
1038         w_object(x, &wf);
1039         Py_XDECREF(wf.strings);
1040         if (wf.str != NULL) {
1041                 char *base = PyString_AS_STRING((PyStringObject *)wf.str);
1042                 if (wf.ptr - base > PY_SSIZE_T_MAX) {
1043                         Py_DECREF(wf.str);
1044                         PyErr_SetString(PyExc_OverflowError,
1045                                         "too much marshall data for a string");
1046                         return NULL;
1047                 }
1048                 _PyString_Resize(&wf.str, (Py_ssize_t)(wf.ptr - base));
1049         }
1050         if (wf.error) {
1051                 Py_XDECREF(wf.str);
1052                 PyErr_SetString(PyExc_ValueError,
1053                                 (wf.error==1)?"unmarshallable object"
1054                                 :"object too deeply nested to marshal");
1055                 return NULL;
1056         }
1057         return wf.str;
1058 }
1059
1060 /* And an interface for Python programs... */
1061
1062 static PyObject *
1063 marshal_dump(PyObject *self, PyObject *args)
1064 {
1065         WFILE wf;
1066         PyObject *x;
1067         PyObject *f;
1068         int version = Py_MARSHAL_VERSION;
1069         if (!PyArg_ParseTuple(args, "OO|i:dump", &x, &f, &version))
1070                 return NULL;
1071         if (!PyFile_Check(f)) {
1072                 PyErr_SetString(PyExc_TypeError,
1073                                 "marshal.dump() 2nd arg must be file");
1074                 return NULL;
1075         }
1076         wf.fp = PyFile_AsFile(f);
1077         wf.str = NULL;
1078         wf.ptr = wf.end = NULL;
1079         wf.error = 0;
1080         wf.depth = 0;
1081         wf.strings = (version > 0) ? PyDict_New() : 0;
1082         wf.version = version;
1083         w_object(x, &wf);
1084         Py_XDECREF(wf.strings);
1085         if (wf.error) {
1086                 PyErr_SetString(PyExc_ValueError,
1087                                 (wf.error==1)?"unmarshallable object"
1088                                 :"object too deeply nested to marshal");
1089                 return NULL;
1090         }
1091         Py_INCREF(Py_None);
1092         return Py_None;
1093 }
1094
1095 static PyObject *
1096 marshal_load(PyObject *self, PyObject *f)
1097 {
1098         RFILE rf;
1099         PyObject *result;
1100         if (!PyFile_Check(f)) {
1101                 PyErr_SetString(PyExc_TypeError,
1102                                 "marshal.load() arg must be file");
1103                 return NULL;
1104         }
1105         rf.fp = PyFile_AsFile(f);
1106         rf.strings = PyList_New(0);
1107         result = read_object(&rf);
1108         Py_DECREF(rf.strings);
1109         return result;
1110 }
1111
1112 static PyObject *
1113 marshal_dumps(PyObject *self, PyObject *args)
1114 {
1115         PyObject *x;
1116         int version = Py_MARSHAL_VERSION;
1117         if (!PyArg_ParseTuple(args, "O|i:dumps", &x, &version))
1118                 return NULL;
1119         return PyMarshal_WriteObjectToString(x, version);
1120 }
1121
1122 static PyObject *
1123 marshal_loads(PyObject *self, PyObject *args)
1124 {
1125         RFILE rf;
1126         char *s;
1127         Py_ssize_t n;
1128         PyObject* result;
1129         if (!PyArg_ParseTuple(args, "s#:loads", &s, &n))
1130                 return NULL;
1131         rf.fp = NULL;
1132         rf.ptr = s;
1133         rf.end = s + n;
1134         rf.strings = PyList_New(0);
1135         result = read_object(&rf);
1136         Py_DECREF(rf.strings);
1137         return result;
1138 }
1139
1140 static PyMethodDef marshal_methods[] = {
1141         {"dump",        marshal_dump,   METH_VARARGS},
1142         {"load",        marshal_load,   METH_O},
1143         {"dumps",       marshal_dumps,  METH_VARARGS},
1144         {"loads",       marshal_loads,  METH_VARARGS},
1145         {NULL,          NULL}           /* sentinel */
1146 };
1147
1148 PyMODINIT_FUNC
1149 PyMarshal_Init(void)
1150 {
1151         PyObject *mod = Py_InitModule("marshal", marshal_methods);
1152         if (mod == NULL)
1153                 return;
1154         PyModule_AddIntConstant(mod, "version", Py_MARSHAL_VERSION);
1155 }