vfs: check userland buffers before reading them.
[haiku.git] / src / add-ons / media / plugins / ffmpeg / AVFormatWriter.cpp
blob017fde4b0111fa7ca64cd41acd42b1b5ca4cd3c9
1 /*
2 * Copyright 2009-2010, Stephan Aßmus <superstippi@gmx.de>
3 * All rights reserved. Distributed under the terms of the GNU L-GPL license.
4 */
6 #include "AVFormatWriter.h"
8 #include <stdio.h>
9 #include <string.h>
10 #include <stdlib.h>
12 #include <new>
14 #include <Application.h>
15 #include <AutoDeleter.h>
16 #include <Autolock.h>
17 #include <ByteOrder.h>
18 #include <MediaIO.h>
19 #include <MediaDefs.h>
20 #include <MediaFormats.h>
21 #include <Roster.h>
23 extern "C" {
24 #include "avformat.h"
27 #include "DemuxerTable.h"
28 #include "EncoderTable.h"
29 #include "gfx_util.h"
32 //#define TRACE_AVFORMAT_WRITER
33 #ifdef TRACE_AVFORMAT_WRITER
34 # define TRACE printf
35 # define TRACE_IO(a...)
36 # define TRACE_PACKET printf
37 #else
38 # define TRACE(a...)
39 # define TRACE_IO(a...)
40 # define TRACE_PACKET(a...)
41 #endif
43 #define ERROR(a...) fprintf(stderr, a)
45 #if LIBAVCODEC_VERSION_INT < ((54 << 16) | (50 << 8))
46 #define AV_CODEC_ID_NONE CODEC_ID_NONE
47 #define AV_PIX_FMT_YUV420P PIX_FMT_YUV420P
48 #endif
51 static const size_t kIOBufferSize = 64 * 1024;
52 // TODO: This could depend on the BMediaFile creation flags, IIRC,
53 // they allow to specify a buffering mode.
55 // NOTE: The following works around some weird bug in libavformat. We
56 // have to open the AVFormatContext->AVStream->AVCodecContext, even though
57 // we are not interested in donig any encoding here!!
58 #define OPEN_CODEC_CONTEXT 1
59 #define GET_CONTEXT_DEFAULTS 0
61 #if LIBAVCODEC_VERSION_INT > ((54 << 16) | (50 << 8))
62 typedef AVCodecID CodecID;
63 #endif
65 // #pragma mark - AVFormatWriter::StreamCookie
68 class AVFormatWriter::StreamCookie {
69 public:
70 StreamCookie(AVFormatContext* context,
71 BLocker* streamLock);
72 virtual ~StreamCookie();
74 status_t Init(media_format* format,
75 const media_codec_info* codecInfo);
77 status_t WriteChunk(const void* chunkBuffer,
78 size_t chunkSize,
79 media_encode_info* encodeInfo);
81 status_t AddTrackInfo(uint32 code, const void* data,
82 size_t size, uint32 flags);
84 private:
85 AVFormatContext* fContext;
86 AVStream* fStream;
87 AVPacket fPacket;
88 // Since different threads may write to the target,
89 // we need to protect the file position and I/O by a lock.
90 BLocker* fStreamLock;
95 AVFormatWriter::StreamCookie::StreamCookie(AVFormatContext* context,
96 BLocker* streamLock)
98 fContext(context),
99 fStream(NULL),
100 fStreamLock(streamLock)
102 av_init_packet(&fPacket);
106 AVFormatWriter::StreamCookie::~StreamCookie()
111 status_t
112 AVFormatWriter::StreamCookie::Init(media_format* format,
113 const media_codec_info* codecInfo)
115 TRACE("AVFormatWriter::StreamCookie::Init()\n");
117 BAutolock _(fStreamLock);
119 fPacket.stream_index = fContext->nb_streams;
120 fStream = avformat_new_stream(fContext, NULL);
121 fStream->id = fPacket.stream_index;
123 if (fStream == NULL) {
124 TRACE(" failed to add new stream\n");
125 return B_ERROR;
128 // TRACE(" fStream->codec: %p\n", fStream->codec);
129 // TODO: This is a hack for now! Use avcodec_find_encoder_by_name()
130 // or something similar...
131 fStream->codec->codec_id = (CodecID)codecInfo->sub_id;
132 if (fStream->codec->codec_id == AV_CODEC_ID_NONE)
133 fStream->codec->codec_id = raw_audio_codec_id_for(*format);
135 // Setup the stream according to the media format...
136 if (format->type == B_MEDIA_RAW_VIDEO) {
137 fStream->codec->codec_type = AVMEDIA_TYPE_VIDEO;
138 #if GET_CONTEXT_DEFAULTS
139 // NOTE: API example does not do this:
140 avcodec_get_context_defaults(fStream->codec);
141 #endif
142 // frame rate
143 fStream->codec->time_base.den = (int)format->u.raw_video.field_rate;
144 fStream->codec->time_base.num = 1;
145 // video size
146 fStream->codec->width = format->u.raw_video.display.line_width;
147 fStream->codec->height = format->u.raw_video.display.line_count;
148 // pixel aspect ratio
149 fStream->sample_aspect_ratio.num
150 = format->u.raw_video.pixel_width_aspect;
151 fStream->sample_aspect_ratio.den
152 = format->u.raw_video.pixel_height_aspect;
153 if (fStream->sample_aspect_ratio.num == 0
154 || fStream->sample_aspect_ratio.den == 0) {
155 av_reduce(&fStream->sample_aspect_ratio.num,
156 &fStream->sample_aspect_ratio.den, fStream->codec->width,
157 fStream->codec->height, 255);
160 fStream->codec->gop_size = 12;
162 fStream->codec->sample_aspect_ratio = fStream->sample_aspect_ratio;
164 // Use the last supported pixel format of the AVCodec, which we hope
165 // is the one with the best quality (true for all currently supported
166 // encoders).
167 // AVCodec* codec = fStream->codec->codec;
168 // for (int i = 0; codec->pix_fmts[i] != PIX_FMT_NONE; i++)
169 // fStream->codec->pix_fmt = codec->pix_fmts[i];
170 fStream->codec->pix_fmt = AV_PIX_FMT_YUV420P;
172 } else if (format->type == B_MEDIA_RAW_AUDIO) {
173 fStream->codec->codec_type = AVMEDIA_TYPE_AUDIO;
174 #if GET_CONTEXT_DEFAULTS
175 // NOTE: API example does not do this:
176 avcodec_get_context_defaults(fStream->codec);
177 #endif
178 // frame rate
179 fStream->codec->sample_rate = (int)format->u.raw_audio.frame_rate;
181 // channels
182 fStream->codec->channels = format->u.raw_audio.channel_count;
184 // set fStream to the audio format we want to use. This is only a hint
185 // (each encoder has a different set of accepted formats)
186 switch (format->u.raw_audio.format) {
187 case media_raw_audio_format::B_AUDIO_FLOAT:
188 fStream->codec->sample_fmt = AV_SAMPLE_FMT_FLT;
189 break;
190 case media_raw_audio_format::B_AUDIO_DOUBLE:
191 fStream->codec->sample_fmt = AV_SAMPLE_FMT_DBL;
192 break;
193 case media_raw_audio_format::B_AUDIO_INT:
194 fStream->codec->sample_fmt = AV_SAMPLE_FMT_S32;
195 break;
196 case media_raw_audio_format::B_AUDIO_SHORT:
197 fStream->codec->sample_fmt = AV_SAMPLE_FMT_S16;
198 break;
199 case media_raw_audio_format::B_AUDIO_UCHAR:
200 fStream->codec->sample_fmt = AV_SAMPLE_FMT_U8;
201 break;
203 case media_raw_audio_format::B_AUDIO_CHAR:
204 default:
205 return B_MEDIA_BAD_FORMAT;
206 break;
209 // Now negociate the actual format with the encoder
210 // First check if the requested format is acceptable
211 AVCodec* codec = avcodec_find_encoder(fStream->codec->codec_id);
213 if (codec == NULL)
214 return B_MEDIA_BAD_FORMAT;
216 const enum AVSampleFormat *p = codec->sample_fmts;
217 for (; *p != -1; p++) {
218 if (*p == fStream->codec->sample_fmt)
219 break;
221 // If not, force one of the acceptable ones
222 if (*p == -1) {
223 fStream->codec->sample_fmt = codec->sample_fmts[0];
225 // And finally set the format struct to the accepted format. It is
226 // then up to the caller to make sure we get data matching that
227 // format.
228 switch (fStream->codec->sample_fmt) {
229 case AV_SAMPLE_FMT_FLT:
230 format->u.raw_audio.format
231 = media_raw_audio_format::B_AUDIO_FLOAT;
232 break;
233 case AV_SAMPLE_FMT_DBL:
234 format->u.raw_audio.format
235 = media_raw_audio_format::B_AUDIO_DOUBLE;
236 break;
237 case AV_SAMPLE_FMT_S32:
238 format->u.raw_audio.format
239 = media_raw_audio_format::B_AUDIO_INT;
240 break;
241 case AV_SAMPLE_FMT_S16:
242 format->u.raw_audio.format
243 = media_raw_audio_format::B_AUDIO_SHORT;
244 break;
245 case AV_SAMPLE_FMT_U8:
246 format->u.raw_audio.format
247 = media_raw_audio_format::B_AUDIO_UCHAR;
248 break;
249 default:
250 return B_MEDIA_BAD_FORMAT;
251 break;
255 if (format->u.raw_audio.channel_mask == 0) {
256 // guess the channel mask...
257 switch (format->u.raw_audio.channel_count) {
258 default:
259 case 2:
260 fStream->codec->channel_layout = AV_CH_LAYOUT_STEREO;
261 break;
262 case 1:
263 fStream->codec->channel_layout = AV_CH_LAYOUT_MONO;
264 break;
265 case 3:
266 fStream->codec->channel_layout = AV_CH_LAYOUT_SURROUND;
267 break;
268 case 4:
269 fStream->codec->channel_layout = AV_CH_LAYOUT_QUAD;
270 break;
271 case 5:
272 fStream->codec->channel_layout = AV_CH_LAYOUT_5POINT0;
273 break;
274 case 6:
275 fStream->codec->channel_layout = AV_CH_LAYOUT_5POINT1;
276 break;
277 case 8:
278 fStream->codec->channel_layout = AV_CH_LAYOUT_7POINT1;
279 break;
280 case 10:
281 fStream->codec->channel_layout = AV_CH_LAYOUT_7POINT1_WIDE;
282 break;
284 } else {
285 // The bits match 1:1 for media_multi_channels and FFmpeg defines.
286 fStream->codec->channel_layout = format->u.raw_audio.channel_mask;
290 // Some formats want stream headers to be separate
291 if ((fContext->oformat->flags & AVFMT_GLOBALHEADER) != 0)
292 fStream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
294 TRACE(" stream->time_base: (%d/%d), codec->time_base: (%d/%d))\n",
295 fStream->time_base.num, fStream->time_base.den,
296 fStream->codec->time_base.num, fStream->codec->time_base.den);
298 #if 0
299 // Write the AVCodecContext pointer to the user data section of the
300 // media_format. For some encoders, it seems to be necessary to use
301 // the AVCodecContext of the AVStream in order to successfully encode
302 // anything and write valid media files. For example some codecs need
303 // to store meta data or global data in the container.
304 app_info appInfo;
305 if (be_app->GetAppInfo(&appInfo) == B_OK) {
306 uchar* userData = format->user_data;
307 *(uint32*)userData = 'ffmp';
308 userData += sizeof(uint32);
309 *(team_id*)userData = appInfo.team;
310 userData += sizeof(team_id);
311 *(AVCodecContext**)userData = fStream->codec;
313 #endif
315 return B_OK;
319 status_t
320 AVFormatWriter::StreamCookie::WriteChunk(const void* chunkBuffer,
321 size_t chunkSize, media_encode_info* encodeInfo)
323 TRACE_PACKET("AVFormatWriter::StreamCookie[%d]::WriteChunk(%p, %ld, "
324 "start_time: %lld)\n", fStream->index, chunkBuffer, chunkSize,
325 encodeInfo->start_time);
327 BAutolock _(fStreamLock);
329 fPacket.data = const_cast<uint8_t*>((const uint8_t*)chunkBuffer);
330 fPacket.size = chunkSize;
332 fPacket.pts = int64_t((double)encodeInfo->start_time
333 * fStream->time_base.den / (1000000.0 * fStream->time_base.num)
334 + 0.5);
336 fPacket.flags = 0;
337 if ((encodeInfo->flags & B_MEDIA_KEY_FRAME) != 0)
338 fPacket.flags |= AV_PKT_FLAG_KEY;
340 TRACE_PACKET(" PTS: %lld (stream->time_base: (%d/%d), "
341 "codec->time_base: (%d/%d))\n", fPacket.pts,
342 fStream->time_base.num, fStream->time_base.den,
343 fStream->codec->time_base.num, fStream->codec->time_base.den);
345 #if 0
346 // TODO: Eventually, we need to write interleaved packets, but
347 // maybe we are only supposed to use this if we have actually
348 // more than one stream. For the moment, this crashes in AVPacket
349 // shuffling inside libavformat. Maybe if we want to use this, we
350 // need to allocate a separate AVPacket and copy the chunk buffer.
351 int result = av_interleaved_write_frame(fContext, &fPacket);
352 if (result < 0)
353 TRACE(" av_interleaved_write_frame(): %d\n", result);
354 #else
355 int result = av_write_frame(fContext, &fPacket);
356 if (result < 0)
357 TRACE(" av_write_frame(): %d\n", result);
358 #endif
360 return result == 0 ? B_OK : B_ERROR;
364 status_t
365 AVFormatWriter::StreamCookie::AddTrackInfo(uint32 code,
366 const void* data, size_t size, uint32 flags)
368 TRACE("AVFormatWriter::StreamCookie::AddTrackInfo(%lu, %p, %ld, %lu)\n",
369 code, data, size, flags);
371 BAutolock _(fStreamLock);
373 return B_NOT_SUPPORTED;
377 // #pragma mark - AVFormatWriter
380 AVFormatWriter::AVFormatWriter()
382 fContext(avformat_alloc_context()),
383 fCodecOpened(false),
384 fHeaderError(-1),
385 fIOContext(NULL),
386 fStreamLock("stream lock")
388 TRACE("AVFormatWriter::AVFormatWriter\n");
392 AVFormatWriter::~AVFormatWriter()
394 TRACE("AVFormatWriter::~AVFormatWriter\n");
396 // Free the streams and close the AVCodecContexts
397 for(unsigned i = 0; i < fContext->nb_streams; i++) {
398 #if OPEN_CODEC_CONTEXT
399 // We only need to close the AVCodecContext when we opened it.
400 // This is experimental, see CommitHeader().
401 if (fCodecOpened)
402 avcodec_close(fContext->streams[i]->codec);
403 #endif
404 av_freep(&fContext->streams[i]->codec);
405 av_freep(&fContext->streams[i]);
408 av_free(fContext);
409 av_free(fIOContext->buffer);
410 av_free(fIOContext);
414 // #pragma mark -
417 status_t
418 AVFormatWriter::Init(const media_file_format* fileFormat)
420 TRACE("AVFormatWriter::Init()\n");
422 uint8* buffer = static_cast<uint8*>(av_malloc(kIOBufferSize));
423 if (buffer == NULL)
424 return B_NO_MEMORY;
426 // Allocate I/O context and initialize it with buffer
427 // and hook functions, pass ourself as cookie.
428 fIOContext = avio_alloc_context(buffer, kIOBufferSize, 1, this,
429 0, _Write, _Seek);
430 if (fIOContext == NULL) {
431 TRACE("av_alloc_put_byte() failed!\n");
432 return B_ERROR;
435 // Setup I/O hooks. This seems to be enough.
436 fContext->pb = fIOContext;
438 // Set the AVOutputFormat according to fileFormat...
439 fContext->oformat = av_guess_format(fileFormat->short_name,
440 fileFormat->file_extension, fileFormat->mime_type);
441 if (fContext->oformat == NULL) {
442 TRACE(" failed to find AVOuputFormat for %s\n",
443 fileFormat->short_name);
444 return B_NOT_SUPPORTED;
447 TRACE(" found AVOuputFormat for %s: %s\n", fileFormat->short_name,
448 fContext->oformat->name);
450 return B_OK;
454 status_t
455 AVFormatWriter::SetCopyright(const char* copyright)
457 TRACE("AVFormatWriter::SetCopyright(%s)\n", copyright);
459 return B_NOT_SUPPORTED;
463 status_t
464 AVFormatWriter::CommitHeader()
466 TRACE("AVFormatWriter::CommitHeader\n");
468 if (fContext == NULL)
469 return B_NO_INIT;
471 if (fCodecOpened)
472 return B_NOT_ALLOWED;
474 #if OPEN_CODEC_CONTEXT
475 for (unsigned i = 0; i < fContext->nb_streams; i++) {
476 AVStream* stream = fContext->streams[i];
477 // NOTE: Experimental, this should not be needed. Especially, since
478 // we have no idea (in the future) what CodecID some encoder uses,
479 // it may be an encoder from a different plugin.
480 AVCodecContext* codecContext = stream->codec;
481 codecContext->strict_std_compliance = -2;
482 AVCodec* codec = avcodec_find_encoder(codecContext->codec_id);
483 if (codec == NULL || avcodec_open2(codecContext, codec, NULL) < 0) {
484 TRACE(" stream[%u] - failed to open AVCodecContext\n", i);
486 TRACE(" stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n",
487 i, stream->time_base.num, stream->time_base.den,
488 stream->codec->time_base.num, stream->codec->time_base.den);
490 #endif
492 // We need to close the codecs we opened, even in case of failure.
493 fCodecOpened = true;
495 fHeaderError = avformat_write_header(fContext, NULL);
496 if (fHeaderError < 0)
497 TRACE(" avformat_write_header(): %d\n", fHeaderError);
499 #ifdef TRACE_AVFORMAT_WRITER
500 TRACE(" wrote header\n");
501 for (unsigned i = 0; i < fContext->nb_streams; i++) {
502 AVStream* stream = fContext->streams[i];
503 TRACE(" stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n",
504 i, stream->time_base.num, stream->time_base.den,
505 stream->codec->time_base.num, stream->codec->time_base.den);
507 #endif // TRACE_AVFORMAT_WRITER
509 return fHeaderError == 0 ? B_OK : B_ERROR;
513 status_t
514 AVFormatWriter::Flush()
516 TRACE("AVFormatWriter::Flush\n");
518 return B_NOT_SUPPORTED;
522 status_t
523 AVFormatWriter::Close()
525 TRACE("AVFormatWriter::Close\n");
527 if (fContext == NULL)
528 return B_NO_INIT;
530 if (!fCodecOpened)
531 return B_NOT_ALLOWED;
533 // From ffmpeg documentation: [av_write_trailer] may only be called
534 // after a successful call to avformat_write_header.
535 if (fHeaderError != 0)
536 return B_ERROR;
538 int result = av_write_trailer(fContext);
539 if (result < 0)
540 TRACE(" av_write_trailer(): %d\n", result);
541 return result == 0 ? B_OK : B_ERROR;
545 status_t
546 AVFormatWriter::AllocateCookie(void** _cookie, media_format* format,
547 const media_codec_info* codecInfo)
549 TRACE("AVFormatWriter::AllocateCookie()\n");
551 if (fCodecOpened)
552 return B_NOT_ALLOWED;
554 BAutolock _(fStreamLock);
556 if (_cookie == NULL)
557 return B_BAD_VALUE;
559 StreamCookie* cookie = new(std::nothrow) StreamCookie(fContext,
560 &fStreamLock);
562 status_t ret = cookie->Init(format, codecInfo);
563 if (ret != B_OK) {
564 delete cookie;
565 return ret;
568 *_cookie = cookie;
569 return B_OK;
573 status_t
574 AVFormatWriter::FreeCookie(void* _cookie)
576 BAutolock _(fStreamLock);
578 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
579 delete cookie;
581 return B_OK;
585 // #pragma mark -
588 status_t
589 AVFormatWriter::SetCopyright(void* cookie, const char* copyright)
591 TRACE("AVFormatWriter::SetCopyright(%p, %s)\n", cookie, copyright);
593 return B_NOT_SUPPORTED;
597 status_t
598 AVFormatWriter::AddTrackInfo(void* _cookie, uint32 code,
599 const void* data, size_t size, uint32 flags)
601 TRACE("AVFormatWriter::AddTrackInfo(%lu, %p, %ld, %lu)\n",
602 code, data, size, flags);
604 if (fHeaderError != 0)
605 return B_ERROR;
607 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
608 return cookie->AddTrackInfo(code, data, size, flags);
612 status_t
613 AVFormatWriter::WriteChunk(void* _cookie, const void* chunkBuffer,
614 size_t chunkSize, media_encode_info* encodeInfo)
616 TRACE_PACKET("AVFormatWriter::WriteChunk(%p, %ld, %p)\n", chunkBuffer,
617 chunkSize, encodeInfo);
619 if (fHeaderError != 0)
620 return B_ERROR;
622 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
623 return cookie->WriteChunk(chunkBuffer, chunkSize, encodeInfo);
627 // #pragma mark - I/O hooks
630 /*static*/ int
631 AVFormatWriter::_Write(void* cookie, uint8* buffer, int bufferSize)
633 TRACE_IO("AVFormatWriter::_Write(%p, %p, %d)\n",
634 cookie, buffer, bufferSize);
636 AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie);
638 ssize_t written = writer->fTarget->Write(buffer, bufferSize);
640 TRACE_IO(" written: %ld\n", written);
641 return (int)written;
646 /*static*/ off_t
647 AVFormatWriter::_Seek(void* cookie, off_t offset, int whence)
649 TRACE_IO("AVFormatWriter::_Seek(%p, %lld, %d)\n",
650 cookie, offset, whence);
652 AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie);
654 BMediaIO* mediaIO = dynamic_cast<BMediaIO*>(writer->fTarget);
655 if (mediaIO == NULL)
656 return -1;
658 // Support for special file size retrieval API without seeking anywhere:
659 if (whence == AVSEEK_SIZE) {
660 off_t size;
661 if (mediaIO->GetSize(&size) == B_OK)
662 return size;
664 return -1;
667 off_t position = mediaIO->Seek(offset, whence);
668 TRACE_IO(" position: %lld\n", position);
669 if (position < 0)
670 return -1;
672 return position;