[FFmpeg-devel,v6,2/2] avcodec/mjpegenc: support writing ICC profiles

diff --git a/libavcodec/ljpegenc.c b/libavcodec/ljpegenc.c
index fad19cbb76..382d291621 100644
--- a/libavcodec/ljpegenc.c
+++ b/libavcodec/ljpegenc.c
@@ -220,7 +220,7 @@  static int ljpeg_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
     const int height = avctx->height;
     const int mb_width  = (width  + s->hsample[0] - 1) / s->hsample[0];
     const int mb_height = (height + s->vsample[0] - 1) / s->vsample[0];
-    int max_pkt_size = AV_INPUT_BUFFER_MIN_SIZE;
+    size_t max_pkt_size = AV_INPUT_BUFFER_MIN_SIZE;
     int ret, header_bits;
 
     if(    avctx->pix_fmt == AV_PIX_FMT_BGR0
@@ -233,12 +233,14 @@  static int ljpeg_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
                         * s->hsample[0] * s->vsample[0];
     }
 
+    if ((ret = ff_mjpeg_add_icc_profile_size(avctx, pict, &max_pkt_size)) < 0)
+        return ret;
     if ((ret = ff_alloc_packet(avctx, pkt, max_pkt_size)) < 0)
         return ret;
 
     init_put_bits(&pb, pkt->data, pkt->size);
 
-    ff_mjpeg_encode_picture_header(avctx, &pb, NULL, &s->scantable,
+    ff_mjpeg_encode_picture_header(avctx, &pb, pict, NULL, &s->scantable,
                                    s->pred, s->matrix, s->matrix);
 
     header_bits = put_bits_count(&pb);
diff --git a/libavcodec/mjpegenc.c b/libavcodec/mjpegenc.c
index a39b990eea..d35195d52e 100644
--- a/libavcodec/mjpegenc.c
+++ b/libavcodec/mjpegenc.c
@@ -80,7 +80,7 @@  static av_cold void init_uni_ac_vlc(const uint8_t huff_size_ac[256],
 
 static void mjpeg_encode_picture_header(MpegEncContext *s)
 {
-    ff_mjpeg_encode_picture_header(s->avctx, &s->pb, s->mjpeg_ctx,
+    ff_mjpeg_encode_picture_header(s->avctx, &s->pb, s->picture->f, s->mjpeg_ctx,
                                    &s->intra_scantable, 0,
                                    s->intra_matrix, s->chroma_intra_matrix);
 
@@ -130,6 +130,7 @@  static void mjpeg_encode_picture_frame(MpegEncContext *s)
     }
 
     bytes_needed = (total_bits + 7) / 8;
+    ff_mjpeg_add_icc_profile_size(s->avctx, s->picture->f, &bytes_needed);
     ff_mpv_reallocate_putbitbuffer(s, bytes_needed, bytes_needed);
 
     for (int i = 0; i < m->huff_ncode; i++) {
diff --git a/libavcodec/mjpegenc_common.c b/libavcodec/mjpegenc_common.c
index 7b82644763..17103051a8 100644
--- a/libavcodec/mjpegenc_common.c
+++ b/libavcodec/mjpegenc_common.c
@@ -131,8 +131,39 @@  static void jpeg_table_header(AVCodecContext *avctx, PutBitContext *p,
     AV_WB16(ptr, size);
 }
 
-static void jpeg_put_comments(AVCodecContext *avctx, PutBitContext *p)
+static const uint8_t icc_hdr_size = 16; /* ICC_PROFILE\0 tag + 4 bytes */
+static const uint16_t icc_chunk_size = UINT16_MAX - icc_hdr_size;
+static const uint8_t icc_max_chunks = UINT8_MAX;
+
+int ff_mjpeg_add_icc_profile_size(AVCodecContext *avctx, const AVFrame *frame,
+                                  size_t *max_pkt_size)
 {
+    const AVFrameSideData *sd;
+    size_t new_pkt_size;
+    int nb_chunks;
+    sd = av_frame_get_side_data(frame, AV_FRAME_DATA_ICC_PROFILE);
+    if (!sd || !sd->size)
+        return 0;
+
+    if (sd->size > icc_max_chunks * icc_chunk_size) {
+        av_log(avctx, AV_LOG_ERROR, "Cannot store %"SIZE_SPECIFIER" byte ICC "
+               "profile: too large for JPEG\n",
+               sd->size);
+        return AVERROR_INVALIDDATA;
+    }
+
+    nb_chunks = (sd->size + icc_chunk_size - 1) / icc_chunk_size;
+    new_pkt_size = *max_pkt_size + nb_chunks * (UINT16_MAX + 2 /* APP2 marker */);
+    if (new_pkt_size < *max_pkt_size) /* overflow */
+        return AVERROR_INVALIDDATA;
+    *max_pkt_size = new_pkt_size;
+    return 0;
+}
+
+static void jpeg_put_comments(AVCodecContext *avctx, PutBitContext *p,
+                              const AVFrame *frame)
+{
+    const AVFrameSideData *sd = NULL;
     int size;
     uint8_t *ptr;
 
@@ -162,6 +193,35 @@  static void jpeg_put_comments(AVCodecContext *avctx, PutBitContext *p)
         put_bits(p, 8, 0); /* thumbnail height */
     }
 
+    /* ICC profile */
+    sd = av_frame_get_side_data(frame, AV_FRAME_DATA_ICC_PROFILE);
+    if (sd && sd->size) {
+        const int nb_chunks = (sd->size + icc_chunk_size - 1) / icc_chunk_size;
+        const uint8_t *data = sd->data;
+        size_t remaining = sd->size;
+        /* must already be checked by the packat allocation code */
+        av_assert0(remaining <= icc_max_chunks * icc_chunk_size);
+        flush_put_bits(p);
+        for (int i = 0; i < nb_chunks; i++) {
+            size = FFMIN(remaining, icc_chunk_size);
+            av_assert1(size > 0);
+            ptr = put_bits_ptr(p);
+            ptr[0] = 0xff; /* chunk marker, not part of icc_hdr_size */
+            ptr[1] = APP2;
+            AV_WB16(ptr+2, size + icc_hdr_size);
+            AV_WL32(ptr+4,  MKTAG('I','C','C','_'));
+            AV_WL32(ptr+8,  MKTAG('P','R','O','F'));
+            AV_WL32(ptr+12, MKTAG('I','L','E','\0'));
+            ptr[16] = i+1;
+            ptr[17] = nb_chunks;
+            memcpy(&ptr[18], data, size);
+            skip_put_bytes(p, size + icc_hdr_size + 2);
+            remaining -= size;
+            data += size;
+        }
+        av_assert1(!remaining);
+    }
+
     /* comment */
     if (!(avctx->flags & AV_CODEC_FLAG_BITEXACT)) {
         put_marker(p, COM);
@@ -214,7 +274,7 @@  void ff_mjpeg_init_hvsample(AVCodecContext *avctx, int hsample[4], int vsample[4
 }
 
 void ff_mjpeg_encode_picture_header(AVCodecContext *avctx, PutBitContext *pb,
-                                    MJpegContext *m,
+                                    const AVFrame *frame, struct MJpegContext *m,
                                     ScanTable *intra_scantable, int pred,
                                     uint16_t luma_intra_matrix[64],
                                     uint16_t chroma_intra_matrix[64])
@@ -234,7 +294,7 @@  void ff_mjpeg_encode_picture_header(AVCodecContext *avctx, PutBitContext *pb,
     if (avctx->codec_id == AV_CODEC_ID_AMV)
         return;
 
-    jpeg_put_comments(avctx, pb);
+    jpeg_put_comments(avctx, pb, frame);
 
     jpeg_table_header(avctx, pb, m, intra_scantable,
                       luma_intra_matrix, chroma_intra_matrix, hsample);
diff --git a/libavcodec/mjpegenc_common.h b/libavcodec/mjpegenc_common.h
index ac753bf153..69e132f223 100644
--- a/libavcodec/mjpegenc_common.h
+++ b/libavcodec/mjpegenc_common.h
@@ -29,8 +29,10 @@ 
 
 struct MJpegContext;
 
+int ff_mjpeg_add_icc_profile_size(AVCodecContext *avctx, const AVFrame *frame,
+                                  size_t *max_pkt_size);
 void ff_mjpeg_encode_picture_header(AVCodecContext *avctx, PutBitContext *pb,
-                                    struct MJpegContext *m,
+                                    const AVFrame *frame, struct MJpegContext *m,
                                     ScanTable *intra_scantable, int pred,
                                     uint16_t luma_intra_matrix[64],
                                     uint16_t chroma_intra_matrix[64]);
diff --git a/libavcodec/mpegvideo_enc.c b/libavcodec/mpegvideo_enc.c
index 71c999fab0..7fe889dec3 100644
--- a/libavcodec/mpegvideo_enc.c
+++ b/libavcodec/mpegvideo_enc.c
@@ -1689,9 +1689,11 @@  int ff_mpv_encode_picture(AVCodecContext *avctx, AVPacket *pkt,
     /* output? */
     if (s->new_picture.f->data[0]) {
         int growing_buffer = context_count == 1 && !pkt->data && !s->data_partitioning;
-        int pkt_size = growing_buffer ? FFMAX(s->mb_width*s->mb_height*64+10000, avctx->internal->byte_buffer_size) - AV_INPUT_BUFFER_PADDING_SIZE
+        size_t pkt_size = growing_buffer ? FFMAX(s->mb_width*s->mb_height*64+10000, avctx->internal->byte_buffer_size) - AV_INPUT_BUFFER_PADDING_SIZE
                                               :
                                               s->mb_width*s->mb_height*(MAX_MB_BYTES+100)+10000;
+        if ((ret = ff_mjpeg_add_icc_profile_size(avctx, s->new_picture.f, &pkt_size)) < 0)
+            return ret;
         if ((ret = ff_alloc_packet(avctx, pkt, pkt_size)) < 0)
             return ret;
         if (s->mb_info) {
diff --git a/tests/fate/image.mak b/tests/fate/image.mak
index c6374c7d8a..70be281411 100644
--- a/tests/fate/image.mak
+++ b/tests/fate/image.mak
@@ -337,9 +337,13 @@  fate-jpg-12bpp: CMD = framecrc -idct simple -i $(TARGET_SAMPLES)/jpg/12bpp.jpg -
 FATE_JPG += fate-jpg-jfif
 fate-jpg-jfif: CMD = framecrc -idct simple -i $(TARGET_SAMPLES)/jpg/20242.jpg
 
+FATE_JPG_TRANSCODE-$(call ENCDEC, MJPEG, IMAGE2) += fate-jpg-icc
+fate-jpg-icc: CMD = transcode png_pipe $(TARGET_SAMPLES)/png1/lena-int_rgb24.png mjpeg "-vf scale" "" "" "-show_frames"
+
 FATE_JPG-$(call DEMDEC, IMAGE2, MJPEG) += $(FATE_JPG)
 FATE_IMAGE += $(FATE_JPG-yes)
-fate-jpg: $(FATE_JPG-yes)
+FATE_IMAGE_TRANSCODE += $(FATE_JPG_TRANSCODE-yes)
+fate-jpg: $(FATE_JPG-yes) $(FATE_JPG_TRANSCODE-yes)
 
 FATE_JPEGLS += fate-jpegls-2bpc
 fate-jpegls-2bpc: CMD = framecrc -idct simple -i $(TARGET_SAMPLES)/jpegls/4.jls
diff --git a/tests/ref/fate/jpg-icc b/tests/ref/fate/jpg-icc
new file mode 100644
index 0000000000..220146555e
--- /dev/null
+++ b/tests/ref/fate/jpg-icc
@@ -0,0 +1,42 @@ 
+0a323df5cdfb9574e329b9831be054a6 *tests/data/fate/jpg-icc.mjpeg
+11010 tests/data/fate/jpg-icc.mjpeg
+#tb 0: 1/25
+#media_type 0: video
+#codec_id 0: rawvideo
+#dimensions 0: 128x128
+#sar 0: 1/1
+0,          0,          0,        1,    49152, 0xaac06b42
+[FRAME]
+media_type=video
+stream_index=0
+key_frame=1
+pts=0
+pts_time=0.000000
+pkt_dts=0
+pkt_dts_time=0.000000
+best_effort_timestamp=0
+best_effort_timestamp_time=0.000000
+pkt_duration=1
+pkt_duration_time=0.040000
+pkt_pos=0
+pkt_size=11010
+width=128
+height=128
+pix_fmt=yuvj444p
+sample_aspect_ratio=1:1
+pict_type=I
+coded_picture_number=0
+display_picture_number=0
+interlaced_frame=0
+top_field_first=0
+repeat_pict=0
+color_range=pc
+color_space=bt470bg
+color_primaries=unknown
+color_transfer=unknown
+chroma_location=center
+[SIDE_DATA]
+side_data_type=ICC profile
+size=3144
+[/SIDE_DATA]
+[/FRAME]