From: Matt Joiner <anacrolix@gmail.com>
Date: Wed, 8 Nov 2017 08:28:03 +0000 (+1100)
Subject: Generalise use of peerExtensionBytes
X-Git-Tag: v1.0.0~339
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=b38996ae2a01564e9089e389b67cf86bf5745f98;p=btrtrc.git

Generalise use of peerExtensionBytes
---

diff --git a/handshake.go b/handshake.go
index b79b5fd6..0afc4a86 100644
--- a/handshake.go
+++ b/handshake.go
@@ -15,6 +15,14 @@ import (
 	"github.com/anacrolix/torrent/mse"
 )
 
+type ExtensionBit uint
+
+const (
+	ExtensionBitDHT      = 0  // http://www.bittorrent.org/beps/bep_0005.html
+	ExtensionBitExtended = 20 // http://www.bittorrent.org/beps/bep_0010.html
+	ExtensionBitFast     = 2  // http://www.bittorrent.org/beps/bep_0006.html
+)
+
 func handshakeWriter(w io.Writer, bb <-chan []byte, done chan<- error) {
 	var err error
 	for b := range bb {
@@ -31,16 +39,24 @@ type (
 	peerID             [20]byte
 )
 
-func (pex *peerExtensionBytes) SupportsExtended() bool {
-	return pex[5]&0x10 != 0
+func (pex peerExtensionBytes) SupportsExtended() bool {
+	return pex.GetBit(ExtensionBitExtended)
+}
+
+func (pex peerExtensionBytes) SupportsDHT() bool {
+	return pex.GetBit(ExtensionBitDHT)
+}
+
+func (pex peerExtensionBytes) SupportsFast() bool {
+	return pex.GetBit(ExtensionBitFast)
 }
 
-func (pex *peerExtensionBytes) SupportsDHT() bool {
-	return pex[7]&0x01 != 0
+func (pex *peerExtensionBytes) SetBit(bit ExtensionBit) {
+	pex[7-bit/8] |= 1 << bit % 8
 }
 
-func (pex *peerExtensionBytes) SupportsFast() bool {
-	return pex[7]&0x04 != 0
+func (pex peerExtensionBytes) GetBit(bit ExtensionBit) bool {
+	return pex[7-bit/8]&(1<<(bit%8)) != 0
 }
 
 type handshakeResult struct {
diff --git a/handshake_test.go b/handshake_test.go
new file mode 100644
index 00000000..b477a2df
--- /dev/null
+++ b/handshake_test.go
@@ -0,0 +1,16 @@
+package torrent
+
+import (
+	"testing"
+
+	"github.com/anacrolix/missinggo"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDefaultExtensionBytes(t *testing.T) {
+	var pex peerExtensionBytes
+	missinggo.CopyExact(&pex, defaultExtensionBytes)
+	assert.True(t, pex.SupportsDHT())
+	assert.True(t, pex.SupportsExtended())
+	assert.False(t, pex.SupportsFast())
+}