Skip to content

Commit

Permalink
Updated frame data plane bytes logic to fix random segfaults
Browse files Browse the repository at this point in the history
  • Loading branch information
@asticode
asticode committed Aug 14, 2024
1 parent 33d572d commit 9e3ff1d
Show file tree
Hide file tree
Showing 9 changed files with 275 additions and 141 deletions.
193 changes: 126 additions & 67 deletions frame_data.go
View file
Original file line number Diff line number Diff line change
@@ -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"
Expand All @@ -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() ([]frameDataFramerPlane, error)
width() int
}

type frameDataFramerPlane 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
Expand Down Expand Up @@ -100,24 +88,24 @@ 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 []frameDataFramerPlane) {
*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 {
*rect = image.Rect(0, 0, w, h)
}
}

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 []frameDataFramerPlane) {
*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 {
Expand All @@ -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 []frameDataFramerPlane) {
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")
}
Expand All @@ -175,31 +170,95 @@ func newFrameDataFrame(f *Frame) *frameDataFrame {
return &frameDataFrame{f: f}
}

func (f *frameDataFrame) height() int {
return f.f.Height()
}
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)
}

func (f *frameDataFrame) imageBufferSize(align int) (int, error) {
return f.f.ImageBufferSize(align)
}
// Invalid buffer size
if s == 0 {
return nil, errors.New("astiav: invalid image buffer size")
}

// Create buffer
b := make([]byte, s)

func (f *frameDataFrame) imageCopyToBuffer(b []byte, align int) (int, error) {
return f.f.ImageCopyToBuffer(b, align)
// 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) linesize(i int) int {
return f.f.Linesize()[i]
func (f *frameDataFrame) height() int {
return f.f.Height()
}

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() ([]frameDataFramerPlane, 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 []frameDataFramerPlane
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, frameDataFramerPlane{
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 {
Expand Down
Loading

0 comments on commit 9e3ff1d

Please sign in to comment.