github.com/anacrolix/torrent@v1.61.0/metainfo/file-tree.go (about)

     1  package metainfo
     2  
     3  import (
     4  	"iter"
     5  	"maps"
     6  	"slices"
     7  
     8  	g "github.com/anacrolix/generics"
     9  
    10  	"github.com/anacrolix/torrent/bencode"
    11  	infohash_v2 "github.com/anacrolix/torrent/types/infohash-v2"
    12  )
    13  
    14  const FileTreePropertiesKey = ""
    15  
    16  type FileTreeFile struct {
    17  	Length     int64  `bencode:"length"`
    18  	PiecesRoot string `bencode:"pieces root"`
    19  }
    20  
    21  // The fields here don't need bencode tags as the marshalling is done manually.
    22  type FileTree struct {
    23  	File FileTreeFile
    24  	Dir  map[string]FileTree
    25  }
    26  
    27  func (ft *FileTree) UnmarshalBencode(bytes []byte) (err error) {
    28  	var dir map[string]bencode.Bytes
    29  	err = bencode.Unmarshal(bytes, &dir)
    30  	if err != nil {
    31  		return
    32  	}
    33  	if propBytes, ok := dir[""]; ok {
    34  		err = bencode.Unmarshal(propBytes, &ft.File)
    35  		if err != nil {
    36  			return
    37  		}
    38  	}
    39  	delete(dir, "")
    40  	g.MakeMapWithCap(&ft.Dir, len(dir))
    41  	for key, bytes := range dir {
    42  		var sub FileTree
    43  		err = sub.UnmarshalBencode(bytes)
    44  		if err != nil {
    45  			return
    46  		}
    47  		ft.Dir[key] = sub
    48  	}
    49  	return
    50  }
    51  
    52  var _ bencode.Unmarshaler = (*FileTree)(nil)
    53  
    54  func (ft *FileTree) MarshalBencode() (bytes []byte, err error) {
    55  	if ft.IsDir() {
    56  		dir := make(map[string]bencode.Bytes, len(ft.Dir))
    57  		for _, key := range ft.orderedKeys() {
    58  			if key == FileTreePropertiesKey {
    59  				continue
    60  			}
    61  			sub := g.MapMustGet(ft.Dir, key)
    62  			subBytes, err := sub.MarshalBencode()
    63  			if err != nil {
    64  				return nil, err
    65  			}
    66  			dir[key] = subBytes
    67  		}
    68  		return bencode.Marshal(dir)
    69  	} else {
    70  		fileBytes, err := bencode.Marshal(ft.File)
    71  		if err != nil {
    72  			return nil, err
    73  		}
    74  		res := map[string]bencode.Bytes{
    75  			"": fileBytes,
    76  		}
    77  		return bencode.Marshal(res)
    78  	}
    79  }
    80  
    81  var _ bencode.Marshaler = (*FileTree)(nil)
    82  
    83  func (ft *FileTree) NumEntries() (num int) {
    84  	num = len(ft.Dir)
    85  	if g.MapContains(ft.Dir, FileTreePropertiesKey) {
    86  		num--
    87  	}
    88  	return
    89  }
    90  
    91  func (ft *FileTree) IsDir() bool {
    92  	return ft.NumEntries() != 0
    93  }
    94  
    95  func (ft *FileTree) orderedKeys() []string {
    96  	return slices.Sorted(maps.Keys(ft.Dir))
    97  }
    98  
    99  func (ft *FileTree) upvertedFiles(pieceLength int64) iter.Seq[FileInfo] {
   100  	var offset int64
   101  	return ft.upvertedFilesInner(pieceLength, nil, &offset)
   102  }
   103  
   104  func (ft *FileTree) upvertedFilesInner(
   105  	pieceLength int64,
   106  	path []string,
   107  	offset *int64,
   108  ) iter.Seq[FileInfo] {
   109  	return func(yield func(FileInfo) bool) {
   110  		if ft.IsDir() {
   111  			for _, key := range ft.orderedKeys() {
   112  				if key == FileTreePropertiesKey {
   113  					continue
   114  				}
   115  				sub := g.MapMustGet(ft.Dir, key)
   116  				for fi := range sub.upvertedFilesInner(pieceLength, append(path, key), offset) {
   117  					if !yield(fi) {
   118  						return
   119  					}
   120  				}
   121  			}
   122  		} else {
   123  			yield(FileInfo{
   124  				Length: ft.File.Length,
   125  				Path:   append([]string(nil), path...),
   126  				// BEP 52 requires paths be UTF-8 if possible.
   127  				PathUtf8:      append([]string(nil), path...),
   128  				PiecesRoot:    ft.PiecesRootAsByteArray(),
   129  				TorrentOffset: *offset,
   130  			})
   131  			// v2 files are piece aligned. This bumps up the offset to the next piece boundary.
   132  			*offset += (ft.File.Length + pieceLength - 1) / pieceLength * pieceLength
   133  		}
   134  
   135  	}
   136  }
   137  
   138  func (ft *FileTree) Walk(path []string, f func(path []string, ft *FileTree)) {
   139  	f(path, ft)
   140  	for key, sub := range ft.Dir {
   141  		if key == FileTreePropertiesKey {
   142  			continue
   143  		}
   144  		sub.Walk(append(path, key), f)
   145  	}
   146  }
   147  
   148  func (ft *FileTree) PiecesRootAsByteArray() (ret g.Option[infohash_v2.T]) {
   149  	if ft.File.PiecesRoot == "" {
   150  		return
   151  	}
   152  	n := copy(ret.Value[:], ft.File.PiecesRoot)
   153  	if n != 32 {
   154  		// Must be 32 bytes for meta version 2 and non-empty files. See BEP 52.
   155  		panic(n)
   156  	}
   157  	ret.Ok = true
   158  	return
   159  }