]> Sergey Matveev's repositories - btrtrc.git/blob - metainfo/info.go
Performance optimizations to calculating bytes left in Torrent
[btrtrc.git] / metainfo / info.go
1 package metainfo
2
3 import (
4         "crypto/sha1"
5         "errors"
6         "fmt"
7         "io"
8         "os"
9         "path/filepath"
10         "strings"
11
12         "github.com/anacrolix/missinggo/slices"
13 )
14
15 // The info dictionary.
16 type Info struct {
17         PieceLength int64  `bencode:"piece length"`
18         Pieces      []byte `bencode:"pieces"`
19         Name        string `bencode:"name"`
20         Length      int64  `bencode:"length,omitempty"`
21         Private     *bool  `bencode:"private,omitempty"`
22         // TODO: Document this field.
23         Source string     `bencode:"source,omitempty"`
24         Files  []FileInfo `bencode:"files,omitempty"`
25 }
26
27 // This is a helper that sets Files and Pieces from a root path and its
28 // children.
29 func (info *Info) BuildFromFilePath(root string) (err error) {
30         info.Name = filepath.Base(root)
31         info.Files = nil
32         err = filepath.Walk(root, func(path string, fi os.FileInfo, err error) error {
33                 if err != nil {
34                         return err
35                 }
36                 if fi.IsDir() {
37                         // Directories are implicit in torrent files.
38                         return nil
39                 } else if path == root {
40                         // The root is a file.
41                         info.Length = fi.Size()
42                         return nil
43                 }
44                 relPath, err := filepath.Rel(root, path)
45                 if err != nil {
46                         return fmt.Errorf("error getting relative path: %s", err)
47                 }
48                 info.Files = append(info.Files, FileInfo{
49                         Path:   strings.Split(relPath, string(filepath.Separator)),
50                         Length: fi.Size(),
51                 })
52                 return nil
53         })
54         if err != nil {
55                 return
56         }
57         slices.Sort(info.Files, func(l, r FileInfo) bool {
58                 return strings.Join(l.Path, "/") < strings.Join(r.Path, "/")
59         })
60         err = info.GeneratePieces(func(fi FileInfo) (io.ReadCloser, error) {
61                 return os.Open(filepath.Join(root, strings.Join(fi.Path, string(filepath.Separator))))
62         })
63         if err != nil {
64                 err = fmt.Errorf("error generating pieces: %s", err)
65         }
66         return
67 }
68
69 func (info *Info) writeFiles(w io.Writer, open func(fi FileInfo) (io.ReadCloser, error)) error {
70         for _, fi := range info.UpvertedFiles() {
71                 r, err := open(fi)
72                 if err != nil {
73                         return fmt.Errorf("error opening %v: %s", fi, err)
74                 }
75                 wn, err := io.CopyN(w, r, fi.Length)
76                 r.Close()
77                 if wn != fi.Length || err != nil {
78                         return fmt.Errorf("error hashing %v: %s", fi, err)
79                 }
80         }
81         return nil
82 }
83
84 // Set info.Pieces by hashing info.Files.
85 func (info *Info) GeneratePieces(open func(fi FileInfo) (io.ReadCloser, error)) error {
86         if info.PieceLength == 0 {
87                 return errors.New("piece length must be non-zero")
88         }
89         pr, pw := io.Pipe()
90         go func() {
91                 err := info.writeFiles(pw, open)
92                 pw.CloseWithError(err)
93         }()
94         defer pr.Close()
95         var pieces []byte
96         for {
97                 hasher := sha1.New()
98                 wn, err := io.CopyN(hasher, pr, info.PieceLength)
99                 if err == io.EOF {
100                         err = nil
101                 }
102                 if err != nil {
103                         return err
104                 }
105                 if wn == 0 {
106                         break
107                 }
108                 pieces = hasher.Sum(pieces)
109                 if wn < info.PieceLength {
110                         break
111                 }
112         }
113         info.Pieces = pieces
114         return nil
115 }
116
117 func (info *Info) TotalLength() (ret int64) {
118         if info.IsDir() {
119                 for _, fi := range info.Files {
120                         ret += fi.Length
121                 }
122         } else {
123                 ret = info.Length
124         }
125         return
126 }
127
128 func (info *Info) NumPieces() int {
129         return len(info.Pieces) / 20
130 }
131
132 func (info *Info) IsDir() bool {
133         return len(info.Files) != 0
134 }
135
136 // The files field, converted up from the old single-file in the parent info
137 // dict if necessary. This is a helper to avoid having to conditionally handle
138 // single and multi-file torrent infos.
139 func (info *Info) UpvertedFiles() []FileInfo {
140         if len(info.Files) == 0 {
141                 return []FileInfo{{
142                         Length: info.Length,
143                         // Callers should determine that Info.Name is the basename, and
144                         // thus a regular file.
145                         Path: nil,
146                 }}
147         }
148         return info.Files
149 }
150
151 func (info *Info) Piece(index int) Piece {
152         return Piece{info, index}
153 }