share/src/swarm/chunks.go

package swarm

import (
	"bytes"
	"io"
	"os"
	"sync"

	"github.com/tzgyn/share/src/utils"
)

// ChunkSize is fixed at 1MB: large enough to keep protocol overhead low,
// small enough that resuming a transfer never re-fetches much.
const ChunkSize = 1 << 20

// NumChunks returns how many ChunkSize pieces a file of size bytes has.
func NumChunks(size int64) int64 {
	if size <= 0 {
		return 0
	}
	return (size + ChunkSize - 1) / ChunkSize
}

func chunkLen(size, offset int64) int64 {
	if offset+ChunkSize > size {
		return size - offset
	}
	return ChunkSize
}

// ChunkStore tracks on-disk chunk presence for one file in the swarm and
// reads/writes it at chunk granularity.
type ChunkStore struct {
	mu   sync.Mutex
	path string
	size int64
	have map[int64]bool // chunk offset -> present
}

// OpenChunkStore opens (or creates) path as a file of exactly size bytes
// for receiving into, and figures out which chunks are already there. This
// reuses the same "a chunk-sized block that's still all zero bytes hasn't
// arrived yet" scan croc itself uses, so resuming an interrupted transfer
// needs no separate manifest of what's already done.
func OpenChunkStore(path string, size int64) (*ChunkStore, error) {
	info, statErr := os.Stat(path)
	fresh := statErr != nil || info.Size() != size

	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0o644)
	if err != nil {
		return nil, err
	}
	if fresh {
		err = f.Truncate(size)
	}
	f.Close()
	if err != nil {
		return nil, err
	}

	cs := &ChunkStore{path: path, size: size, have: make(map[int64]bool)}
	if !fresh {
		if err := cs.scanExisting(); err != nil {
			return nil, err
		}
	}
	return cs, nil
}

// NewSeedChunkStore wraps an already-complete local file being shared into
// the swarm: no scanning, every chunk is simply present. The file is never
// truncated or modified.
func NewSeedChunkStore(path string, size int64) *ChunkStore {
	cs := &ChunkStore{path: path, size: size, have: make(map[int64]bool)}
	for offset := int64(0); offset < size; offset += ChunkSize {
		cs.have[offset] = true
	}
	return cs
}

func (cs *ChunkStore) scanExisting() error {
	f, err := os.Open(cs.path)
	if err != nil {
		return err
	}
	defer f.Close()

	zero := make([]byte, ChunkSize)
	buf := make([]byte, ChunkSize)
	for offset := int64(0); offset < cs.size; offset += ChunkSize {
		l := chunkLen(cs.size, offset)
		if _, err := f.ReadAt(buf[:l], offset); err != nil && err != io.EOF {
			return err
		}
		if !bytes.Equal(buf[:l], zero[:l]) {
			cs.have[offset] = true
		}
	}
	return nil
}

func (cs *ChunkStore) HaveOffsets() []int64 {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	list := make([]int64, 0, len(cs.have))
	for offset := range cs.have {
		list = append(list, offset)
	}
	return list
}

func (cs *ChunkStore) HasChunk(offset int64) bool {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	return cs.have[offset]
}

func (cs *ChunkStore) Complete() bool {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	return int64(len(cs.have)) == NumChunks(cs.size)
}

func (cs *ChunkStore) ReadChunk(offset int64) ([]byte, error) {
	f, err := os.Open(cs.path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	buf := make([]byte, chunkLen(cs.size, offset))
	if _, err := f.ReadAt(buf, offset); err != nil && err != io.EOF {
		return nil, err
	}
	return buf, nil
}

func (cs *ChunkStore) WriteChunk(offset int64, data []byte) error {
	f, err := os.OpenFile(cs.path, os.O_WRONLY, 0o644)
	if err != nil {
		return err
	}
	defer f.Close()
	if _, err := f.WriteAt(data, offset); err != nil {
		return err
	}
	cs.mu.Lock()
	cs.have[offset] = true
	cs.mu.Unlock()
	return nil
}

// VerifyHash confirms a completed file matches the expected whole-file
// hash from the manifest.
func (cs *ChunkStore) VerifyHash(expected []byte) (bool, error) {
	got, err := utils.HashFile(cs.path, "xxhash")
	if err != nil {
		return false, err
	}
	return bytes.Equal(got, expected), nil
}