download_test.go

Documentation: github.com/concurrency-8/torrent

     1  package torrent
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"fmt"
     7  	"log"
     8  	"math"
     9  	"net"
    10  	"os"
    11  	"sync"
    12  	"testing"
    13  
    14  	"github.com/concurrency-8/parser"
    15  	"github.com/concurrency-8/piece"
    16  	"github.com/concurrency-8/queue"
    17  	"github.com/concurrency-8/tracker"
    18  	"github.com/stretchr/testify/assert"
    19  )
    20  
    21  func setLogs() {
    22  	Info = log.New(os.Stdout, "Testing ", 0)
    23  	Error = log.New(os.Stderr, "Testing ", 0)
    24  }
    25  
    26  /*
    27  // TestOnWholeMessage tests torrent/download.go : onWholeMessage(*kwargs)
    28  func TestOnWholeMessage(t *testing.T) {
    29  	setLogs()
    30  	fmt.Println("Testing torrent/download.go : onWholeMessage(*kwargs)")
    31  	num := 20
    32  	messages := make([][]byte, num)
    33  	wholeMessage := new(bytes.Buffer)
    34  	for i := 0; i < int(num); i++ {
    35  		rand.Seed(int64(i))
    36  		length := rand.Intn(256)
    37  		var message []byte
    38  		if i == 0 {
    39  			message = make([]byte, length+49)
    40  		} else {
    41  			message = make([]byte, length+4)
    42  		}
    43  
    44  		rand.Read(message)
    45  		message[0] = uint8(length)
    46  		binary.Write(wholeMessage, binary.BigEndian, message)
    47  		messages[i] = message
    48  	}
    49  
    50  	client, server := net.Pipe() // create a client and server connection
    51  
    52  	go func() {
    53  		server.Write(wholeMessage.Bytes())
    54  		server.Close() // close after writing out all data
    55  	}()
    56  
    57  	i := 0
    58  	_, err := onWholeMessage(tracker.Peer{}, client, func(peer tracker.Peer, b []byte,
    59  		client net.Conn,
    60  		pieces *piece.PieceTracker,
    61  		queue *queue.Queue,
    62  		report *tracker.ClientStatusReport) error { // mock Message Handler
    63  		if i == 0 {
    64  			assert.Equal(t, len(b), int(messages[i][0])+49, "length not equal")
    65  		} else {
    66  			assert.Equal(t, len(b), int(messages[i][0])+4, "length not equal")
    67  		}
    68  
    69  		assert.Equal(t, b, messages[i], "message received not same")
    70  		i++
    71  		return assert.AnError
    72  	}, piece.NewPieceTracker(tracker.GetRandomClientReport().TorrentFile), nil, nil) // TODO add tests for other message handlers
    73  
    74  	//exclude EOF errors, due to closing a connection.
    75  	assert.Equal(t, err, fmt.Errorf("EOF"), "Not EOF error")
    76  	assert.Equal(t, i, int(num), "Number of messages received are not equal")
    77  }
    78  
    79  /*
    80  func TestDownload(t *testing.T) {
    81  	fmt.Println("Testing torrent/download.go : Download()")
    82  	torrentfile, err := parser.ParseFromFile("../test_torrents/ubuntu.iso.torrent")
    83  	assert.Nil(t, err, "Opening torrent file failed.")
    84  	u, err := url.Parse(torrentfile.Announce[0])
    85  	assert.Nil(t, err, "Parsing announce URL failed")
    86  	statusreport := tracker.GetRandomClientReport()
    87  	resp, err := tracker.GetPeers(u, statusreport)
    88  	assert.Nil(t, err, "GetPeers returned error")
    89  	assert.NotEmpty(t, resp.Peers, "Empty Peer list")
    90  	for _, peer := range resp.Peers {
    91  		err = Download(peer, statusreport)
    92  		fmt.Println(err)
    93  		assert.Nil(t, err, "Download returned error", err)
    94  		if err == nil {
    95  			break
    96  		}
    97  	}
    98  }
    99  */
   100  
   101  // TestChokeHandler tests handling choking protocol - TODO : Adapt to the new function
   102  // The new function tries to handshake again, just in case the peer decides to unchoke
   103  // func TestChokeHandler(t *testing.T) {
   104  // 	client, _ := net.Pipe()
   105  
   106  // 	ChokeHandler(tracker.Peer{}, client, nil, nil)
   107  
   108  // 	_, err := client.Read(make([]byte, 4))
   109  // 	assert.Equal(t, err, fmt.Errorf("io: read/write on closed pipe"), "ChokeHandler failed")
   110  // }
   111  
   112  // TestUnchokeHandler tests handling of unchoking protocol
   113  func TestUnChokeHandler(t *testing.T) {
   114  	setLogs()
   115  	queue := queue.NewQueue(parser.TorrentFile{})
   116  	UnchokeHandler(tracker.Peer{}, nil, nil, queue)
   117  	assert.Equal(t, queue.Choked, false, "Choked attribute not set properly")
   118  }
   119  
   120  //TestRequestPiece tests function RequestPiece
   121  func TestRequestPiece(t *testing.T) {
   122  
   123  	file, _ := parser.ParseFromFile(parser.GetTorrentFileList()[0])
   124  	pieces := piece.NewPieceTracker(file)
   125  	queue := queue.NewQueue(file)
   126  	queue.Choked = false
   127  	pieceBlock := parser.RandomPieceBlock(file)
   128  	queue.Enqueue(pieceBlock.Index)
   129  	length := queue.Length()
   130  	client, server := net.Pipe()
   131  	fmt.Println(pieceBlock)
   132  	go func() {
   133  		for i := 0; i < length; i++ {
   134  			RequestPiece(tracker.Peer{}, server, pieces, queue)
   135  		}
   136  		defer server.Close()
   137  	}()
   138  
   139  	for i := 0; i < length; i++ {
   140  		resp := make([]byte, 17)
   141  		respLen, _ := client.Read(resp)
   142  		assert.Equal(t, respLen, 17, "Full message not received")
   143  		size, id, payload := ParseMsg(bytes.NewBuffer(resp))
   144  		assert.Equal(t, size, uint32(13), "Request: Size not equal")
   145  		assert.Equal(t, id, uint8(6), "Request: Message ID different")
   146  		assert.Equal(t, uint32(payload["index"].(uint32)), pieceBlock.Index, "Request: index field of payload not same")
   147  		assert.Equal(t, uint32(payload["begin"].(uint32)), uint32(i)*parser.BLOCK_LEN, "Request: begin field of payload not same")
   148  	}
   149  }
   150  
   151  func TestHaveHandler(t *testing.T) {
   152  	var flag sync.WaitGroup
   153  	flag.Add(1)
   154  	fmt.Println("Testing torrent/download.go : HaveHandler")
   155  	file, _ := parser.ParseFromFile(parser.GetTorrentFileList()[0])
   156  	pieces := piece.NewPieceTracker(file)
   157  	queue := queue.NewQueue(file)
   158  	queue.Choked = false
   159  	pieceBlock := parser.RandomPieceBlock(file)
   160  	client, server := net.Pipe()
   161  	actualsamplemsg, err := BuildHave(pieceBlock.Index)
   162  	assert.Nil(t, err, "error writing to Buffer in BuildHave")
   163  	go func() {
   164  		resp := make([]byte, 20)
   165  		_, err = server.Write(actualsamplemsg.Bytes())
   166  		flag.Wait()
   167  		respLen, err := server.Read(resp)
   168  		assert.Nil(t, err, "Error reading from server")
   169  		size, id, _ := ParseMsg(bytes.NewBuffer(resp[:respLen]))
   170  		assert.Equal(t, uint8(6), id, "Invalid id after reading from pipe.")
   171  		assert.Equal(t, uint32(13), size, "Invalid size")
   172  		defer server.Close()
   173  
   174  	}()
   175  	resp := make([]byte, 20)
   176  	buffer := new(bytes.Buffer)
   177  	respLen, err := client.Read(resp)
   178  	flag.Done()
   179  	assert.Nil(t, err, "Error reading from pipe")
   180  	err = binary.Write(buffer, binary.BigEndian, resp[:respLen])
   181  	assert.Nil(t, err, "Error writing to buffer.")
   182  	size, id, payload := ParseMsg(buffer)
   183  	assert.Equal(t, uint8(4), id, "Invalid id after reading from pipe.")
   184  	assert.Equal(t, uint32(5), size, "Invalid size")
   185  	assert.NotEmpty(t, payload["payload"], "Empty piece index in payload.")
   186  	var pieceIndex uint32
   187  	pieceIndex, err = HaveHandler(tracker.Peer{}, client, pieces, queue, payload)
   188  	assert.Nil(t, err, "Error in HaveHandler")
   189  	assert.Equal(t, pieceBlock.Index, pieceIndex, "Piece Index doesn't match.")
   190  	assert.True(t, pieces.Requested[pieceBlock.Index][0], "Requested not set.")
   191  }
   192  
   193  func TestBitFieldHandler(t *testing.T) {
   194  	var flag sync.WaitGroup
   195  	flag.Add(1)
   196  	fmt.Println("Testing torrent/download.go : BitFieldHandler")
   197  	file, _ := parser.ParseFromFile(parser.GetTorrentFileList()[0])
   198  	//each piece has 20 byte hash, irrespective of size.
   199  	npieces := uint32(len(file.Piece) / 20)
   200  	nbytes := uint(math.Ceil(float64(npieces) / float64(8)))
   201  	msg := new(bytes.Buffer)
   202  	binary.Write(msg, binary.BigEndian, uint32(nbytes+1))
   203  	binary.Write(msg, binary.BigEndian, uint8(5))
   204  	binary.Write(msg, binary.BigEndian, getRandomByteArr(nbytes))
   205  	actualmsg := msg.Bytes()
   206  	pieces := piece.NewPieceTracker(file)
   207  	queue := queue.NewQueue(file)
   208  	queue.Choked = false
   209  	client, server := net.Pipe()
   210  	go func() {
   211  		resp := make([]byte, nbytes+1)
   212  		_, err := server.Write(actualmsg)
   213  		flag.Wait()
   214  		assert.Nil(t, err, "Error writing to pipe.")
   215  		respLen, err := server.Read(resp)
   216  		assert.Nil(t, err, "Error reading from server")
   217  		size, id, _ := ParseMsg(bytes.NewBuffer(resp[:respLen]))
   218  		assert.Equal(t, uint8(6), id, "Invalid id after reading from pipe.")
   219  		assert.Equal(t, uint32(13), size, "Invalid size")
   220  		defer server.Close()
   221  
   222  	}()
   223  	resp := make([]byte, nbytes+10)
   224  	respLen, err := client.Read(resp)
   225  	flag.Done()
   226  	assert.Nil(t, err, "Error reading from Pipe")
   227  	buffer := new(bytes.Buffer)
   228  	err = binary.Write(buffer, binary.BigEndian, int32(nbytes+1))
   229  	assert.Nil(t, err, "Error writing to buffer.")
   230  	err = binary.Write(buffer, binary.BigEndian, int8(5))
   231  	assert.Nil(t, err, "Error writing to buffer.")
   232  	err = binary.Write(buffer, binary.BigEndian, resp[:respLen])
   233  	assert.Nil(t, err, "Error writing to buffer.")
   234  	size, id, payload := ParseMsg(buffer)
   235  	assert.Equal(t, uint8(5), id, "Invalid id after reading from Pipe")
   236  	assert.Equal(t, uint32(nbytes+1), size, "Invalid size")
   237  	assert.NotEmpty(t, payload["payload"], "Empty pieces in payload")
   238  	err = BitFieldHandler(tracker.Peer{}, client, pieces, queue, payload)
   239  	assert.Nil(t, err, "Error in BitFieldHandler")
   240  	// For each item in the queue, assert into the received field.
   241  	for i := 0; queue.Length() > 0; i++ {
   242  		nextitem, err := queue.Peek()
   243  		assert.Nil(t, err, "Error peeking into queue")
   244  		err = queue.Dequeue()
   245  		assert.Nil(t, err, "Error Dequeueing from queue")
   246  		index := nextitem.Index / 8
   247  		//index into the byte
   248  		offset := nextitem.Index % 8
   249  		//offset in the byte
   250  		p := uint8(1 << (7 - offset))
   251  		//p is used to perform bitwise and on the byte value.
   252  		assert.Equal(t, p, uint8(actualmsg[index])&p)
   253  	}
   254  
   255  }
   256
View as plain text