Slightly reworked the final write loop.
[ruby-vorbistagger.git] / ext / vcedit.c
1 /*
2  * Copyright (C) 2000-2001 Michael Smith (msmith at xiph org)
3  *
4  * This library is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU Lesser General Public
6  * License as published by the Free Software Foundation, version 2.
7  *
8  * This library is distributed in the hope that it will be useful,
9  * but WITHOUT ANY WARRANTY; without even the implied warranty of
10  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
11  * Lesser General Public License for more details.
12  *
13  * You should have received a copy of the GNU Lesser General Public
14  * License along with this library; if not, write to the Free Software
15  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
16  * MA 02110-1301 USA
17  */
18
19 #include <stdio.h>
20 #include <stdbool.h>
21 #include <stdlib.h>
22 #include <string.h>
23 #include <errno.h>
24 #include <limits.h>
25 #include <unistd.h>
26 #include <ogg/ogg.h>
27 #include <vorbis/codec.h>
28 #include <assert.h>
29
30 #include "vcedit.h"
31
32 #define CHUNKSIZE 4096
33
34 struct vcedit_state_St {
35         int refcount;
36
37         ogg_sync_state oy;
38         ogg_stream_state os;
39
40         vorbis_comment vc;
41         vorbis_info vi;
42
43         FILE *in;
44         bool opened;
45         long serial;
46         unsigned char *mainbuf;
47         unsigned char *bookbuf;
48         int     mainlen;
49         int     booklen;
50         char *vendor;
51         int prevW;
52         int extrapage;
53         int eosin;
54
55         char filename[0];
56 };
57
58 static void
59 vcedit_state_free (vcedit_state *state)
60 {
61         free (state->mainbuf);
62         free (state->bookbuf);
63         free (state->vendor);
64
65         if (state->in) {
66                 fclose (state->in);
67                 state->in = NULL;
68         }
69
70         free (state);
71 }
72
73 static bool
74 vcedit_state_init (vcedit_state *state, const char *filename)
75 {
76         state->refcount = 1;
77
78         strcpy (state->filename, filename);
79
80         return true;
81 }
82
83 vcedit_state *
84 vcedit_state_new (const char *filename)
85 {
86         vcedit_state *state;
87         size_t len;
88
89         len = strlen (filename);
90         if (len > PATH_MAX)
91                 return NULL;
92
93         state = malloc (sizeof (vcedit_state) + len + 1);
94         if (!state)
95                 return NULL;
96
97         memset (state, 0, sizeof (vcedit_state));
98
99         if (!vcedit_state_init (state, filename)) {
100                 vcedit_state_free (state);
101                 return NULL;
102         }
103
104         return state;
105 }
106
107 vorbis_comment *
108 vcedit_comments (vcedit_state *state)
109 {
110         return state->opened ? &state->vc : NULL;
111 }
112
113 static void
114 vcedit_clear_internals (vcedit_state *state)
115 {
116         ogg_stream_clear (&state->os);
117         ogg_sync_clear (&state->oy);
118
119         vorbis_info_clear (&state->vi);
120         vorbis_comment_clear (&state->vc);
121
122         free (state->vendor);
123         state->vendor = NULL;
124
125         free (state->mainbuf);
126         state->mainbuf = NULL;
127         state->mainlen = 0;
128
129         free (state->bookbuf);
130         state->bookbuf = NULL;
131         state->booklen = 0;
132
133         state->serial = 0;
134         state->opened = false;
135 }
136
137 void
138 vcedit_state_ref (vcedit_state *state)
139 {
140         state->refcount++;
141 }
142
143 void
144 vcedit_state_unref (vcedit_state *state)
145 {
146         if (--state->refcount)
147                 return;
148
149         if (state->opened)
150                 vcedit_clear_internals (state);
151
152         vcedit_state_free (state);
153 }
154
155 /* Next two functions pulled straight from libvorbis, apart from one change
156  * - we don't want to overwrite the vendor string.
157  */
158 static void
159 _v_writestring (oggpack_buffer *o, char *s, int len)
160 {
161         while (len--)
162                 oggpack_write (o, *s++, 8);
163 }
164
165 static int
166 _commentheader_out (vorbis_comment *vc, char *vendor, ogg_packet *op)
167 {
168         int i;
169
170         oggpack_buffer opb;
171
172         oggpack_writeinit (&opb);
173
174         /* preamble */
175         oggpack_write (&opb, 0x03, 8);
176         _v_writestring (&opb, "vorbis", 6);
177
178         /* vendor */
179         oggpack_write (&opb, strlen (vendor), 32);
180         _v_writestring (&opb, vendor, strlen (vendor));
181
182         /* comments */
183         oggpack_write (&opb, vc->comments, 32);
184
185         for (i = 0; i < vc->comments; i++) {
186                 if (!vc->user_comments[i])
187                         oggpack_write (&opb, 0, 32);
188                 else {
189                         oggpack_write (&opb, vc->comment_lengths[i], 32);
190                         _v_writestring (&opb, vc->user_comments[i],
191                                         vc->comment_lengths[i]);
192                 }
193         }
194
195         oggpack_write (&opb, 1, 1);
196
197         op->packet = _ogg_malloc (oggpack_bytes (&opb));
198         memcpy (op->packet, opb.buffer, oggpack_bytes (&opb));
199
200         op->bytes = oggpack_bytes (&opb);
201         op->b_o_s = 0;
202         op->e_o_s = 0;
203         op->granulepos = 0;
204
205         oggpack_writeclear (&opb);
206
207         return 0;
208 }
209
210 static int
211 _blocksize (vcedit_state *s, ogg_packet *p)
212 {
213         int this, ret = 0;
214
215         this = vorbis_packet_blocksize (&s->vi, p);
216
217         if (s->prevW)
218                 ret = (this + s->prevW) / 4;
219
220         s->prevW = this;
221
222         return ret;
223 }
224
225 static int
226 _fetch_next_packet (vcedit_state *s, ogg_packet *p, ogg_page *page)
227 {
228         char *buffer;
229         int result, bytes;
230
231         result = ogg_stream_packetout (&s->os, p);
232         if (result == 1)
233                 return 1;
234
235         if (s->eosin)
236                 return 0;
237
238         while (ogg_sync_pageout (&s->oy, page) != 1) {
239                 buffer = ogg_sync_buffer (&s->oy, CHUNKSIZE);
240                 bytes = fread (buffer, 1, CHUNKSIZE, s->in);
241                 ogg_sync_wrote (&s->oy, bytes);
242
243                 if (!bytes && feof (s->in))
244                         return 0;
245         }
246
247         if (ogg_page_eos (page))
248                 s->eosin = 1;
249         else if (ogg_page_serialno (page) != s->serial) {
250                 s->eosin = 1;
251                 s->extrapage = 1;
252                 return 0;
253         }
254
255         ogg_stream_pagein (&s->os, page);
256
257         return _fetch_next_packet (s, p, page);
258 }
259
260 vcedit_error
261 vcedit_open (vcedit_state *state)
262 {
263         vcedit_error ret;
264         ogg_packet *header;
265         ogg_packet header_main, header_comments, header_codebooks;
266         ogg_page og;
267         char *buffer;
268         size_t bytes, total = 0;
269         int i = 0;
270
271         state->in = fopen (state->filename, "rb");
272         if (!state->in)
273                 return VCEDIT_ERR_OPEN;
274
275         ogg_sync_init (&state->oy);
276
277         do {
278                 /* Bail if we don't find data in the first 40 kB */
279                 if (feof (state->in) || total >= (CHUNKSIZE * 10)) {
280                         ogg_sync_clear (&state->oy);
281
282                         return VCEDIT_ERR_INVAL;
283                 }
284
285                 buffer = ogg_sync_buffer (&state->oy, CHUNKSIZE);
286
287                 bytes = fread (buffer, 1, CHUNKSIZE, state->in);
288                 total += bytes;
289
290                 ogg_sync_wrote (&state->oy, bytes);
291         } while (ogg_sync_pageout (&state->oy, &og) != 1);
292
293         state->serial = ogg_page_serialno (&og);
294
295         ogg_stream_init (&state->os, state->serial);
296         vorbis_info_init (&state->vi);
297         vorbis_comment_init (&state->vc);
298
299         if (ogg_stream_pagein (&state->os, &og) < 0) {
300                 ret = VCEDIT_ERR_INVAL;
301                 goto err;
302         }
303
304         if (ogg_stream_packetout (&state->os, &header_main) != 1) {
305                 ret = VCEDIT_ERR_INVAL;
306                 goto err;
307         }
308
309         if (vorbis_synthesis_headerin (&state->vi, &state->vc, &header_main) < 0) {
310                 ret = VCEDIT_ERR_INVAL;
311                 goto err;
312         }
313
314         state->mainlen = header_main.bytes;
315         state->mainbuf = malloc (state->mainlen);
316         memcpy (state->mainbuf, header_main.packet, header_main.bytes);
317
318         header = &header_comments;
319
320         while (i < 2) {
321                 while (i < 2) {
322                         int result;
323
324                         result = ogg_sync_pageout (&state->oy, &og);
325                         if (!result)
326                                 break; /* Too little data so far */
327
328                         if (result != 1)
329                                 continue;
330
331                         ogg_stream_pagein (&state->os, &og);
332
333                         while (i < 2) {
334                                 result = ogg_stream_packetout (&state->os, header);
335                                 if (!result)
336                                         break;
337
338                                 if (result != 1) {
339                                         ret = VCEDIT_ERR_INVAL;
340                                         goto err;
341                                 }
342
343                                 vorbis_synthesis_headerin (&state->vi, &state->vc, header);
344
345                                 if (i++ == 1) {
346                                         state->booklen = header->bytes;
347                                         state->bookbuf = malloc (state->booklen);
348                                         memcpy (state->bookbuf, header->packet, header->bytes);
349                                 }
350
351                                 header = &header_codebooks;
352                         }
353                 }
354
355                 buffer = ogg_sync_buffer (&state->oy, CHUNKSIZE);
356                 bytes = fread (buffer, 1, CHUNKSIZE, state->in);
357
358                 if (!bytes && feof (state->in) && i < 2) {
359                         ret = VCEDIT_ERR_INVAL;
360                         goto err;
361                 }
362
363                 ogg_sync_wrote (&state->oy, bytes);
364         }
365
366         /* Copy the vendor tag */
367         state->vendor = strdup (state->vc.vendor);
368
369         /* Headers are done! */
370         state->opened = true;
371
372         return VCEDIT_ERR_SUCCESS;
373
374 err:
375         vcedit_clear_internals (state);
376
377         return ret;
378 }
379
380 vcedit_error
381 vcedit_write (vcedit_state *state)
382 {
383         ogg_stream_state streamout;
384         ogg_packet header_main, header_comments, header_codebooks, op;
385         ogg_page ogout, ogin;
386         ogg_int64_t granpos = 0;
387         FILE *out;
388         char *buffer, tmpfile[PATH_MAX];
389         int s, result, bytes, needflush = 0, needout = 0;
390         size_t tmp;
391
392         if (!state->opened)
393                 return VCEDIT_ERR_INVAL;
394
395         strcpy (tmpfile, state->filename);
396         strcat (tmpfile, ".XXXXXX");
397
398         s = mkstemp (tmpfile);
399         if (s == -1)
400                 return VCEDIT_ERR_TMPFILE;
401
402         out = fdopen (s, "wb");
403         if (!out) {
404                 unlink (tmpfile);
405                 close (s);
406
407                 return VCEDIT_ERR_TMPFILE;
408         }
409
410         state->prevW = state->extrapage = state->eosin = 0;
411
412         header_main.bytes = state->mainlen;
413         header_main.packet = state->mainbuf;
414         header_main.b_o_s = 1;
415         header_main.e_o_s = 0;
416         header_main.granulepos = 0;
417
418         header_codebooks.bytes = state->booklen;
419         header_codebooks.packet = state->bookbuf;
420         header_codebooks.b_o_s = 0;
421         header_codebooks.e_o_s = 0;
422         header_codebooks.granulepos = 0;
423
424         ogg_stream_init (&streamout, state->serial);
425
426         _commentheader_out (&state->vc, state->vendor, &header_comments);
427
428         ogg_stream_packetin (&streamout, &header_main);
429         ogg_stream_packetin (&streamout, &header_comments);
430         ogg_stream_packetin (&streamout, &header_codebooks);
431
432         while ((result = ogg_stream_flush (&streamout, &ogout))) {
433                 tmp = fwrite (ogout.header, 1, ogout.header_len, out);
434                 if (tmp != (size_t) ogout.header_len)
435                         goto cleanup;
436
437                 tmp = fwrite (ogout.body, 1, ogout.body_len, out);
438                 if (tmp != (size_t) ogout.body_len)
439                         goto cleanup;
440         }
441
442         while (_fetch_next_packet (state, &op, &ogin)) {
443                 int size;
444
445                 size = _blocksize (state, &op);
446                 granpos += size;
447
448                 if (needflush) {
449                         if (ogg_stream_flush (&streamout, &ogout)) {
450                                 tmp = fwrite (ogout.header, 1, ogout.header_len, out);
451                                 if (tmp != (size_t) ogout.header_len)
452                                         goto cleanup;
453
454                                 tmp = fwrite (ogout.body, 1, ogout.body_len, out);
455                                 if (tmp != (size_t) ogout.body_len)
456                                         goto cleanup;
457                         }
458                 } else if (needout) {
459                         if (ogg_stream_pageout (&streamout, &ogout)) {
460                                 tmp = fwrite (ogout.header, 1, ogout.header_len, out);
461                                 if (tmp != (size_t) ogout.header_len)
462                                         goto cleanup;
463
464                                 tmp = fwrite (ogout.body, 1, ogout.body_len, out);
465                                 if (tmp != (size_t) ogout.body_len)
466                                         goto cleanup;
467                         }
468                 }
469
470                 needflush = needout = 0;
471
472                 if (op.granulepos == -1) {
473                         op.granulepos = granpos;
474                         ogg_stream_packetin (&streamout, &op);
475                 } else {
476                         /* granulepos is set, validly. Use it, and force a flush to
477                          * account for shortened blocks (vcut) when appropriate
478                          */
479                         if (granpos > op.granulepos) {
480                                 granpos = op.granulepos;
481                                 ogg_stream_packetin (&streamout, &op);
482                                 needflush = 1;
483                         } else {
484                                 ogg_stream_packetin (&streamout, &op);
485                                 needout = 1;
486                         }
487                 }
488         }
489
490         streamout.e_o_s = 1;
491
492         while (ogg_stream_flush (&streamout, &ogout)) {
493                 tmp = fwrite (ogout.header, 1, ogout.header_len, out);
494                 if (tmp != (size_t) ogout.header_len)
495                         goto cleanup;
496
497                 tmp = fwrite (ogout.body, 1, ogout.body_len, out);
498                 if (tmp != (size_t) ogout.body_len)
499                         goto cleanup;
500         }
501
502         if (state->extrapage) {
503                 tmp = fwrite (ogin.header, 1, ogin.header_len, out);
504                 if (tmp != (size_t) ogin.header_len)
505                         goto cleanup;
506
507                 tmp = fwrite (ogin.body, 1, ogin.body_len, out);
508                 if (tmp != (size_t) ogin.body_len)
509                         goto cleanup;
510         }
511
512         /* clear it, because not all paths to here do */
513         state->eosin = 0;
514
515         do {
516                 /* We copy the rest of the stream (other logical streams)
517                  * through, a page at a time.
518                  */
519                 while ((result = ogg_sync_pageout (&state->oy, &ogout))) {
520                         if (result != 1)
521                                 continue;
522
523                         /* Don't bother going through the rest, we can just
524                          * write the page out now
525                          */
526                         tmp = fwrite (ogout.header, 1, ogout.header_len, out);
527                         if (tmp != (size_t) ogout.header_len)
528                                 goto cleanup;
529
530                         tmp = fwrite (ogout.body, 1, ogout.body_len, out);
531                         if (tmp != (size_t) ogout.body_len)
532                                 goto cleanup;
533                 }
534
535                 buffer = ogg_sync_buffer (&state->oy, CHUNKSIZE);
536                 bytes = fread (buffer, 1, CHUNKSIZE, state->in);
537                 ogg_sync_wrote (&state->oy, bytes);
538
539                 state->eosin = !bytes && feof (state->in);
540         } while (!state->eosin);
541
542         fclose (out);
543         fclose (state->in);
544
545         unlink (state->filename);
546         rename (tmpfile, state->filename);
547
548 cleanup:
549         ogg_stream_clear (&streamout);
550
551     /* We don't ogg_packet_clear() this, because the memory was
552          * allocated in _commentheader_out(), so we mirror that here
553          */
554     _ogg_free (header_comments.packet);
555
556         free (state->mainbuf);
557         free (state->bookbuf);
558
559     state->mainbuf = state->bookbuf = NULL;
560
561         if (!state->eosin)
562                 return VCEDIT_ERR_INVAL;
563
564         vcedit_clear_internals (state);
565
566         return (vcedit_open (state) == VCEDIT_ERR_SUCCESS) ?
567                VCEDIT_ERR_SUCCESS : VCEDIT_ERR_REOPEN;
568 }