metainfo/file-tree.go

   1 package metainfo
   2
   3 import (
   4         "sort"
   5
   6         g "github.com/anacrolix/generics"
   7         "golang.org/x/exp/maps"
   8
   9         "github.com/anacrolix/torrent/bencode"
  10 )
  11
  12 const FileTreePropertiesKey = ""
  13
  14 type FileTreeFile struct {
  15         Length     int64  `bencode:"length"`
  16         PiecesRoot string `bencode:"pieces root"`
  17 }
  18
  19 // The fields here don't need bencode tags as the marshalling is done manually.
  20 type FileTree struct {
  21         File FileTreeFile
  22         Dir  map[string]FileTree
  23 }
  24
  25 func (ft *FileTree) UnmarshalBencode(bytes []byte) (err error) {
  26         var dir map[string]bencode.Bytes
  27         err = bencode.Unmarshal(bytes, &dir)
  28         if err != nil {
  29                 return
  30         }
  31         if propBytes, ok := dir[""]; ok {
  32                 err = bencode.Unmarshal(propBytes, &ft.File)
  33                 if err != nil {
  34                         return
  35                 }
  36         }
  37         delete(dir, "")
  38         g.MakeMapWithCap(&ft.Dir, len(dir))
  39         for key, bytes := range dir {
  40                 var sub FileTree
  41                 err = sub.UnmarshalBencode(bytes)
  42                 if err != nil {
  43                         return
  44                 }
  45                 ft.Dir[key] = sub
  46         }
  47         return
  48 }
  49
  50 var _ bencode.Unmarshaler = (*FileTree)(nil)
  51
  52 func (ft *FileTree) NumEntries() (num int) {
  53         num = len(ft.Dir)
  54         if g.MapContains(ft.Dir, FileTreePropertiesKey) {
  55                 num--
  56         }
  57         return
  58 }
  59
  60 func (ft *FileTree) IsDir() bool {
  61         return ft.NumEntries() != 0
  62 }
  63
  64 func (ft *FileTree) orderedKeys() []string {
  65         keys := maps.Keys(ft.Dir)
  66         sort.Strings(keys)
  67         return keys
  68 }
  69
  70 func (ft *FileTree) upvertedFiles(pieceLength int64, out func(fi FileInfo)) {
  71         var offset int64
  72         ft.upvertedFilesInner(pieceLength, nil, &offset, out)
  73 }
  74
  75 func (ft *FileTree) upvertedFilesInner(
  76         pieceLength int64,
  77         path []string,
  78         offset *int64,
  79         out func(fi FileInfo),
  80 ) {
  81         if ft.IsDir() {
  82                 for _, key := range ft.orderedKeys() {
  83                         if key == FileTreePropertiesKey {
  84                                 continue
  85                         }
  86                         sub := g.MapMustGet(ft.Dir, key)
  87                         sub.upvertedFilesInner(pieceLength, append(path, key), offset, out)
  88                 }
  89         } else {
  90                 out(FileInfo{
  91                         Length: ft.File.Length,
  92                         Path:   append([]string(nil), path...),
  93                         // BEP 52 requires paths be UTF-8 if possible.
  94                         PathUtf8:      append([]string(nil), path...),
  95                         PiecesRoot:    ft.PiecesRootAsByteArray(),
  96                         TorrentOffset: *offset,
  97                 })
  98                 *offset += (ft.File.Length + pieceLength - 1) / pieceLength * pieceLength
  99         }
 100 }
 101
 102 func (ft *FileTree) Walk(path []string, f func(path []string, ft *FileTree)) {
 103         f(path, ft)
 104         for key, sub := range ft.Dir {
 105                 if key == FileTreePropertiesKey {
 106                         continue
 107                 }
 108                 sub.Walk(append(path, key), f)
 109         }
 110 }
 111
 112 func (ft *FileTree) PiecesRootAsByteArray() (ret g.Option[[32]byte]) {
 113         if ft.File.PiecesRoot == "" {
 114                 return
 115         }
 116         n := copy(ret.Value[:], ft.File.PiecesRoot)
 117         if n != 32 {
 118                 // Must be 32 bytes for meta version 2 and non-empty files. See BEP 52.
 119                 panic(n)
 120         }
 121         ret.Ok = true
 122         return
 123 }