message_utils_test.go

Documentation: github.com/concurrency-8/torrent

     1  package torrent
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"fmt"
     7  	"math/rand"
     8  	"os"
     9  	"testing"
    10  
    11  	"github.com/concurrency-8/parser"
    12  	"github.com/concurrency-8/tracker"
    13  	"github.com/stretchr/testify/assert"
    14  )
    15  
    16  func getTorrentFileList() []string {
    17  	return []string{"../ubuntu.iso.torrent", "../big-buck-bunny.torrent"}
    18  }
    19  
    20  func getRandomClientReport() (report *tracker.ClientStatusReport) {
    21  	torrent, _ := parser.ParseFromFile(getTorrentFileList()[0])
    22  	report = &tracker.ClientStatusReport{}
    23  	report.TorrentFile = torrent
    24  	report.PeerID = string(getRandomByteArr(20))
    25  	report.Left = torrent.Length
    26  	report.Port = uint16(6464)
    27  	report.Event = ""
    28  	return report
    29  }
    30  
    31  func getRandomByteArr(size uint) []byte {
    32  	temp := make([]byte, size)
    33  	_, err := rand.Read(temp)
    34  	if err != nil {
    35  		fmt.Fprintf(os.Stderr, "Unable to generate Crypto random byte array")
    36  	}
    37  	return temp
    38  }
    39  
    40  func TestBuildHandshake(t *testing.T) {
    41  	assert := assert.New(t)
    42  	csr := getRandomClientReport()
    43  	handshake, err := BuildHandshake(*csr)
    44  
    45  	assert.Nil(err)
    46  
    47  	handshakeReader := bytes.NewReader(handshake.Bytes())
    48  
    49  	var pstrlen uint8
    50  	assert.Nil(binary.Read(handshakeReader, binary.BigEndian, &pstrlen))
    51  
    52  	// Checks for the length and hence confirming that the pstrlen and pstr part are in sync
    53  	assert.Equal(int(49+pstrlen), len(handshake.Bytes()))
    54  
    55  	pstr := make([]byte, pstrlen)
    56  	assert.Nil(binary.Read(handshakeReader, binary.BigEndian, &pstr))
    57  
    58  	var reserved [8]byte
    59  	assert.Nil(binary.Read(handshakeReader, binary.BigEndian, &reserved))
    60  
    61  	var infohashFromBuf, infohashFromFile [20]byte
    62  	assert.Nil(binary.Read(handshakeReader, binary.BigEndian, &infohashFromBuf))
    63  	copy(infohashFromFile[:], csr.TorrentFile.InfoHash)
    64  	assert.Equal(infohashFromFile, infohashFromBuf)
    65  
    66  	var peerID [20]byte
    67  	assert.Nil(binary.Read(handshakeReader, binary.BigEndian, &peerID))
    68  	assert.Equal(csr.PeerID, string(peerID[:]))
    69  }
    70  
    71  func TestBuildKeepAlive(t *testing.T) {
    72  	alive := BuildKeepAlive()
    73  	assert.Equal(t, *bytes.NewBuffer(make([]byte, 4)), *alive)
    74  }
    75  
    76  func TestBuildChoke(t *testing.T) {
    77  	choke, err := BuildChoke()
    78  
    79  	assert.Nil(t, err)
    80  
    81  	// Length of choke buffer
    82  	assert.Equal(t, 5, len(choke.Bytes()))
    83  
    84  	chokeReader := bytes.NewReader(choke.Bytes())
    85  
    86  	// Read length
    87  	var length uint32
    88  	assert.Nil(t, binary.Read(chokeReader, binary.BigEndian, &length))
    89  	assert.Equal(t, uint32(1), length)
    90  
    91  	// Read message type (choke)
    92  	var messageType uint8
    93  	assert.Nil(t, binary.Read(chokeReader, binary.BigEndian, &messageType))
    94  	assert.Equal(t, uint8(0), messageType)
    95  }
    96  
    97  func TestBuildUnchoke(t *testing.T) {
    98  	unchoke, err := BuildUnchoke()
    99  
   100  	assert.Nil(t, err)
   101  
   102  	// Length of unchoke buffer
   103  	assert.Equal(t, 5, len(unchoke.Bytes()))
   104  
   105  	unchokeReader := bytes.NewReader(unchoke.Bytes())
   106  
   107  	// Read length
   108  	var length uint32
   109  	assert.Nil(t, binary.Read(unchokeReader, binary.BigEndian, &length))
   110  	assert.Equal(t, uint32(1), length)
   111  
   112  	// Read message type (unchoke)
   113  	var messageType uint8
   114  	assert.Nil(t, binary.Read(unchokeReader, binary.BigEndian, &messageType))
   115  	assert.Equal(t, uint8(1), messageType)
   116  }
   117  
   118  func TestBuildInterested(t *testing.T) {
   119  
   120  	assert := assert.New(t)
   121  
   122  	interested, err := BuildInterested()
   123  
   124  	assert.Nil(err)
   125  	// Length of interested buffer
   126  	assert.Equal(5, len(interested.Bytes()))
   127  
   128  	interestedReader := bytes.NewReader(interested.Bytes())
   129  
   130  	// Read length
   131  	var length uint32
   132  	assert.Nil(binary.Read(interestedReader, binary.BigEndian, &length))
   133  	assert.Equal(uint32(1), length)
   134  
   135  	// Read message type (interested)
   136  	var messageType uint8
   137  	assert.Nil(binary.Read(interestedReader, binary.BigEndian, &messageType))
   138  	assert.Equal(uint8(2), messageType)
   139  
   140  }
   141  
   142  func TestBuildUninterested(t *testing.T) {
   143  	assert := assert.New(t)
   144  
   145  	uninterested, err := BuildUninterested()
   146  
   147  	assert.Nil(err)
   148  	// Length of uninterested buffer
   149  	assert.Equal(5, len(uninterested.Bytes()))
   150  
   151  	uninterestedReader := bytes.NewReader(uninterested.Bytes())
   152  
   153  	// Read length
   154  	var length uint32
   155  	assert.Nil(binary.Read(uninterestedReader, binary.BigEndian, &length))
   156  	assert.Equal(uint32(1), length)
   157  
   158  	// Read message type (uninterested)
   159  	var messageType uint8
   160  	assert.Nil(binary.Read(uninterestedReader, binary.BigEndian, &messageType))
   161  	assert.Equal(uint8(3), messageType)
   162  }
   163  
   164  func TestBuildHave(t *testing.T) {
   165  	assert := assert.New(t)
   166  
   167  	payload := uint32(rand.Uint32())
   168  	have, err := BuildHave(payload)
   169  
   170  	assert.Nil(err)
   171  	// Length of have buffer
   172  	assert.Equal(9, len(have.Bytes()))
   173  
   174  	haveReader := bytes.NewReader(have.Bytes())
   175  
   176  	// Read length
   177  	var length uint32
   178  	assert.Nil(binary.Read(haveReader, binary.BigEndian, &length))
   179  	assert.Equal(uint32(5), length)
   180  
   181  	// Read message type (have)
   182  	var messageType uint8
   183  	assert.Nil(binary.Read(haveReader, binary.BigEndian, &messageType))
   184  	assert.Equal(uint8(4), messageType)
   185  
   186  	// Read payload
   187  	var payloadRead uint32
   188  	assert.Nil(binary.Read(haveReader, binary.BigEndian, &payloadRead))
   189  	assert.Equal(payload, payloadRead)
   190  }
   191  
   192  func TestBuildRequest(t *testing.T) {
   193  	assert := assert.New(t)
   194  
   195  	piece := GetRandomPiece()
   196  
   197  	request, err := BuildRequest(piece)
   198  
   199  	assert.Nil(err)
   200  
   201  	// Length of request buffer
   202  	assert.Equal(17, len(request.Bytes()))
   203  
   204  	requestReader := bytes.NewReader(request.Bytes())
   205  
   206  	// Read length of message
   207  	var length uint32
   208  	assert.Nil(binary.Read(requestReader, binary.BigEndian, &length))
   209  	assert.Equal(uint32(13), length)
   210  
   211  	// Read message type
   212  	var messageType uint8
   213  	assert.Nil(binary.Read(requestReader, binary.BigEndian, &messageType))
   214  	assert.Equal(uint8(6), messageType)
   215  
   216  	// Read piece index
   217  	var pieceIndex uint32
   218  	assert.Nil(binary.Read(requestReader, binary.BigEndian, &pieceIndex))
   219  	assert.Equal(piece.Index, pieceIndex)
   220  
   221  	// Read piece begin point
   222  	var pieceBegin uint32
   223  	assert.Nil(binary.Read(requestReader, binary.BigEndian, &pieceBegin))
   224  	assert.Equal(piece.Begin, pieceBegin)
   225  
   226  	// Read piece length
   227  	var pieceLength uint32
   228  	assert.Nil(binary.Read(requestReader, binary.BigEndian, &pieceLength))
   229  	assert.Equal(piece.Length, pieceLength)
   230  }
   231  
   232  // func TestBuildPiece(t *testing.T) {
   233  // 	assert := assert.New(t)
   234  
   235  // 	samplePiece := GetRandomPiece()
   236  
   237  // 	builtPiece, err := BuildPiece(samplePiece)
   238  
   239  // 	assert.Nil(err)
   240  
   241  // 	// Length of piece buffer
   242  // 	assert.Equal(len(samplePiece.Block.Bytes())+13, len(builtPiece.Bytes()))
   243  
   244  // 	pieceReader := bytes.NewReader(builtPiece.Bytes())
   245  
   246  // 	// Read length of message
   247  // 	var length uint32
   248  // 	assert.Nil(binary.Read(pieceReader, binary.BigEndian, &length))
   249  // 	assert.Equal(uint32(len(samplePiece.Block.Bytes())+9), length)
   250  
   251  // 	// Read message type
   252  // 	var messageType uint8
   253  // 	assert.Nil(binary.Read(pieceReader, binary.BigEndian, &messageType))
   254  // 	assert.Equal(uint8(7), messageType)
   255  
   256  // 	// Read piece index
   257  // 	var pieceIndexRead uint32
   258  // 	assert.Nil(binary.Read(pieceReader, binary.BigEndian, &pieceIndexRead))
   259  // 	assert.Equal(samplePiece.Index, pieceIndexRead)
   260  
   261  // 	// Read Begin
   262  // 	var pieceBeginRead uint32
   263  // 	assert.Nil(binary.Read(pieceReader, binary.BigEndian, &pieceBeginRead))
   264  // 	assert.Equal(samplePiece.Begin, pieceBeginRead)
   265  
   266  // 	// Read Buffer
   267  // 	blockRead := make([]byte, len(samplePiece.Block.Bytes()))
   268  // 	assert.Nil(binary.Read(pieceReader, binary.BigEndian, &blockRead))
   269  // 	assert.Equal(samplePiece.Block.Bytes(), blockRead)
   270  // }
   271  
   272  func TestBuildCancel(t *testing.T) {
   273  	assert := assert.New(t)
   274  
   275  	samplePiece := GetRandomPiece()
   276  
   277  	cancel, err := BuildCancel(samplePiece)
   278  
   279  	assert.Nil(err)
   280  
   281  	// Length of cancel buffer
   282  	assert.Equal(17, len(cancel.Bytes()))
   283  
   284  	cancelReader := bytes.NewReader(cancel.Bytes())
   285  
   286  	// Read length of message
   287  	var length uint32
   288  	assert.Nil(binary.Read(cancelReader, binary.BigEndian, &length))
   289  	assert.Equal(uint32(13), length)
   290  
   291  	// Message type = 8
   292  	var messageType uint8
   293  	assert.Nil(binary.Read(cancelReader, binary.BigEndian, &messageType))
   294  	assert.Equal(uint8(8), messageType)
   295  
   296  	// piece index
   297  	var pieceIndex uint32
   298  	assert.Nil(binary.Read(cancelReader, binary.BigEndian, &pieceIndex))
   299  	assert.Equal(samplePiece.Index, pieceIndex)
   300  
   301  	// piece begin
   302  	var pieceBegin uint32
   303  	assert.Nil(binary.Read(cancelReader, binary.BigEndian, &pieceBegin))
   304  	assert.Equal(samplePiece.Begin, pieceBegin)
   305  
   306  	// piece length
   307  	var pieceLength uint32
   308  	assert.Nil(binary.Read(cancelReader, binary.BigEndian, &pieceLength))
   309  	assert.Equal(samplePiece.Length, pieceLength)
   310  }
   311  
   312  func TestBuildPort(t *testing.T) {
   313  	assert := assert.New(t)
   314  	port := uint16(rand.Intn(90000) + 10000)
   315  
   316  	portBuf, err := BuildPort(port)
   317  	assert.Nil(err)
   318  
   319  	// Length of port buffer
   320  	assert.Equal(7, len(portBuf.Bytes()))
   321  
   322  	portBufReader := bytes.NewReader(portBuf.Bytes())
   323  
   324  	// Length of message
   325  	var length uint32
   326  	assert.Nil(binary.Read(portBufReader, binary.BigEndian, &length))
   327  	assert.Equal(uint32(3), length)
   328  
   329  	// Message type = 9
   330  	var messageType uint8
   331  	assert.Nil(binary.Read(portBufReader, binary.BigEndian, &messageType))
   332  	assert.Equal(uint8(9), messageType)
   333  
   334  	// port
   335  	var portReadFromBuf uint16
   336  	assert.Nil(binary.Read(portBufReader, binary.BigEndian, &portReadFromBuf))
   337  	assert.Equal(port, portReadFromBuf)
   338  }
   339  
   340  // ParseMsg test message parser
   341  func TestParseMsg(t *testing.T) {
   342  
   343  	// BuildChoke
   344  	chokeMessage, _ := BuildChoke()
   345  	size, id, payload := ParseMsg(chokeMessage)
   346  	assert.Equal(t, size, uint32(1), "choke : Size not equal")
   347  	assert.Equal(t, id, uint8(0), "choke : Message ID different")
   348  	assert.Equal(t, len(payload), 0, "choke : length of payload not zero")
   349  
   350  	//BuildUnchoke
   351  	unchokeMessage, _ := BuildUnchoke()
   352  	size, id, payload = ParseMsg(unchokeMessage)
   353  	assert.Equal(t, size, uint32(1), "unchoke : Size not equal")
   354  	assert.Equal(t, id, uint8(1), "unchoke : Message ID different")
   355  	assert.Equal(t, len(payload), 0, "unchoke : length of payload not zero")
   356  
   357  	// BuildInterested
   358  	interestedMessage, _ := BuildInterested()
   359  	size, id, payload = ParseMsg(interestedMessage)
   360  	assert.Equal(t, size, uint32(1), "Interested : Size not equal")
   361  	assert.Equal(t, id, uint8(2), "Interested : Message ID different")
   362  	assert.Equal(t, len(payload), 0, "Interested: length of payload not zero")
   363  
   364  	// BuildUninerested
   365  	uninterestedMessage, _ := BuildUninterested()
   366  	size, id, payload = ParseMsg(uninterestedMessage)
   367  	assert.Equal(t, size, uint32(1), "UnInterested: Size not equal")
   368  	assert.Equal(t, id, uint8(3), "Unterested: Message ID different")
   369  	assert.Equal(t, len(payload), 0, "UnInterested: length of payload not zero")
   370  
   371  	// BuildHave
   372  	p1 := rand.Uint32()
   373  	haveMessage, _ := BuildHave(p1)
   374  	size, id, payload = ParseMsg(haveMessage)
   375  	assert.Equal(t, size, uint32(5), "Have: Size not equal")
   376  	assert.Equal(t, id, uint8(4), "Have: Message ID different")
   377  	var p2 uint32
   378  	binary.Read(payload["payload"].(*bytes.Buffer), binary.BigEndian, &p2)
   379  	assert.Equal(t, p2, p1, "Have: length of payload not zero")
   380  
   381  	file, _ := parser.ParseFromFile(parser.GetTorrentFileList()[0])
   382  	pieceBlock := parser.RandomPieceBlock(file)
   383  
   384  	// BuildRequest
   385  	requestMessage, _ := BuildRequest(pieceBlock)
   386  	size, id, payload = ParseMsg(requestMessage)
   387  	assert.Equal(t, size, uint32(13), "Request: Size not equal")
   388  	assert.Equal(t, id, uint8(6), "Request: Message ID different")
   389  	assert.Equal(t, uint32(payload["index"].(uint32)), pieceBlock.Index, "Request: index field of payload not same")
   390  	assert.Equal(t, uint32(payload["begin"].(uint32)), pieceBlock.Begin, "Request: begin field of payload not same")
   391  	var length uint32
   392  	binary.Read(payload["length"].(*bytes.Buffer), binary.BigEndian, &length)
   393  	assert.Equal(t, length, pieceBlock.Length, "Request: length field of payload not same")
   394  
   395  	// BuildCancel
   396  	cancelMessage, _ := BuildCancel(pieceBlock)
   397  	size, id, payload = ParseMsg(cancelMessage)
   398  	assert.Equal(t, size, uint32(13), "Cancel: Size not equal")
   399  	assert.Equal(t, id, uint8(8), "Cancel: Message ID different")
   400  	assert.Equal(t, uint32(payload["index"].(uint32)), pieceBlock.Index, "Cancel: index field of payload not same")
   401  	assert.Equal(t, uint32(payload["begin"].(uint32)), pieceBlock.Begin, "Cancel: begin field of payload not same")
   402  	binary.Read(payload["length"].(*bytes.Buffer), binary.BigEndian, &length)
   403  	assert.Equal(t, length, pieceBlock.Length, "Cancel: length field of payload not same")
   404  
   405  	// BuildPort
   406  	port1 := uint16(rand.Uint32())
   407  	portMessage, _ := BuildPort(port1)
   408  	size, id, payload = ParseMsg(portMessage)
   409  	assert.Equal(t, size, uint32(3), "Port: Size not equal")
   410  	assert.Equal(t, id, uint8(9), "Port Message ID different")
   411  	var port2 uint16
   412  	binary.Read(payload["payload"].(*bytes.Buffer), binary.BigEndian, &port2)
   413  	assert.Equal(t, p2, p1, "Port: length of payload not zero")
   414  
   415  }
   416
View as plain text