tstor/daemons/torrent/dup_cache.go

package torrent

import (
	"path"
	"slices"
	"sync"

	"git.kmsign.ru/royalcat/tstor/pkg/slicesutils"
	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/types/infohash"
)

type dupInfo struct {
	infohash infohash.T
	fileinfo metainfo.FileInfo
}

type dupIndex struct {
	mu         sync.RWMutex
	torrents   map[infohash.T][]metainfo.FileInfo
	sha1       map[string][]dupInfo   // bittorrent v1
	piecesRoot map[[32]byte][]dupInfo // bittorrent v2
}

func newDupIndex() *dupIndex {
	return &dupIndex{
		torrents:   map[infohash.T][]metainfo.FileInfo{},
		sha1:       map[string][]dupInfo{},
		piecesRoot: map[[32]byte][]dupInfo{},
	}
}

func (c *dupIndex) AddFile(fileinfo metainfo.FileInfo, ih infohash.T) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.torrents[ih] = append(c.torrents[ih], fileinfo)

	if fileinfo.Sha1 != "" {
		c.sha1[fileinfo.Sha1] = append(c.sha1[fileinfo.Sha1], dupInfo{fileinfo: fileinfo, infohash: ih})
	}

	if fileinfo.PiecesRoot.Ok {
		c.piecesRoot[fileinfo.PiecesRoot.Value] = append(c.piecesRoot[fileinfo.PiecesRoot.Value], dupInfo{fileinfo: fileinfo, infohash: ih})
	}
}

func (c *dupIndex) DuplicateFiles(fileinfo metainfo.FileInfo, ih infohash.T) []dupInfo {
	c.mu.RLock()
	defer c.mu.RUnlock()

	if fileinfo.Sha1 != "" {
		if dups, ok := c.sha1[fileinfo.Sha1]; ok {
			return slices.Clone(dups)
		}
	}

	if fileinfo.PiecesRoot.Ok {
		if dups, ok := c.piecesRoot[fileinfo.PiecesRoot.Value]; ok {
			return slices.Clone(dups)
		}
	}

	return []dupInfo{}
}

func (c *dupIndex) Includes(ih infohash.T, files []metainfo.FileInfo) []dupInfo {
	c.mu.RLock()
	defer c.mu.RUnlock()

	out := []dupInfo{}

	for ih, v := range c.torrents {
		intersection := slicesutils.IntersectionFunc(files, v, func(a, b metainfo.FileInfo) bool {
			mostly := path.Join(a.BestPath()...) == path.Join(b.BestPath()...) && a.Length == b.Length
			if a.Sha1 != "" && b.Sha1 != "" {
				return mostly && a.Sha1 == b.Sha1
			}
			if a.PiecesRoot.Ok && b.PiecesRoot.Ok {
				return mostly && a.PiecesRoot.Value == b.PiecesRoot.Value
			}

			return mostly
		})

		for _, v := range intersection {
			out = append(out, dupInfo{infohash: ih, fileinfo: v})
		}
	}

	return []dupInfo{}
}