Updated frame data plane bytes logic to fix random segfaults

frame_data.go

@@ -1,6 +1,14 @@
 package astiav
 
+//#cgo pkg-config: libavutil
+//#include <libavutil/imgutils.h>
 //#include <stdint.h>
+/*
+ptrdiff_t astiavFFAlign(int i, int align)
+{
+	return FFALIGN(i, align);
+}
+*/
 import "C"
 import (
 	"errors"
@@ -14,44 +22,24 @@ type FrameData struct {
 }
 
 type frameDataFramer interface {
+	bytes(align int) ([]byte, error)
 	height() int
-	imageBufferSize(align int) (int, error)
-	imageCopyToBuffer(b []byte, align int) (int, error)
-	linesize(i int) int
 	pixelFormat() PixelFormat
-	planeBytes(i int) []byte
+	planes() ([]frameDataPlane, error)
 	width() int
 }
 
+type frameDataPlane struct {
+	bytes    []byte
+	linesize int
+}
+
 func newFrameData(f frameDataFramer) *FrameData {
 	return &FrameData{f: f}
 }
 
 func (d *FrameData) Bytes(align int) ([]byte, error) {
-	switch {
-	// Video
-	case d.f.height() > 0 && d.f.width() > 0:
-		// Get buffer size
-		s, err := d.f.imageBufferSize(align)
-		if err != nil {
-			return nil, fmt.Errorf("astiav: getting image buffer size failed: %w", err)
-		}
-
-		// Invalid buffer size
-		if s == 0 {
-			return nil, errors.New("astiav: invalid image buffer size")
-		}
-
-		// Create buffer
-		b := make([]byte, s)
-
-		// Copy image to buffer
-		if _, err = d.f.imageCopyToBuffer(b, align); err != nil {
-			return nil, fmt.Errorf("astiav: copying image to buffer failed: %w", err)
-		}
-		return b, nil
-	}
-	return nil, errors.New("astiav: frame type not implemented")
+	return d.f.bytes(align)
 }
 
 // Always returns non-premultiplied formats when dealing with alpha channels, however this might not
@@ -100,9 +88,9 @@ func (d *FrameData) imageYCbCrSubsampleRatio() image.YCbCrSubsampleRatio {
 	return image.YCbCrSubsampleRatio444
 }
 
-func (d *FrameData) toImagePix(pix *[]uint8, stride *int, rect *image.Rectangle) {
-	*pix = d.f.planeBytes(0)
-	if v := d.f.linesize(0); *stride != v {
+func (d *FrameData) toImagePix(pix *[]uint8, stride *int, rect *image.Rectangle, planes []frameDataPlane) {
+	*pix = planes[0].bytes
+	if v := planes[0].linesize; *stride != v {
 		*stride = v
 	}
 	if w, h := d.f.width(), d.f.height(); rect.Dy() != w || rect.Dx() != h {
@@ -110,14 +98,14 @@ func (d *FrameData) toImagePix(pix *[]uint8, stride *int, rect *image.Rectangle)
 	}
 }
 
-func (d *FrameData) toImageYCbCr(y, cb, cr *[]uint8, yStride, cStride *int, subsampleRatio *image.YCbCrSubsampleRatio, rect *image.Rectangle) {
-	*y = d.f.planeBytes(0)
-	*cb = d.f.planeBytes(1)
-	*cr = d.f.planeBytes(2)
-	if v := d.f.linesize(0); *yStride != v {
+func (d *FrameData) toImageYCbCr(y, cb, cr *[]uint8, yStride, cStride *int, subsampleRatio *image.YCbCrSubsampleRatio, rect *image.Rectangle, planes []frameDataPlane) {
+	*y = planes[0].bytes
+	*cb = planes[1].bytes
+	*cr = planes[2].bytes
+	if v := planes[0].linesize; *yStride != v {
 		*yStride = v
 	}
-	if v := d.f.linesize(1); *cStride != v {
+	if v := planes[1].linesize; *cStride != v {
 		*cStride = v
 	}
 	if v := d.imageYCbCrSubsampleRatio(); *subsampleRatio != v {
@@ -128,37 +116,44 @@ func (d *FrameData) toImageYCbCr(y, cb, cr *[]uint8, yStride, cStride *int, subs
 	}
 }
 
-func (d *FrameData) toImageYCbCrA(y, cb, cr, a *[]uint8, yStride, cStride, aStride *int, subsampleRatio *image.YCbCrSubsampleRatio, rect *image.Rectangle) {
-	d.toImageYCbCr(y, cb, cr, yStride, cStride, subsampleRatio, rect)
-	*a = d.f.planeBytes(3)
-	if v := d.f.linesize(3); *aStride != v {
+func (d *FrameData) toImageYCbCrA(y, cb, cr, a *[]uint8, yStride, cStride, aStride *int, subsampleRatio *image.YCbCrSubsampleRatio, rect *image.Rectangle, planes []frameDataPlane) {
+	d.toImageYCbCr(y, cb, cr, yStride, cStride, subsampleRatio, rect, planes)
+	*a = planes[3].bytes
+	if v := planes[3].linesize; *aStride != v {
 		*aStride = v
 	}
 }
 
 func (d *FrameData) ToImage(dst image.Image) error {
+	// Get planes
+	planes, err := d.f.planes()
+	if err != nil {
+		return fmt.Errorf("astiav: getting planes failed: %w", err)
+	}
+
+	// Update image
 	if v, ok := dst.(*image.Alpha); ok {
-		d.toImagePix(&v.Pix, &v.Stride, &v.Rect)
+		d.toImagePix(&v.Pix, &v.Stride, &v.Rect, planes)
 	} else if v, ok := dst.(*image.Alpha16); ok {
-		d.toImagePix(&v.Pix, &v.Stride, &v.Rect)
+		d.toImagePix(&v.Pix, &v.Stride, &v.Rect, planes)
 	} else if v, ok := dst.(*image.CMYK); ok {
-		d.toImagePix(&v.Pix, &v.Stride, &v.Rect)
+		d.toImagePix(&v.Pix, &v.Stride, &v.Rect, planes)
 	} else if v, ok := dst.(*image.Gray); ok {
-		d.toImagePix(&v.Pix, &v.Stride, &v.Rect)
+		d.toImagePix(&v.Pix, &v.Stride, &v.Rect, planes)
 	} else if v, ok := dst.(*image.Gray16); ok {
-		d.toImagePix(&v.Pix, &v.Stride, &v.Rect)
+		d.toImagePix(&v.Pix, &v.Stride, &v.Rect, planes)
 	} else if v, ok := dst.(*image.NRGBA); ok {
-		d.toImagePix(&v.Pix, &v.Stride, &v.Rect)
+		d.toImagePix(&v.Pix, &v.Stride, &v.Rect, planes)
 	} else if v, ok := dst.(*image.NRGBA64); ok {
-		d.toImagePix(&v.Pix, &v.Stride, &v.Rect)
+		d.toImagePix(&v.Pix, &v.Stride, &v.Rect, planes)
 	} else if v, ok := dst.(*image.NYCbCrA); ok {
-		d.toImageYCbCrA(&v.Y, &v.Cb, &v.Cr, &v.A, &v.YStride, &v.CStride, &v.AStride, &v.SubsampleRatio, &v.Rect)
+		d.toImageYCbCrA(&v.Y, &v.Cb, &v.Cr, &v.A, &v.YStride, &v.CStride, &v.AStride, &v.SubsampleRatio, &v.Rect, planes)
 	} else if v, ok := dst.(*image.RGBA); ok {
-		d.toImagePix(&v.Pix, &v.Stride, &v.Rect)
+		d.toImagePix(&v.Pix, &v.Stride, &v.Rect, planes)
 	} else if v, ok := dst.(*image.RGBA64); ok {
-		d.toImagePix(&v.Pix, &v.Stride, &v.Rect)
+		d.toImagePix(&v.Pix, &v.Stride, &v.Rect, planes)
 	} else if v, ok := dst.(*image.YCbCr); ok {
-		d.toImageYCbCr(&v.Y, &v.Cb, &v.Cr, &v.YStride, &v.CStride, &v.SubsampleRatio, &v.Rect)
+		d.toImageYCbCr(&v.Y, &v.Cb, &v.Cr, &v.YStride, &v.CStride, &v.SubsampleRatio, &v.Rect, planes)
 	} else {
 		return errors.New("astiav: image format is not handled")
 	}
@@ -175,31 +170,95 @@ func newFrameDataFrame(f *Frame) *frameDataFrame {
 	return &frameDataFrame{f: f}
 }
 
+func (f *frameDataFrame) bytes(align int) ([]byte, error) {
+	switch {
+	// Video
+	case f.height() > 0 && f.width() > 0:
+		// Get buffer size
+		s, err := f.f.ImageBufferSize(align)
+		if err != nil {
+			return nil, fmt.Errorf("astiav: getting image buffer size failed: %w", err)
+		}
+
+		// Invalid buffer size
+		if s == 0 {
+			return nil, errors.New("astiav: invalid image buffer size")
+		}
+
+		// Create buffer
+		b := make([]byte, s)
+
+		// Copy image to buffer
+		if _, err = f.f.ImageCopyToBuffer(b, align); err != nil {
+			return nil, fmt.Errorf("astiav: copying image to buffer failed: %w", err)
+		}
+		return b, nil
+	}
+	return nil, errors.New("astiav: frame type not implemented")
+}
+
 func (f *frameDataFrame) height() int {
 	return f.f.Height()
 }
 
-func (f *frameDataFrame) imageBufferSize(align int) (int, error) {
-	return f.f.ImageBufferSize(align)
-}
-
-func (f *frameDataFrame) imageCopyToBuffer(b []byte, align int) (int, error) {
-	return f.f.ImageCopyToBuffer(b, align)
-}
-
-func (f *frameDataFrame) linesize(i int) int {
-	return f.f.Linesize()[i]
-}
-
 func (f *frameDataFrame) pixelFormat() PixelFormat {
 	return f.f.PixelFormat()
 }
 
-func (f *frameDataFrame) planeBytes(i int) []byte {
-	return bytesFromC(func(size *C.size_t) *C.uint8_t {
-		*size = C.size_t(int(f.f.c.linesize[i]) * f.f.Height())
-		return f.f.c.data[i]
+// Using bytesFromC on f.c.data caused random segfaults
+func (f *frameDataFrame) planes() ([]frameDataPlane, error) {
+	// Get bytes
+	const align = 1
+	b, err := f.bytes(align)
+	if err != nil {
+		return nil, fmt.Errorf("astiav: getting bytes failed: %w", err)
+	}
+
+	switch {
+	// Video
+	case f.height() > 0 && f.width() > 0:
+		// Below is mostly inspired by https://github.com/FFmpeg/FFmpeg/blob/n5.1.2/libavutil/imgutils.c#L466
+
+		// Get linesize
+		var linesize [4]C.int
+		if err := newError(C.av_image_fill_linesizes(&linesize[0], (C.enum_AVPixelFormat)(f.f.c.format), f.f.c.width)); err != nil {
+			return nil, fmt.Errorf("astiav: getting linesize failed: %w", err)
+		}
+
+		// Align linesize
+		var alignedLinesize [4]C.ptrdiff_t
+		for i := 0; i < 4; i++ {
+			alignedLinesize[i] = C.astiavFFAlign(linesize[i], C.int(align))
+		}
+
+		// Get plane sizes
+		var planeSizes [4]C.size_t
+		if err := newError(C.av_image_fill_plane_sizes(&planeSizes[0], (C.enum_AVPixelFormat)(f.f.c.format), f.f.c.height, &alignedLinesize[0])); err != nil {
+			return nil, fmt.Errorf("astiav: getting plane sizes failed: %w", err)
+		}
+
+		// Loop through plane sizes
+		var ps []frameDataPlane
+		start := 0
+		for idx, planeSize := range planeSizes {
+			// Get end
+			end := start + int(planeSize)
+			if len(b) < end {
+				return nil, fmt.Errorf("astiav: buffer length %d is invalid for [%d:%d]", len(b), start, end)
+			}
+
+			// Append plane
+			ps = append(ps, frameDataPlane{
+				bytes:    b[start:end],
+				linesize: int(alignedLinesize[idx]),
 			})
+
+			// Update start
+			start += int(planeSize)
+		}
+		return ps, nil
+	}
+	return nil, errors.New("astiav: frame type not implemented")
 }
 
 func (f *frameDataFrame) width() int {

frame_data_test.go

@@ -1,8 +1,8 @@
 package astiav
 
 import (
+	"fmt"
 	"image"
-	"image/png"
 	"os"
 	"testing"
 
@@ -12,37 +12,27 @@ import (
 type mockedFrameDataFrame struct {
 	h          int
 	imageBytes []byte
-	linesizes   []int
 	pf         PixelFormat
-	planesBytes [][]byte
+	planes_    []frameDataPlane
 	w          int
 }
 
 var _ frameDataFramer = (*mockedFrameDataFrame)(nil)
 
+func (f *mockedFrameDataFrame) bytes(align int) ([]byte, error) {
+	return f.imageBytes, nil
+}
+
 func (f *mockedFrameDataFrame) height() int {
 	return f.h
 }
 
-func (f *mockedFrameDataFrame) imageBufferSize(align int) (int, error) {
-	return len(f.imageBytes), nil
-}
-
-func (f *mockedFrameDataFrame) imageCopyToBuffer(b []byte, align int) (int, error) {
-	copy(b, f.imageBytes)
-	return len(f.imageBytes), nil
-}
-
-func (f *mockedFrameDataFrame) linesize(i int) int {
-	return f.linesizes[i]
-}
-
 func (f *mockedFrameDataFrame) pixelFormat() PixelFormat {
 	return f.pf
 }
 
-func (f *mockedFrameDataFrame) planeBytes(i int) []byte {
-	return f.planesBytes[i]
+func (f *mockedFrameDataFrame) planes() ([]frameDataPlane, error) {
+	return f.planes_, nil
 }
 
 func (f *mockedFrameDataFrame) width() int {
@@ -123,10 +113,8 @@ func TestFrameDataInternal(t *testing.T) {
 		}
 	}
 
-	fdf.imageBytes = []byte{0, 1, 2, 3}
-	_, err := fd.Bytes(0)
-	require.Error(t, err)
 	fdf.h = 1
+	fdf.imageBytes = []byte{0, 1, 2, 3}
 	fdf.w = 2
 	b, err := fd.Bytes(0)
 	require.NoError(t, err)
@@ -136,9 +124,8 @@ func TestFrameDataInternal(t *testing.T) {
 		e           image.Image
 		err         bool
 		i           image.Image
-		linesizes   []int
 		pixelFormat PixelFormat
-		planesBytes [][]byte
+		planes      []frameDataPlane
 	}{
 		{
 			e: &image.Alpha{
@@ -147,9 +134,13 @@ func TestFrameDataInternal(t *testing.T) {
 				Rect:   image.Rect(0, 0, 2, 1),
 			},
 			i:           &image.Alpha{},
-			linesizes:   []int{1},
 			pixelFormat: PixelFormatRgba,
-			planesBytes: [][]byte{{0, 1, 2, 3}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1, 2, 3},
+					linesize: 1,
+				},
+			},
 		},
 		{
 			e: &image.Alpha16{
@@ -158,9 +149,13 @@ func TestFrameDataInternal(t *testing.T) {
 				Rect:   image.Rect(0, 0, 2, 1),
 			},
 			i:           &image.Alpha16{},
-			linesizes:   []int{1},
 			pixelFormat: PixelFormatRgba,
-			planesBytes: [][]byte{{0, 1, 2, 3}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1, 2, 3},
+					linesize: 1,
+				},
+			},
 		},
 		{
 			e: &image.CMYK{
@@ -169,9 +164,13 @@ func TestFrameDataInternal(t *testing.T) {
 				Rect:   image.Rect(0, 0, 2, 1),
 			},
 			i:           &image.CMYK{},
-			linesizes:   []int{1},
 			pixelFormat: PixelFormatRgba,
-			planesBytes: [][]byte{{0, 1, 2, 3}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1, 2, 3},
+					linesize: 1,
+				},
+			},
 		},
 		{
 			e: &image.Gray{
@@ -180,9 +179,13 @@ func TestFrameDataInternal(t *testing.T) {
 				Rect:   image.Rect(0, 0, 2, 1),
 			},
 			i:           &image.Gray{},
-			linesizes:   []int{1},
 			pixelFormat: PixelFormatRgba,
-			planesBytes: [][]byte{{0, 1, 2, 3}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1, 2, 3},
+					linesize: 1,
+				},
+			},
 		},
 		{
 			e: &image.Gray16{
@@ -191,9 +194,13 @@ func TestFrameDataInternal(t *testing.T) {
 				Rect:   image.Rect(0, 0, 2, 1),
 			},
 			i:           &image.Gray16{},
-			linesizes:   []int{1},
 			pixelFormat: PixelFormatRgba,
-			planesBytes: [][]byte{{0, 1, 2, 3}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1, 2, 3},
+					linesize: 1,
+				},
+			},
 		},
 		{
 			e: &image.NRGBA{
@@ -202,9 +209,13 @@ func TestFrameDataInternal(t *testing.T) {
 				Rect:   image.Rect(0, 0, 2, 1),
 			},
 			i:           &image.NRGBA{},
-			linesizes:   []int{1},
 			pixelFormat: PixelFormatRgba,
-			planesBytes: [][]byte{{0, 1, 2, 3}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1, 2, 3},
+					linesize: 1,
+				},
+			},
 		},
 		{
 			e: &image.NRGBA64{
@@ -213,9 +224,13 @@ func TestFrameDataInternal(t *testing.T) {
 				Rect:   image.Rect(0, 0, 2, 1),
 			},
 			i:           &image.NRGBA64{},
-			linesizes:   []int{1},
 			pixelFormat: PixelFormatRgba,
-			planesBytes: [][]byte{{0, 1, 2, 3}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1, 2, 3},
+					linesize: 1,
+				},
+			},
 		},
 		{
 			e: &image.NYCbCrA{
@@ -232,9 +247,25 @@ func TestFrameDataInternal(t *testing.T) {
 				},
 			},
 			i:           &image.NYCbCrA{},
-			linesizes:   []int{1, 2, 3, 4},
 			pixelFormat: PixelFormatYuv444P,
-			planesBytes: [][]byte{{0, 1}, {2, 3}, {4, 5}, {6, 7}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1},
+					linesize: 1,
+				},
+				{
+					bytes:    []byte{2, 3},
+					linesize: 2,
+				},
+				{
+					bytes:    []byte{4, 5},
+					linesize: 3,
+				},
+				{
+					bytes:    []byte{6, 7},
+					linesize: 4,
+				},
+			},
 		},
 		{
 			e: &image.RGBA{
@@ -243,9 +274,13 @@ func TestFrameDataInternal(t *testing.T) {
 				Rect:   image.Rect(0, 0, 2, 1),
 			},
 			i:           &image.RGBA{},
-			linesizes:   []int{1},
 			pixelFormat: PixelFormatRgba,
-			planesBytes: [][]byte{{0, 1, 2, 3}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1, 2, 3},
+					linesize: 1,
+				},
+			},
 		},
 		{
 			e: &image.RGBA64{
@@ -254,9 +289,13 @@ func TestFrameDataInternal(t *testing.T) {
 				Rect:   image.Rect(0, 0, 2, 1),
 			},
 			i:           &image.RGBA64{},
-			linesizes:   []int{1},
 			pixelFormat: PixelFormatRgba,
-			planesBytes: [][]byte{{0, 1, 2, 3}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1, 2, 3},
+					linesize: 1,
+				},
+			},
 		},
 		{
 			e: &image.YCbCr{
@@ -269,14 +308,25 @@ func TestFrameDataInternal(t *testing.T) {
 				Rect:           image.Rect(0, 0, 2, 1),
 			},
 			i:           &image.YCbCr{},
-			linesizes:   []int{1, 2, 3},
 			pixelFormat: PixelFormatYuv420P,
-			planesBytes: [][]byte{{0, 1}, {2, 3}, {4, 5}},
+			planes: []frameDataPlane{
+				{
+					bytes:    []byte{0, 1},
+					linesize: 1,
+				},
+				{
+					bytes:    []byte{2, 3},
+					linesize: 2,
+				},
+				{
+					bytes:    []byte{4, 5},
+					linesize: 3,
+				},
+			},
 		},
 	} {
-		fdf.linesizes = v.linesizes
 		fdf.pf = v.pixelFormat
-		fdf.planesBytes = v.planesBytes
+		fdf.planes_ = v.planes
 		err = fd.ToImage(v.i)
 		if v.err {
 			require.Error(t, err)
@@ -287,30 +337,37 @@ func TestFrameDataInternal(t *testing.T) {
 }
 
 func TestFrameData(t *testing.T) {
-	const (
-		name = "image-rgba"
-		ext  = "png"
-	)
-	f, err := globalHelper.inputLastFrame(name+"."+ext, MediaTypeVideo)
+	for _, v := range []struct {
+		ext  string
+		name string
+	}{
+		{
+			ext:  "png",
+			name: "image-rgba",
+		},
+		{
+			ext:  "h264",
+			name: "video-yuv420p",
+		},
+	} {
+		f, err := globalHelper.inputLastFrame(v.name+"."+v.ext, MediaTypeVideo)
 		require.NoError(t, err)
 		fd := f.Data()
 
 		b1, err := fd.Bytes(1)
 		require.NoError(t, err)
-
-	b2, err := os.ReadFile("testdata/" + name + "-bytes")
+		b2 := []byte(fmt.Sprintf("%+v", b1))
+		b3, err := os.ReadFile("testdata/" + v.name + "-bytes")
 		require.NoError(t, err)
-	require.Equal(t, b1, b2)
-
-	f1, err := os.Open("testdata/" + name + "." + ext)
-	require.NoError(t, err)
-	defer f1.Close()
+		require.Equal(t, b2, b3)
 
 		i1, err := fd.GuessImageFormat()
 		require.NoError(t, err)
-	require.NoError(t, err)
 		require.NoError(t, fd.ToImage(i1))
-	i2, err := png.Decode(f1)
+		b4 := []byte(fmt.Sprintf("%+v", i1))
+		b5, err := os.ReadFile("testdata/" + v.name + "-struct")
 		require.NoError(t, err)
-	require.Equal(t, i1, i2)
+		require.Equal(t, b4, b5)
+
+	}
 }