From a65755a79bbf875ac53ad1a6f3415f85c9744291 Mon Sep 17 00:00:00 2001
From: cpb8010 <cpb8010@gmail.com>
Date: Tue, 5 Nov 2013 00:31:20 -0500
Subject: [PATCH] Parsing peer_id before whitelist check

---
 server/announce.go    |  2 +-
 server/server.go      | 10 ++++++
 server/server_test.go | 81 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+), 1 deletion(-)
 create mode 100644 server/server_test.go

diff --git a/server/announce.go b/server/announce.go
index 92468d3..c45ca26 100644
--- a/server/announce.go
+++ b/server/announce.go
@@ -41,7 +41,7 @@ func (s Server) serveAnnounce(w http.ResponseWriter, r *http.Request) {
 		}
 
 		// Check if the user's client is whitelisted
-		whitelisted, err := tx.ClientWhitelisted(peerID)
+		whitelisted, err := tx.ClientWhitelisted(parsePeerID(peerID))
 		if err != nil {
 			log.Panicf("server: %s", err)
 		}
diff --git a/server/server.go b/server/server.go
index c2bf126..5683cc9 100644
--- a/server/server.go
+++ b/server/server.go
@@ -135,3 +135,13 @@ func validateUser(tx tracker.Conn, dir string) (*storage.User, error) {
 
 	return user, nil
 }
+
+// Takes a peer_id and returns a ClientID
+func parsePeerID(peerID string) (clientID string) {
+	if peerID[0] == '-' {
+		clientID = peerID[1:7]
+	} else {
+		clientID = peerID[0:6]
+	}
+	return
+}
diff --git a/server/server_test.go b/server/server_test.go
new file mode 100644
index 0000000..e239355
--- /dev/null
+++ b/server/server_test.go
@@ -0,0 +1,81 @@
+// Copyright 2013 The Chihaya Authors. All rights reserved.
+// Use of this source code is governed by the BSD 2-Clause license,
+// which can be found in the LICENSE file.
+
+package server
+
+import (
+	"testing"
+)
+
+type PeerClientPair struct {
+	peerId   string
+	clientId string
+}
+
+var (
+	AzTestClients = []PeerClientPair{
+		{"-AZ3034-6wfG2wk6wWLc", "AZ3034"},
+		{"-AZ3042-6ozMq5q6Q3NX", "AZ3042"},
+		{"-BS5820-oy4La2MWGEFj", "BS5820"},
+		{"-AR6360-6oZyyMWoOOBe", "AR6360"},
+		{"-AG2083-s1hiF8vGAAg0", "AG2083"},
+		{"-AG3003-lEl2Mm4NEO4n", "AG3003"},
+		{"-MR1100-00HS~T7*65rm", "MR1100"},
+		{"-LK0140-ATIV~nbEQAMr", "LK0140"},
+		{"-KT2210-347143496631", "KT2210"},
+		{"-TR0960-6ep6svaa61r4", "TR0960"},
+		{"-XX1150-dv220cotgj4d", "XX1150"},
+		{"-AZ2504-192gwethivju", "AZ2504"},
+		{"-KT4310-3L4UvarKuqIu", "KT4310"},
+		{"-AZ2060-0xJQ02d4309O", "AZ2060"},
+		{"-BD0300-2nkdf08Jd890", "BD0300"},
+		{"-A~0010-a9mn9DFkj39J", "A~0010"},
+		{"-UT2300-MNu93JKnm930", "UT2300"},
+		{"-UT2300-KT4310KT4301", "UT2300"},
+	}
+
+	ShadowTestClients = []PeerClientPair{
+		{"T03A0----f089kjsdf6e", "T03A0-"},
+		{"S58B-----nKl34GoNb75", "S58B--"},
+		{"M4-4-0--9aa757Efd5Bl", "M4-4-0"},
+	}
+
+	OtherTestClients = []PeerClientPair{
+		{"AZ2500BTeYUzyabAfo6U", "AZ2500"}, // BitTyrant
+		{"exbc0JdSklm834kj9Udf", "exbc0J"}, // Old BitComet
+		{"FUTB0L84j542mVc84jkd", "FUTB0L"}, // Alt BitComet
+		{"XBT054d-8602Jn83NnF9", "XBT054"}, // XBT
+		{"OP1011affbecbfabeefb", "OP1011"}, // Opera
+		{"-ML2.7.2-kgjjfkd9762", "ML2.7."}, // MLDonkey
+		{"-BOWA0C-SDLFJWEIORNM", "BOWA0C"}, // Bits on Wheels
+		{"Q1-0-0--dsn34DFn9083", "Q1-0-0"}, // Queen Bee
+		{"Q1-10-0-Yoiumn39BDfO", "Q1-10-"}, // Queen Bee Alt
+		{"346------SDFknl33408", "346---"}, // TorreTopia
+		{"QVOD0054ABFFEDCCDEDB", "QVOD00"}, // Qvod
+	}
+)
+
+func TestParseAZClientID(t *testing.T) {
+	for _, pair := range AzTestClients {
+		if pair.clientId != parsePeerID(pair.peerId) {
+			t.Error("PeerID incorrectly parsed", pair)
+		}
+	}
+}
+
+func TestParseShadowClientID(t *testing.T) {
+	for _, pair := range ShadowTestClients {
+		if pair.clientId != parsePeerID(pair.peerId) {
+			t.Error("PeerID incorrectly parsed", pair)
+		}
+	}
+}
+
+func TestParseOtherClientID(t *testing.T) {
+	for _, pair := range OtherTestClients {
+		if pair.clientId != parsePeerID(pair.peerId) {
+			t.Error("PeerID incorrectly parsed", pair)
+		}
+	}
+}