metadata insertion sample

insert_meta_sample.go

package main

import (
	"io"
	"os"

	"github.com/abema/go-mp4"
	"github.com/sunfish-shogi/bufseekio"
)

func main() {
	input, err := os.Open(os.Args[1])
	if err != nil {
		panic(err)
	}
	defer input.Close()
	output, err := os.Create(os.Args[2])
	if err != nil {
		panic(err)
	}
	defer output.Close()

	var ilstExists bool
	var mdatOffsetDiff int64
	var stcoOffsets []int64

	r := bufseekio.NewReadSeeker(input, 1024*1024, 3)
	w := mp4.NewWriter(output)
	_, err = mp4.ReadBoxStructure(r, func(h *mp4.ReadHandle) (interface{}, error) {
		switch h.BoxInfo.Type {
		// 1. moov, trak, mdia, minf, stbl, udta
		case mp4.BoxTypeMoov(),
			mp4.BoxTypeTrak(),
			mp4.BoxTypeMdia(),
			mp4.BoxTypeMinf(),
			mp4.BoxTypeStbl(),
			mp4.BoxTypeUdta(),
			mp4.BoxTypeMeta(),
			mp4.BoxTypeIlst():
			_, err := w.StartBox(&h.BoxInfo)
			if err != nil {
				return nil, err
			}
			if _, err := h.Expand(); err != nil {
				return nil, err
			}
			// 1-a. [only moov box] add udta box if not exists
			if h.BoxInfo.Type == mp4.BoxTypeMoov() && !ilstExists {
				path := []mp4.BoxType{mp4.BoxTypeUdta(), mp4.BoxTypeMeta()}
				for _, boxType := range path {
					if _, err := w.StartBox(&mp4.BoxInfo{Type: boxType}); err != nil {
						return nil, err
					}
				}
				ctx := h.BoxInfo.Context
				ctx.UnderUdta = true
				if _, err := w.StartBox(&mp4.BoxInfo{Type: mp4.BoxTypeHdlr()}); err != nil {
					return nil, err
				}
				hdlr := &mp4.Hdlr{
					HandlerType: [4]byte{'m', 'd', 'i', 'r'},
				}
				if _, err := mp4.Marshal(w, hdlr, ctx); err != nil {
					return nil, err
				}
				if _, err := w.EndBox(); err != nil {
					return nil, err
				}
				if _, err := w.StartBox(&mp4.BoxInfo{Type: mp4.BoxTypeIlst()}); err != nil {
					return nil, err
				}
				ctx.UnderIlst = true
				if err := addMeta(w, ctx); err != nil {
					return nil, err
				}
				if _, err := w.EndBox(); err != nil {
					return nil, err
				}
				for range path {
					if _, err := w.EndBox(); err != nil {
						return nil, err
					}
				}
			}
			// 1-b. [only ilst box] add metadatas
			if h.BoxInfo.Type == mp4.BoxTypeIlst() {
				ctx := h.BoxInfo.Context
				ctx.UnderIlst = true
				if err := addMeta(w, ctx); err != nil {
					return nil, err
				}
				ilstExists = true
			}
			if _, err = w.EndBox(); err != nil {
				return nil, err
			}
		// 2. otherwise
		default:
			// 2-a. [only stco box] keep offset
			if h.BoxInfo.Type == mp4.BoxTypeStco() {
				offset, _ := w.Seek(0, io.SeekCurrent)
				stcoOffsets = append(stcoOffsets, offset)
			}
			// 2-b. [only mdat box] keep difference of offsets
			if h.BoxInfo.Type == mp4.BoxTypeMdat() {
				iOffset := int64(h.BoxInfo.Offset)
				oOffset, _ := w.Seek(0, io.SeekCurrent)
				mdatOffsetDiff = oOffset - iOffset
			}
			// copy box without modification
			if err := w.CopyBox(r, &h.BoxInfo); err != nil {
				return nil, err
			}
		}

		return nil, nil
	})
	if err != nil {
		panic(err)
	}

	// if mdat box is moved, update stco box
	if mdatOffsetDiff != 0 {
		for _, stcoOffset := range stcoOffsets {
			// seek to stco box header
			if _, err := output.Seek(stcoOffset, io.SeekStart); err != nil {
				panic(err)
			}
			// read box header
			bi, err := mp4.ReadBoxInfo(output)
			if err != nil {
				panic(err)
			}
			// read stco box payload
			var stco mp4.Stco
			if _, err := mp4.Unmarshal(output, bi.Size-bi.HeaderSize, &stco, bi.Context); err != nil {
				panic(err)
			}
			// update chunk offsets
			for i := range stco.ChunkOffset {
				stco.ChunkOffset[i] += uint32(mdatOffsetDiff)
			}
			// seek to stco box payload
			if _, err := bi.SeekToPayload(output); err != nil {
				panic(err)
			}
			// write stco box payload
			if _, err := mp4.Marshal(output, &stco, bi.Context); err != nil {
				panic(err)
			}
		}
	}
}

func addMeta(w *mp4.Writer, ctx mp4.Context) error {
	metaBoxType := mp4.BoxType{0xA9, 'a', 'l', 'b'}
	if _, err := w.StartBox(&mp4.BoxInfo{Type: metaBoxType}); err != nil {
		return err
	}
	if _, err := w.StartBox(&mp4.BoxInfo{Type: mp4.BoxTypeData()}); err != nil {
		return err
	}
	var boxData = &mp4.Data{
		DataType: mp4.DataTypeStringUTF8,
		Data:     []byte("Hello, World!"),
	}
	dataCtx := ctx
	dataCtx.UnderIlstMeta = true
	if _, err := mp4.Marshal(w, boxData, dataCtx); err != nil {
		return err
	}
	if _, err := w.EndBox(); err != nil {
		return err
	}
	_, err := w.EndBox()
	return err
}