feat: saving functionality now provided by file

Author: raklaptudirm

cmd/mtor/main.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"time"
 
+	"github.com/raklaptudirm/mtor/internal/build"
 	"github.com/raklaptudirm/mtor/pkg/file"
 	"github.com/raklaptudirm/mtor/pkg/torrent"
 )
@@ -43,51 +44,22 @@ func main() {
 
 	fmt.Printf("torrent %x - %d pieces\n", t.InfoHash, len(t.PieceHashes))
 
-	ps := Pieces{
-		size:   t.PieceLength,
-		buffer: make([]byte, t.Length),
-	}
-	err = t.Download(&ps, config)
+	ps := build.PieceManager
+	err = ps.Init()
 	if err != nil {
 		fmt.Println(err)
-		return
 	}
+	defer ps.Close()
 
-	w, err := os.Create(string(f.Info.Name))
+	err = t.DownloadPieces(ps, config)
 	if err != nil {
 		fmt.Println(err)
 		return
 	}
-	defer w.Close()
-
-	for i := range t.PieceHashes {
-		block, err := ps.Get(i)
-		if err != nil {
-			fmt.Println(err)
-			return
-		}
 
-		_, err = w.Write(block)
-		if err != nil {
-			fmt.Println(err)
-			return
-		}
+	err = f.Save(ps, ".") // save in cwd
+	if err != nil {
+		fmt.Println(err)
+		return
 	}
 }
-
-// structure to put stuff together
-type Pieces struct {
-	size   int
-	length int
-	buffer []byte
-}
-
-func (p *Pieces) Put(index int, block []byte) error {
-	p.length++
-	copy(p.buffer[index*p.size:], block)
-	return nil
-}
-
-func (p *Pieces) Get(index int) ([]byte, error) {
-	return p.buffer[index*p.size : (index+1)*p.size], nil
-}

internal/build/env.go

@@ -0,0 +1,5 @@
+package build
+
+import "github.com/raklaptudirm/mtor/internal/manager"
+
+var PieceManager = manager.New()

internal/manager/piece.go

@@ -0,0 +1,87 @@
+// Copyright © 2021 Rak Laptudirm <raklaptudirm@gmail.com>
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package manager
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"path"
+)
+
+// piece represents the piece manager.
+type piece struct {
+	src string // storage directory
+}
+
+// ErrManagerClosed is returned when the manager is not initialized,
+// or closed.
+var ErrManagerClosed = errors.New("the manager is closed")
+
+// Init initializes the manager.
+func (p *piece) Init() error {
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return err
+	}
+
+	// create storage directory
+	dir, err := os.MkdirTemp(home, "mtor pieces ")
+	if err != nil {
+		return err
+	}
+
+	p.src = dir
+	return nil
+}
+
+// Put stores a piece in the manager.
+func (p *piece) Put(i int, buf []byte) error {
+	if p.isClosed() {
+		return ErrManagerClosed
+	}
+
+	file := path.Join(p.src, fmt.Sprintf("%x", i))
+	return os.WriteFile(file, buf, 0600)
+}
+
+// Get fetches a piece from the manager.
+func (p *piece) Get(i int) ([]byte, error) {
+	if p.isClosed() {
+		return nil, ErrManagerClosed
+	}
+
+	file := path.Join(p.src, fmt.Sprintf("%x", i))
+	return os.ReadFile(file)
+}
+
+// Close closes the manager.
+func (p *piece) Close() error {
+	if p.isClosed() {
+		return ErrManagerClosed
+	}
+
+	// free space
+	return os.RemoveAll(p.src)
+}
+
+// isClosed checks if the manager is closed.
+func (p *piece) isClosed() bool {
+	return p.src == ""
+}
+
+// New returns a new and un-initialzed instance of the manager.
+func New() *piece {
+	return &piece{}
+}

pkg/file/file.go

@@ -19,6 +19,8 @@ import (
 	"fmt"
 	"io"
 	"math/rand"
+	"os"
+	"path"
 	"time"
 
 	"github.com/jackpal/bencode-go"
@@ -59,6 +61,44 @@ type singleFile struct {
 	Path   []string `bencode:"path"`   // path of the file
 }
 
+// Save saves the torrent as a file or directory, fetching pieces from the
+// provided piece manager.
+func (f *file) Save(pieces torrent.PieceManager, dst string) error {
+	if f.isSingleFile() {
+		return f.saveSingleFile(pieces, dst)
+	}
+
+	// TODO: add support for saving multifile torrents
+	return nil
+}
+
+// saveSingleFile saves a single-file torrent as a file, fetching the pieces
+// from the provided piece manager.
+func (f *file) saveSingleFile(pieces torrent.PieceManager, dst string) error {
+	length := len(f.Info.Pieces) / 20 // each hash is 20 bytes
+
+	file, err := os.Create(path.Join(dst, f.Info.Name))
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	// get each piece
+	for i := 0; i < length; i++ {
+		piece, err := pieces.Get(i)
+		if err != nil {
+			return err
+		}
+
+		// write each piece to the file
+		_, err = file.Write(piece)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // Torrent converts a file into a torrent.Torrent.
 func (f *file) Torrent() (*torrent.Torrent, error) {
 	hash, err := f.Info.hash()

pkg/torrent/download.go

@@ -286,9 +286,9 @@ func (t *Torrent) pieceLen(index int) int {
 	return t.PieceLength
 }
 
-// Download downloads the t torrent and stores the downloaded pieces into
-// the provided PieceManager.
-func (t *Torrent) Download(p PieceManager, c *DownloadConfig) error {
+// DownloadPieces downloads the pieces of the provided torrent and stores
+// them into the provided PieceManager.
+func (t *Torrent) DownloadPieces(p PieceManager, c *DownloadConfig) error {
 	start := time.Now()
 
 	err := t.newDownload(p, c).start()

pkg/torrent/piece.go

@@ -87,8 +87,12 @@ func (p *pieceProgress) readMessage() error {
 // PieceManager represents an interface which can handle the storage of the
 // torrent's pieces.
 type PieceManager interface {
+	// Init initializes the manager to start storing pieces.
+	Init() error
 	// Put stores the buffer data with the provided piece index.
 	Put(int, []byte) error
 	// Get gets the data of the provided piece index.
 	Get(int) ([]byte, error)
+	// Close destroy's the manager's data. Call this when done.
+	Close() error
 }