message_utils.go

Documentation: github.com/concurrency-8/torrent

     1  package torrent
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  
     7  	// "github.com/concurrency-8/queue"
     8  
     9  	"github.com/concurrency-8/parser"
    10  	"github.com/concurrency-8/tracker"
    11  )
    12  
    13  // Payload refers to the message payload that is to be sent.
    14  type Payload map[string]interface{}
    15  
    16  // BuildHandshake returns a pointer to a buffer.
    17  // Buffer looks like:
    18  //	uint8		: pstrlen	- Length of pstr
    19  //	[pstrlen]byte	: pstr		- pstr, the string identifier of the protocol
    20  //	[8]byte		: reserved	- 8 reserved bytes
    21  //	[20]byte	: infohash	- SHA1 hash of the info key in the metainfo file. Same as the info hash transmitted in tracker requests
    22  //	[20]byte	: peerID	- 20 byte unique ID for the client. Usually the same peerID transmitted in tracker requests
    23  // In version 1.0 of the BitTorrent protocol, pstrlen = 19, and pstr = "BitTorrent protocol"
    24  func BuildHandshake(report tracker.ClientStatusReport) (handshake *bytes.Buffer, err error) {
    25  	handshake = new(bytes.Buffer)
    26  
    27  	// pstrlen
    28  	if err = binary.Write(handshake, binary.BigEndian, uint8(19)); err != nil {
    29  		return
    30  	}
    31  
    32  	// pstr
    33  	if err = binary.Write(handshake, binary.BigEndian, []byte("BitTorrent protocol")); err != nil {
    34  		return
    35  	}
    36  
    37  	// reserved
    38  	if err = binary.Write(handshake, binary.BigEndian, uint64(0)); err != nil {
    39  		return
    40  	}
    41  
    42  	// infohash
    43  	var infohashFromFile [20]byte
    44  	copy(infohashFromFile[:], report.TorrentFile.InfoHash)
    45  	if err = binary.Write(handshake, binary.BigEndian, infohashFromFile); err != nil {
    46  		return
    47  	}
    48  
    49  	// peerID
    50  	if err = binary.Write(handshake, binary.BigEndian, []byte(report.PeerID)); err != nil {
    51  		return
    52  	}
    53  
    54  	return
    55  }
    56  
    57  // BuildKeepAlive returns pointer to an empty buffer (4 bytes)
    58  func BuildKeepAlive() (keepAlive *bytes.Buffer) {
    59  	keepAlive = bytes.NewBuffer(make([]byte, 4))
    60  
    61  	return
    62  }
    63  
    64  // BuildChoke returns pointer to a buffer.
    65  // Buffer looks like:
    66  //	uint32	: length	- Length of remaining part(message) = 1
    67  //	uint8	: messageType	- For choke, messageType = 0
    68  func BuildChoke() (choke *bytes.Buffer, err error) {
    69  	choke = new(bytes.Buffer)
    70  
    71  	if err = binary.Write(choke, binary.BigEndian, uint32(1)); err != nil {
    72  		return
    73  	}
    74  
    75  	if err = binary.Write(choke, binary.BigEndian, uint8(0)); err != nil {
    76  		return
    77  	}
    78  
    79  	return
    80  }
    81  
    82  // BuildUnchoke returns pointer to a buffer.
    83  // Buffer looks like:
    84  //	uint32	: length	- Length of remaining part(message) = 1
    85  //	uint8	: messageType	- For unchoke, messageType = 1
    86  func BuildUnchoke() (unchoke *bytes.Buffer, err error) {
    87  	unchoke = new(bytes.Buffer)
    88  
    89  	if err = binary.Write(unchoke, binary.BigEndian, uint32(1)); err != nil {
    90  		return
    91  	}
    92  
    93  	if err = binary.Write(unchoke, binary.BigEndian, uint8(1)); err != nil {
    94  		return
    95  	}
    96  
    97  	return
    98  }
    99  
   100  // BuildInterested returns pointer to a buffer.
   101  // Buffer looks like:
   102  //	uint32	: length	- Length of remaining part(message) = 1
   103  //	uint8	: messageType	- For interested, messageType = 2
   104  func BuildInterested() (interested *bytes.Buffer, err error) {
   105  	interested = new(bytes.Buffer)
   106  
   107  	if err = binary.Write(interested, binary.BigEndian, uint32(1)); err != nil {
   108  		return
   109  	}
   110  
   111  	if err = binary.Write(interested, binary.BigEndian, uint8(2)); err != nil {
   112  		return
   113  	}
   114  
   115  	return
   116  }
   117  
   118  // BuildUninterested returns pointer to a buffer.
   119  // Buffer looks like:
   120  //	uint32	: length	- Length of remaining part(message) = 1
   121  //	uint8	: messageType	- For uninterested, messageType = 3
   122  func BuildUninterested() (uninterested *bytes.Buffer, err error) {
   123  	uninterested = new(bytes.Buffer)
   124  
   125  	if err = binary.Write(uninterested, binary.BigEndian, uint32(1)); err != nil {
   126  		return
   127  	}
   128  
   129  	if err = binary.Write(uninterested, binary.BigEndian, uint8(3)); err != nil {
   130  		return
   131  	}
   132  
   133  	return
   134  }
   135  
   136  // BuildHave returns pointer to a buffer. This takes uint32 payload(piece index) as an argument
   137  // Buffer looks like:
   138  //	uint32	: length	- Length of remaining part(message) = 5
   139  //	uint8	: messageType	- for have, messageType = 4
   140  //	uint32	: piece index	- payload
   141  func BuildHave(payload uint32) (have *bytes.Buffer, err error) {
   142  	have = new(bytes.Buffer)
   143  
   144  	if err = binary.Write(have, binary.BigEndian, uint32(5)); err != nil {
   145  		return
   146  	}
   147  
   148  	if err = binary.Write(have, binary.BigEndian, uint8(4)); err != nil {
   149  		return
   150  	}
   151  
   152  	if err = binary.Write(have, binary.BigEndian, uint32(payload)); err != nil {
   153  		return
   154  	}
   155  
   156  	return
   157  }
   158  
   159  // BuildRequest returns pointer to a buffer. This takes parser.PieceBlock as an argument
   160  //	uint32	: length	- Length of remaining part(message) = 13
   161  //	uint8	: messageType	- for request, message = 5
   162  //	uint32	: piece index	- parser.PieceBlock.Index for payload
   163  //	uint32	: piece begin	- parser.PieceBlock.Begin for payload
   164  //	uint32	: piece length	- parser.PieceBlock.Length for payload
   165  func BuildRequest(payload parser.PieceBlock) (request *bytes.Buffer, err error) {
   166  	request = new(bytes.Buffer)
   167  
   168  	// Length of message
   169  	if err = binary.Write(request, binary.BigEndian, uint32(13)); err != nil {
   170  		return
   171  	}
   172  	// message type
   173  	if err = binary.Write(request, binary.BigEndian, uint8(6)); err != nil {
   174  		return
   175  	}
   176  	// piece index
   177  	if err = binary.Write(request, binary.BigEndian, uint32(payload.Index)); err != nil {
   178  		return
   179  	}
   180  	// piece begin
   181  	if err = binary.Write(request, binary.BigEndian, uint32(payload.Begin)); err != nil {
   182  		return
   183  	}
   184  	// piece length
   185  	if err = binary.Write(request, binary.BigEndian, uint32(payload.Length)); err != nil {
   186  		return
   187  	}
   188  
   189  	return
   190  }
   191  
   192  // BuildPiece returns pointer to a buffer having the piece. Takes the parser.PieceBlock object as an arg
   193  //	uint32	: length	- length of remaining part (message) = payload length + 9
   194  //	uint8	: messageType	- for piece, type = 7
   195  //	uint32	: piece index	- parser.PieceBlock.Index for payload
   196  //	uint32	: piece begin	- parser.PieceBlock.Begin for payload
   197  //	[]byte	: piece		- the data of the piece, parser.PieceBlock.Block for payload
   198  // func BuildPiece(payload parser.PieceBlock) (piece *bytes.Buffer, err error) {
   199  // 	piece = new(bytes.Buffer)
   200  
   201  // 	// Length of message (Has the piece)
   202  // 	if err = binary.Write(piece, binary.BigEndian, uint32(len(payload.Block.Bytes())+9)); err != nil {
   203  // 		return
   204  // 	}
   205  
   206  // 	// Message type
   207  // 	if err = binary.Write(piece, binary.BigEndian, uint8(7)); err != nil {
   208  // 		return
   209  // 	}
   210  
   211  // 	// piece index
   212  // 	if err = binary.Write(piece, binary.BigEndian, uint32(payload.Index)); err != nil {
   213  // 		return
   214  // 	}
   215  
   216  // 	// piece begin
   217  // 	if err = binary.Write(piece, binary.BigEndian, uint32(payload.Begin)); err != nil {
   218  // 		return
   219  // 	}
   220  
   221  // 	// piece
   222  // 	if err = binary.Write(piece, binary.BigEndian, payload.Block.Bytes()); err != nil {
   223  // 		return
   224  // 	}
   225  
   226  // 	return
   227  // }
   228  
   229  // BuildCancel returns pointer to a buffer. Takes parser.PieceBlock object as arg
   230  //	uint32	: length	- Length of the remaining message = 13
   231  //	uint8	: messageType	- for cancel, messageType = 8
   232  //	uint32	: piece index	- parser.PieceBlock.Index for payload
   233  //	uint32	: piece begin	- parser.PieceBlock.Begin for payload
   234  //	uint32	: piece length	- parser.PieceBlock.Length for payload
   235  func BuildCancel(payload parser.PieceBlock) (cancelBuf *bytes.Buffer, err error) {
   236  	cancelBuf = new(bytes.Buffer)
   237  
   238  	// Length of Message
   239  	if err = binary.Write(cancelBuf, binary.BigEndian, uint32(13)); err != nil {
   240  		return
   241  	}
   242  
   243  	// Message type - Cancel
   244  	if err = binary.Write(cancelBuf, binary.BigEndian, uint8(8)); err != nil {
   245  		return
   246  	}
   247  
   248  	// piece index
   249  	if err = binary.Write(cancelBuf, binary.BigEndian, payload.Index); err != nil {
   250  		return
   251  	}
   252  
   253  	// piece begin
   254  	if err = binary.Write(cancelBuf, binary.BigEndian, payload.Begin); err != nil {
   255  		return
   256  	}
   257  
   258  	// piece length
   259  	if err = binary.Write(cancelBuf, binary.BigEndian, payload.Length); err != nil {
   260  		return
   261  	}
   262  
   263  	return
   264  }
   265  
   266  // BuildPort returns a pointer to a buffer. Takes uint16 port as arg
   267  //	uint32	: length	- length of remaining message = 3
   268  //	uint8	: messageType	- for port, messageType = 9
   269  //	uint16	: port		- the argument, port
   270  func BuildPort(port uint16) (portBuf *bytes.Buffer, err error) {
   271  	portBuf = new(bytes.Buffer)
   272  
   273  	// Length of message
   274  	if err = binary.Write(portBuf, binary.BigEndian, uint32(3)); err != nil {
   275  		return
   276  	}
   277  	// Message type = 9
   278  	if err = binary.Write(portBuf, binary.BigEndian, uint8(9)); err != nil {
   279  		return
   280  	}
   281  	// listen-port
   282  	if err = binary.Write(portBuf, binary.BigEndian, uint16(port)); err != nil {
   283  		return
   284  	}
   285  	return
   286  }
   287  
   288  // ParseMsg parses a message
   289  func ParseMsg(msg *bytes.Buffer) (size uint32, id uint8, payload Payload) {
   290  	payload = make(Payload)
   291  	binary.Read(msg, binary.BigEndian, &size)
   292  	// fmt.Println("message:", msg.Bytes())
   293  	if size > 0 {
   294  		binary.Read(msg, binary.BigEndian, &id)
   295  	}
   296  	if size > 1 {
   297  
   298  		if id == 6 || id == 7 || id == 8 {
   299  			rest := bytes.NewBuffer(msg.Bytes()[8:])
   300  			var index, begin uint32
   301  			binary.Read(msg, binary.BigEndian, &index)
   302  			binary.Read(msg, binary.BigEndian, &begin)
   303  			payload["index"] = index
   304  			payload["begin"] = begin
   305  
   306  			if id == 7 {
   307  				payload["block"] = rest
   308  				// fmt.Println(rest.Bytes()[:10])
   309  			} else {
   310  				payload["length"] = rest
   311  			}
   312  		} else {
   313  			payload["payload"] = msg
   314  		}
   315  	}
   316  
   317  	return
   318  
   319  }
   320
View as plain text