aria2's control file format implementation in Go.

aria2ctrlfile.go

// Package aria2ctrlfile implements [aria2]'s [control file format] (version 1).
// See https://github.com/aria2/aria2/blob/release-1.37.0/src/DefaultBtProgressInfoFile.cc
// for implementation details.
//
// [aria2]: https://github.com/aria2/aria2
// [control file format]: https://aria2.github.io/manual/en/html/technical-notes.html#control-file-aria2-format
package aria2ctrlfile

import (
	"encoding/binary"
	"errors"
	"fmt"
	"os"
)

type CtrlFile struct {
	file *os.File
}

type Data struct {
	InfoHashCheckEnabled bool
	InfoHash             []byte
	PieceLength          uint32
	TotalLength          uint64
	UploadLength         uint64
	Bitfield             []byte
	InFlightPieces       []Piece
}

type Piece struct {
	Index    uint32
	Length   uint32
	Bitfield []byte
}

var (
	ErrVersion = errors.New("unsupported control file version")
)

func New(file *os.File) *CtrlFile {
	return &CtrlFile{
		file: file,
	}
}

func (c *CtrlFile) Load() (*Data, error) {
	d := &Data{}

	var version uint16
	err := c.read(&version)
	if err != nil {
		return nil, fmt.Errorf("could not read version: %w", err)
	}
	if version != 1 {
		return nil, fmt.Errorf("%w (want %o, got %o)", ErrVersion, uint16(1), version)
	}

	extension := make([]byte, 4)
	err = c.read(extension)
	if err != nil {
		return nil, fmt.Errorf("could not read extension: %w", err)
	}
	d.InfoHashCheckEnabled = extension[3]&1 == 1

	var infoHashLen uint32
	err = c.read(&infoHashLen)
	if err != nil {
		return nil, fmt.Errorf("could not read info hash length: %w", err)
	}

	d.InfoHash = make([]byte, infoHashLen)
	err = c.read(d.InfoHash)
	if err != nil {
		return nil, fmt.Errorf("could not read info hash: %w", err)
	}

	err = c.read(&d.PieceLength)
	if err != nil {
		return nil, fmt.Errorf("could not read piece length: %w", err)
	}

	err = c.read(&d.TotalLength)
	if err != nil {
		return nil, fmt.Errorf("could not read total length: %w", err)
	}

	err = c.read(&d.UploadLength)
	if err != nil {
		return nil, fmt.Errorf("could not read upload length: %w", err)
	}

	var bitfieldLen uint32
	err = c.read(&bitfieldLen)
	if err != nil {
		return nil, fmt.Errorf("could not read bitfield length: %w", err)
	}

	d.Bitfield = make([]byte, bitfieldLen)
	err = c.read(d.Bitfield)
	if err != nil {
		return nil, fmt.Errorf("could not read bitfield: %w", err)
	}

	var inFlightPieceCount uint32
	err = c.read(&inFlightPieceCount)
	if err != nil {
		return nil, fmt.Errorf("could not read in-flight piece count: %w", err)
	}

	d.InFlightPieces = make([]Piece, inFlightPieceCount)

	for i := 0; i < int(inFlightPieceCount); i++ {
		err = c.read(&d.InFlightPieces[i].Index)
		if err != nil {
			return nil, fmt.Errorf("could not read the index of in-flight piece #%d: %w", i, err)
		}
		err = c.read(&d.InFlightPieces[i].Length)
		if err != nil {
			return nil, fmt.Errorf("could not read the length of in-flight piece #%d: %w", i, err)
		}
		var bitfieldLen uint32
		err = c.read(&bitfieldLen)
		if err != nil {
			return nil, fmt.Errorf("could not read the bitfield length of in-flight piece #%d: %w", i, err)
		}
		d.InFlightPieces[i].Bitfield = make([]byte, bitfieldLen)
		err = c.read(d.InFlightPieces[i].Bitfield)
		if err != nil {
			return nil, fmt.Errorf("could not read the bitfield of in-flight piece #%d: %w", i, err)
		}
	}

	return d, nil
}

func (c *CtrlFile) Save(d *Data) error {

	var version uint16 = 1
	err := c.write(version)
	if err != nil {
		return fmt.Errorf("could not write version: %w", err)
	}

	extension := make([]byte, 4)
	if d.InfoHashCheckEnabled {
		extension[3] = 1
	}
	err = c.write(extension)
	if err != nil {
		return fmt.Errorf("could not write extension: %w", err)
	}

	err = c.write(uint32(len(d.InfoHash)))
	if err != nil {
		return fmt.Errorf("could not write info hash length: %w", err)
	}

	err = c.write(d.InfoHash)
	if err != nil {
		return fmt.Errorf("could not write info hash: %w", err)
	}

	err = c.write(d.PieceLength)
	if err != nil {
		return fmt.Errorf("could not write piece length: %w", err)
	}

	err = c.write(d.TotalLength)
	if err != nil {
		return fmt.Errorf("could not write total length: %w", err)
	}

	err = c.write(d.UploadLength)
	if err != nil {
		return fmt.Errorf("could not write upload length: %w", err)
	}

	err = c.write(uint32(len(d.Bitfield)))
	if err != nil {
		return fmt.Errorf("could not write bitfield length: %w", err)
	}

	err = c.write(d.Bitfield)
	if err != nil {
		return fmt.Errorf("could not write bitfield: %w", err)
	}

	err = c.read(uint32(len(d.InFlightPieces)))
	if err != nil {
		return fmt.Errorf("could not write in-flight piece count: %w", err)
	}

	for i, piece := range d.InFlightPieces {
		err = c.write(piece.Index)
		if err != nil {
			return fmt.Errorf("could not write the index of in-flight piece #%d: %w", i, err)
		}
		err = c.write(piece.Length)
		if err != nil {
			return fmt.Errorf("could not write the length of in-flight piece #%d: %w", i, err)
		}
		err = c.write(uint32(len(piece.Bitfield)))
		if err != nil {
			return fmt.Errorf("could not write the bitfield length of in-flight piece #%d: %w", i, err)
		}
		err = c.write(piece.Bitfield)
		if err != nil {
			return fmt.Errorf("could not write the bitfield of in-flight piece #%d: %w", i, err)
		}
	}

	return nil
}

func (c *CtrlFile) read(data any) error {
	return binary.Read(c.file, binary.BigEndian, data)
}

func (c *CtrlFile) write(data any) error {
	return binary.Write(c.file, binary.BigEndian, data)
}