diff --git a/libavcodec/cbs.c b/libavcodec/cbs.c
index 0badb192d9..aa2e4904d0 100644
--- a/libavcodec/cbs.c
+++ b/libavcodec/cbs.c
@@ -283,6 +283,28 @@ int ff_cbs_read(CodedBitstreamContext *ctx,
     return cbs_read_fragment_content(ctx, frag);
 }
 
+int ff_cbs_default_write_unit_data(CodedBitstreamContext *ctx,
+                                   CodedBitstreamUnit *unit,
+                                   PutBitContext *pbc)
+{
+    int err;
+
+    if (put_bits_count(pbc) % 8)
+        unit->data_bit_padding = 8 - put_bits_count(pbc) % 8;
+    else
+        unit->data_bit_padding = 0;
+
+    flush_put_bits(pbc);
+
+    err = ff_cbs_alloc_unit_data(ctx, unit, put_bits_count(pbc) / 8);
+    if (err < 0)
+        return err;
+
+    memcpy(unit->data, pbc->buf, unit->data_size);
+
+    return 0;
+}
+
 static int cbs_write_unit_data(CodedBitstreamContext *ctx,
                                CodedBitstreamUnit *unit)
 {
@@ -318,22 +340,6 @@ static int cbs_write_unit_data(CodedBitstreamContext *ctx,
         return ret;
     }
 
-    // Overflow but we didn't notice.
-    av_assert0(put_bits_count(&pbc) <= 8 * ctx->write_buffer_size);
-
-    if (put_bits_count(&pbc) % 8)
-        unit->data_bit_padding = 8 - put_bits_count(&pbc) % 8;
-    else
-        unit->data_bit_padding = 0;
-
-    flush_put_bits(&pbc);
-
-    ret = ff_cbs_alloc_unit_data(ctx, unit, put_bits_count(&pbc) / 8);
-    if (ret < 0)
-        return ret;
-
-    memcpy(unit->data, ctx->write_buffer, unit->data_size);
-
     return 0;
 }
 
diff --git a/libavcodec/cbs_av1.c b/libavcodec/cbs_av1.c
index bbe4461130..7ff5539fb8 100644
--- a/libavcodec/cbs_av1.c
+++ b/libavcodec/cbs_av1.c
@@ -1217,7 +1217,7 @@ static int cbs_av1_write_obu(CodedBitstreamContext *ctx,
     // OBU data must be byte-aligned.
     av_assert0(put_bits_count(pbc) % 8 == 0);
 
-    return 0;
+    return ff_cbs_default_write_unit_data(ctx, unit, pbc);
 }
 
 static int cbs_av1_assemble_fragment(CodedBitstreamContext *ctx,
diff --git a/libavcodec/cbs_h2645.c b/libavcodec/cbs_h2645.c
index 88fa0029cd..d3d5342bad 100644
--- a/libavcodec/cbs_h2645.c
+++ b/libavcodec/cbs_h2645.c
@@ -1265,7 +1265,7 @@ static int cbs_h264_write_nal_unit(CodedBitstreamContext *ctx,
         return AVERROR_PATCHWELCOME;
     }
 
-    return 0;
+    return ff_cbs_default_write_unit_data(ctx, unit, pbc);
 }
 
 static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
@@ -1377,7 +1377,7 @@ static int cbs_h265_write_nal_unit(CodedBitstreamContext *ctx,
         return AVERROR_PATCHWELCOME;
     }
 
-    return 0;
+    return ff_cbs_default_write_unit_data(ctx, unit, pbc);
 }
 
 static int cbs_h2645_assemble_fragment(CodedBitstreamContext *ctx,
diff --git a/libavcodec/cbs_internal.h b/libavcodec/cbs_internal.h
index 4c5a535ca6..96c2a1ac80 100644
--- a/libavcodec/cbs_internal.h
+++ b/libavcodec/cbs_internal.h
@@ -44,7 +44,9 @@ typedef struct CodedBitstreamType {
     int (*read_unit)(CodedBitstreamContext *ctx,
                      CodedBitstreamUnit *unit);
 
-    // Write the data bitstream from unit->content into pbc.
+    // Write the data bitstream from unit->content. This includes
+    // allocating the buffer as well as setting all of unit's data-fields.
+    // The supplied PutBitContext may be used as temporary buffer.
     // Return value AVERROR(ENOSPC) indicates that pbc was too small.
     int (*write_unit)(CodedBitstreamContext *ctx,
                       CodedBitstreamUnit *unit,
@@ -60,6 +62,14 @@ typedef struct CodedBitstreamType {
 } CodedBitstreamType;
 
 
+// Helper function to convert pbc's content into a unit's data.
+// Allocates and copies data and sets all of unit's data-fields.
+
+int ff_cbs_default_write_unit_data(CodedBitstreamContext *ctx,
+                                   CodedBitstreamUnit *unit,
+                                   PutBitContext *pbc);
+
+
 // Helper functions for trace output.
 
 void ff_cbs_trace_header(CodedBitstreamContext *ctx,
diff --git a/libavcodec/cbs_jpeg.c b/libavcodec/cbs_jpeg.c
index 64fe70beab..da22bdf720 100644
--- a/libavcodec/cbs_jpeg.c
+++ b/libavcodec/cbs_jpeg.c
@@ -384,10 +384,16 @@ static int cbs_jpeg_write_unit(CodedBitstreamContext *ctx,
                                CodedBitstreamUnit *unit,
                                PutBitContext *pbc)
 {
+    int err;
+
     if (unit->type == JPEG_MARKER_SOS)
-        return cbs_jpeg_write_scan   (ctx, unit, pbc);
+        err = cbs_jpeg_write_scan   (ctx, unit, pbc);
     else
-        return cbs_jpeg_write_segment(ctx, unit, pbc);
+        err = cbs_jpeg_write_segment(ctx, unit, pbc);
+    if (err < 0)
+        return err;
+
+    return ff_cbs_default_write_unit_data(ctx, unit, pbc);
 }
 
 static int cbs_jpeg_assemble_fragment(CodedBitstreamContext *ctx,
diff --git a/libavcodec/cbs_mpeg2.c b/libavcodec/cbs_mpeg2.c
index 13d871cc89..ba5e2417c2 100644
--- a/libavcodec/cbs_mpeg2.c
+++ b/libavcodec/cbs_mpeg2.c
@@ -374,10 +374,16 @@ static int cbs_mpeg2_write_unit(CodedBitstreamContext *ctx,
                                 CodedBitstreamUnit *unit,
                                 PutBitContext *pbc)
 {
+    int err;
+
     if (MPEG2_START_IS_SLICE(unit->type))
-        return cbs_mpeg2_write_slice (ctx, unit, pbc);
+        err = cbs_mpeg2_write_slice (ctx, unit, pbc);
     else
-        return cbs_mpeg2_write_header(ctx, unit, pbc);
+        err = cbs_mpeg2_write_header(ctx, unit, pbc);
+    if (err < 0)
+        return err;
+
+    return ff_cbs_default_write_unit_data(ctx, unit, pbc);
 }
 
 static int cbs_mpeg2_assemble_fragment(CodedBitstreamContext *ctx,
diff --git a/libavcodec/cbs_vp9.c b/libavcodec/cbs_vp9.c
index 42e4dcf5ac..a74c4f4cec 100644
--- a/libavcodec/cbs_vp9.c
+++ b/libavcodec/cbs_vp9.c
@@ -544,7 +544,7 @@ static int cbs_vp9_write_unit(CodedBitstreamContext *ctx,
         skip_put_bytes(pbc, frame->data_size);
     }
 
-    return 0;
+    return ff_cbs_default_write_unit_data(ctx, unit, pbc);
 }
 
 static int cbs_vp9_assemble_fragment(CodedBitstreamContext *ctx,
