Added alloc io context

class_test.go

@@ -57,7 +57,7 @@ func TestClassers(t *testing.T) {
 	f.Free()
 	fmc1.Free()
 	fmc2.CloseInput()
-	require.NoError(t, ic.Closep())
+	require.NoError(t, ic.Free())
 	ssc.Free()
 	require.Equal(t, cl, len(classers.p))
 }

format_context_test.go

@@ -45,7 +45,7 @@ func TestFormatContext(t *testing.T) {
 	defer fc3.Free()
 	c, err := OpenIOContext("testdata/video.mp4", NewIOContextFlags(IOContextFlagRead))
 	require.NoError(t, err)
-	defer c.Closep() //nolint:errcheck
+	defer c.Free() //nolint:errcheck
 	fc3.SetPb(c)
 	fc3.SetStrictStdCompliance(StrictStdComplianceExperimental)
 	require.NotNil(t, fc3.Pb())

io_context.c

@@ -0,0 +1,17 @@
+#include "io_context.h"
+#include <stdint.h>
+
+int astiavIOContextReadFunc(void *opaque, uint8_t *buf, int buf_size)
+{
+    return goAstiavIOContextReadFunc(opaque, buf, buf_size);
+}
+
+int64_t astiavIOContextSeekFunc(void *opaque, int64_t offset, int whence)
+{
+    return goAstiavIOContextSeekFunc(opaque, offset, whence);
+}
+
+int astiavIOContextWriteFunc(void *opaque, uint8_t *buf, int buf_size)
+{
+    return goAstiavIOContextWriteFunc(opaque, buf, buf_size);
+}

io_context.go

@@ -2,12 +2,20 @@ package astiav
 
 //#cgo pkg-config: libavformat
 //#include <libavformat/avformat.h>
+//#include "io_context.h"
 import "C"
-import "unsafe"
+import (
+	"errors"
+	"fmt"
+	"sync"
+	"unsafe"
+)
 
 // https://github.com/FFmpeg/FFmpeg/blob/n5.0/libavformat/avio.h#L161
 type IOContext struct {
-	c *C.struct_AVIOContext
+	buffer    unsafe.Pointer
+	c         *C.struct_AVIOContext
+	handlerID unsafe.Pointer
 }
 
 func newIOContextFromC(c *C.struct_AVIOContext) *IOContext {
@@ -21,6 +29,68 @@ func newIOContextFromC(c *C.struct_AVIOContext) *IOContext {
 
 var _ Classer = (*IOContext)(nil)
 
+type IOContextReadFunc func(b []byte) (n int, err error)
+
+type IOContextSeekFunc func(offset int64, whence int) (n int64, err error)
+
+type IOContextWriteFunc func(b []byte) (n int, err error)
+
+func AllocIOContext(bufferSize int, readFunc IOContextReadFunc, seekFunc IOContextSeekFunc, writeFunc IOContextWriteFunc) (ic *IOContext, err error) {
+	// Invalid buffer size
+	if bufferSize <= 0 {
+		err = errors.New("astiav: buffer size <= 0")
+		return
+	}
+
+	// Alloc buffer
+	buffer := C.av_malloc(C.size_t(bufferSize))
+	if buffer == nil {
+		err = errors.New("astiav: allocating buffer failed")
+		return
+	}
+
+	// Make sure buffer is freed in case of error
+	defer func() {
+		if err != nil {
+			C.av_free(buffer)
+		}
+	}()
+
+	// Since go doesn't allow c to store pointers to go data, we need to create this C pointer
+	handlerID := C.malloc(C.size_t(1))
+	if handlerID == nil {
+		err = errors.New("astiav: allocating handler id failed")
+		return
+	}
+
+	// Make sure handler id is freed in case of error
+	defer func() {
+		if err != nil {
+			C.free(handlerID)
+		}
+	}()
+
+	// Alloc io context
+	cic := C.avio_alloc_context((*C.uchar)(buffer), C.int(bufferSize), 1, handlerID, (*[0]byte)(C.astiavIOContextReadFunc), (*[0]byte)(C.astiavIOContextWriteFunc), (*[0]byte)(C.astiavIOContextSeekFunc))
+	if cic == nil {
+		err = errors.New("astiav: allocating io context failed: %w")
+		return
+	}
+
+	// Create io context
+	ic = newIOContextFromC(cic)
+
+	// Store buffer and handler
+	ic.buffer = buffer
+	ic.handlerID = handlerID
+	ioContextHandlers.set(handlerID, &ioContextHandler{
+		r: readFunc,
+		s: seekFunc,
+		w: writeFunc,
+	})
+	return
+}
+
 func OpenIOContext(filename string, flags IOContextFlags) (*IOContext, error) {
 	cfi := C.CString(filename)
 	defer C.free(unsafe.Pointer(cfi))
@@ -35,17 +105,168 @@ func (ic *IOContext) Class() *Class {
 	return newClassFromC(unsafe.Pointer(ic.c))
 }
 
-func (ic *IOContext) Closep() error {
+func (ic *IOContext) Free() error {
 	classers.del(ic)
 	if ic.c != nil {
-		return newError(C.avio_closep(&ic.c))
+		if ic.buffer == nil {
+			if err := newError(C.avio_closep(&ic.c)); err != nil {
+				return err
+			}
+		} else {
+			C.av_free(ic.buffer)
+			C.free(ic.handlerID)
+			C.avio_context_free(&ic.c)
+		}
 	}
 	return nil
 }
 
+func (ic *IOContext) Read(b []byte) (n int, err error) {
+	// Nothing to read
+	if b == nil || len(b) <= 0 {
+		return
+	}
+
+	// Alloc buffer
+	buf := C.av_malloc(C.size_t(len(b)))
+	if buf == nil {
+		err = errors.New("astiav: allocating buffer failed")
+		return
+	}
+
+	// Make sure buffer is freed
+	defer C.av_free(buf)
+
+	// Read
+	ret := C.avio_read_partial(ic.c, (*C.uchar)(unsafe.Pointer(buf)), C.int(len(b)))
+	if err = newError(ret); err != nil {
+		err = fmt.Errorf("astiav: reading failed: %w", err)
+		return
+	}
+
+	// Copy
+	C.memcpy(unsafe.Pointer(&b[0]), unsafe.Pointer(buf), C.size_t(ret))
+	n = int(ret)
+	return
+}
+
+func (ic *IOContext) Seek(offset int64, whence int) (int64, error) {
+	ret := C.avio_seek(ic.c, C.int64_t(offset), C.int(whence))
+	if err := newError(C.int(ret)); err != nil {
+		return 0, err
+	}
+	return int64(ret), nil
+}
+
 func (ic *IOContext) Write(b []byte) {
-	if b == nil {
+	// Nothing to write
+	if b == nil || len(b) <= 0 {
 		return
 	}
+
+	// Write
 	C.avio_write(ic.c, (*C.uchar)(unsafe.Pointer(&b[0])), C.int(len(b)))
 }
+
+func (ic *IOContext) Flush() {
+	C.avio_flush(ic.c)
+}
+
+type ioContextHandler struct {
+	r IOContextReadFunc
+	s IOContextSeekFunc
+	w IOContextWriteFunc
+}
+
+var ioContextHandlers = newioContextHandlerPool()
+
+type ioContextHandlerPool struct {
+	m sync.Mutex
+	p map[unsafe.Pointer]*ioContextHandler
+}
+
+func newioContextHandlerPool() *ioContextHandlerPool {
+	return &ioContextHandlerPool{p: make(map[unsafe.Pointer]*ioContextHandler)}
+}
+
+func (p *ioContextHandlerPool) set(id unsafe.Pointer, h *ioContextHandler) {
+	p.m.Lock()
+	defer p.m.Unlock()
+	p.p[id] = h
+}
+
+func (p *ioContextHandlerPool) get(id unsafe.Pointer) (h *ioContextHandler, ok bool) {
+	p.m.Lock()
+	defer p.m.Unlock()
+	h, ok = p.p[id]
+	return
+}
+
+//export goAstiavIOContextReadFunc
+func goAstiavIOContextReadFunc(opaque unsafe.Pointer, buf *C.uint8_t, bufSize C.int) C.int {
+	// Get handler
+	h, ok := ioContextHandlers.get(opaque)
+	if !ok {
+		return C.AVERROR_UNKNOWN
+	}
+
+	// No callback
+	if h.r == nil {
+		return C.int(0)
+	}
+
+	// Create go buffer
+	b := make([]byte, int(bufSize), int(bufSize))
+
+	// Read
+	n, err := h.r(b)
+	if err != nil {
+		return C.AVERROR_UNKNOWN
+	}
+
+	// Copy
+	C.memcpy(unsafe.Pointer(buf), unsafe.Pointer(&b[0]), C.size_t(n))
+	return C.int(n)
+}
+
+//export goAstiavIOContextSeekFunc
+func goAstiavIOContextSeekFunc(opaque unsafe.Pointer, offset C.int64_t, whence C.int) C.int64_t {
+	// Get handler
+	h, ok := ioContextHandlers.get(opaque)
+	if !ok {
+		return C.AVERROR_UNKNOWN
+	}
+
+	// No callback
+	if h.s == nil {
+		return offset
+	}
+
+	// Seek
+	n, err := h.s(int64(offset), int(whence))
+	if err != nil {
+		return C.int64_t(C.AVERROR_UNKNOWN)
+	}
+	return C.int64_t(n)
+}
+
+//export goAstiavIOContextWriteFunc
+func goAstiavIOContextWriteFunc(opaque unsafe.Pointer, buf *C.uint8_t, bufSize C.int) C.int {
+	// Get handler
+	h, ok := ioContextHandlers.get(opaque)
+	if !ok {
+		return C.AVERROR_UNKNOWN
+	}
+
+	// No callback
+	if h.w == nil {
+		return bufSize
+	}
+
+	// Write
+	n, err := h.w(C.GoBytes(unsafe.Pointer(buf), bufSize))
+	if err != nil {
+		return C.AVERROR_UNKNOWN
+	}
+	return C.int(n)
+}

io_context.h

@@ -0,0 +1,9 @@
+#include <stdint.h>
+
+extern int goAstiavIOContextReadFunc(void *opaque, uint8_t *buf, int buf_size);
+extern int64_t goAstiavIOContextSeekFunc(void *opaque, int64_t offset, int whence);
+extern int goAstiavIOContextWriteFunc(void *opaque, uint8_t *buf, int buf_size);
+
+int astiavIOContextReadFunc(void *opaque, uint8_t *buf, int buf_size);
+int64_t astiavIOContextSeekFunc(void *opaque, int64_t offset, int whence);
+int astiavIOContextWriteFunc(void *opaque, uint8_t *buf, int buf_size);

io_context_test.go

@@ -1,6 +1,7 @@
 package astiav
 
 import (
+	"bytes"
 	"os"
 	"path/filepath"
 	"testing"
@@ -9,6 +10,32 @@ import (
 )
 
 func TestIOContext(t *testing.T) {
+	var seeked bool
+	written := bytes.NewBuffer([]byte{})
+	c, err := AllocIOContext(8, func(b []byte) (int, error) {
+		return bytes.NewBuffer(b).Write([]byte("read"))
+	}, func(offset int64, whence int) (n int64, err error) {
+		seeked = true
+		return offset, nil
+	}, func(b []byte) (int, error) {
+		return written.Write(b)
+	})
+	require.NoError(t, err)
+	defer c.Free()
+	b := make([]byte, 6)
+	n, err := c.Read(b)
+	require.NoError(t, err)
+	require.Equal(t, 4, n)
+	_, err = c.Seek(2, 0)
+	require.NoError(t, err)
+	require.True(t, seeked)
+	b = []byte("write")
+	c.Write([]byte("write"))
+	c.Flush()
+	require.Equal(t, b, written.Bytes())
+}
+
+func TestOpenIOContext(t *testing.T) {
 	path := filepath.Join(t.TempDir(), "iocontext.txt")
 	c, err := OpenIOContext(path, NewIOContextFlags(IOContextFlagWrite))
 	require.NoError(t, err)
@@ -17,8 +44,7 @@ func TestIOContext(t *testing.T) {
 	require.Equal(t, "AVIOContext", cl.Name())
 	c.Write(nil)
 	c.Write([]byte("test"))
-	err = c.Closep()
-	require.NoError(t, err)
+	require.NoError(t, c.Free())
 	b, err := os.ReadFile(path)
 	require.NoError(t, err)
 	require.Equal(t, "test", string(b))