Merge pull request #157 from GrappigPanda/issue-125

Fixes #125

announce/announce.go

@@ -2,13 +2,15 @@ package announce
 
 import (
 	"fmt"
+	r "github.com/GrappigPanda/notorious/kvStoreInterfaces"
 	"net/http"
 	"net/url"
 	"strconv"
 	"strings"
-    r "github.com/GrappigPanda/notorious/kvStoreInterfaces"
 )
 
+// ParseAnnounceData handles getting the annunce data from a remote client and
+// parses it into an acceptable data structure.
 func (a *AnnounceData) ParseAnnounceData(req *http.Request) (err error) {
 	query := req.URL.Query()
 
@@ -161,11 +163,11 @@ func (a *AnnounceData) removeFromKVStorage(subkey string) {
 }
 
 func (a *AnnounceData) infoHashExists() bool {
-	return r.RedisGetBoolKeyVal(a.InfoHash)
+	return r.RedisGetBoolKeyVal(nil, a.InfoHash)
 }
 
 func (a *AnnounceData) createInfoHashKey() {
-	r.CreateNewTorrentKey(a.InfoHash)
+	r.CreateNewTorrentKey(nil, a.InfoHash)
 }
 
 // ParseInfoHash parses the encoded info hash. Such a simple solution for a

announce/definitions.go

@@ -14,8 +14,8 @@ const (
 type AnnounceData struct {
 	InfoHash string //20 byte sha1 hash
 	PeerID   string //max len 20
-	IP        string //optional
-	Event     string // TorrentEvent
+	IP       string //optional
+	Event    string // TorrentEvent
 
 	Port uint64 // port number the peer is listening
 	// on

config/config.go

@@ -2,7 +2,7 @@ package config
 
 import (
 	"github.com/spf13/viper"
-        "strconv"
+	"strconv"
 )
 
 // ConfigStruct holds the values that our config file holds
@@ -25,20 +25,20 @@ func LoadConfig() ConfigStruct {
 	viper.AddConfigPath("../")
 	viper.AddConfigPath("/etc/")
 
-        err := viper.ReadInConfig()
-        if err != nil {
-                panic("Failed to open config file")
-        }
+	err := viper.ReadInConfig()
+	if err != nil {
+		panic("Failed to open config file")
+	}
 
-        if viper.GetBool("UseEnvVariables") == true {
-                viper.AutomaticEnv()
-                viper.BindEnv("mysqluser")
-        }
+	if viper.GetBool("UseEnvVariables") == true {
+		viper.AutomaticEnv()
+		viper.BindEnv("mysqluser")
+	}
 
-        whitelist, err := strconv.ParseBool(viper.Get("whitelist").(string))
-        if err != nil {
-                whitelist = false
-        }
+	whitelist, err := strconv.ParseBool(viper.Get("whitelist").(string))
+	if err != nil {
+		whitelist = false
+	}
 
 	if viper.Get("MySQLPass").(string) != "" {
 		return ConfigStruct{

database/database.go

@@ -9,18 +9,6 @@ import (
 	_ "github.com/jinzhu/gorm/dialects/mysql"
 )
 
-// formatConnectStrings concatenates the data from the config file into a
-// usable MySQL connection string.
-func formatConnectString(c config.ConfigStruct) string {
-	return fmt.Sprintf("%v:%v@tcp(%v:%v)/%v?parseTime=true",
-		c.MySQLUser,
-		c.MySQLPass,
-		c.MySQLHost,
-		c.MySQLPort,
-		c.MySQLDB,
-	)
-}
-
 // OpenConnection does as its name dictates and opens a connection to the
 // MysqlHost listed in the config
 func OpenConnection() (db *gorm.DB, err error) {
@@ -36,6 +24,8 @@ func OpenConnection() (db *gorm.DB, err error) {
 
 // InitDB initializes database tables.
 func InitDB(db *gorm.DB) {
+	db = assertOpenConnection(db)
+
 	db.CreateTable(&White_Torrent{})
 	db.CreateTable(&Torrent{})
 	db.CreateTable(&TrackerStats{})
@@ -44,11 +34,8 @@ func InitDB(db *gorm.DB) {
 
 // AddWhitelistedTorrent adds a torrent to the whitelist so that they may be
 // used by the tracker in the future.
-func (t *White_Torrent) AddWhitelistedTorrent() bool {
-	db, err := OpenConnection()
-	if err != nil {
-		err = err
-	}
+func (t *White_Torrent) AddWhitelistedTorrent(db *gorm.DB) bool {
+	db = assertOpenConnection(db)
 
 	db.Create(t)
 	return db.NewRecord(t)
@@ -57,11 +44,9 @@ func (t *White_Torrent) AddWhitelistedTorrent() bool {
 // GetTorrent retrieves a torrent by its infoHash from the generic torrent
 // table in the database. Note: there's also a whitelisted torrent table
 // (`white_torrent`).
-func GetTorrent(infoHash string) (t *Torrent, err error) {
-	db, err := OpenConnection()
-	if err != nil {
-		err = err
-	}
+func GetTorrent(infoHash string) (db *gorm.DB, t *Torrent, err error) {
+	db = assertOpenConnection(db)
+
 	t = &Torrent{}
 
 	db.Where("info_hash = ?", infoHash).Find(&t)
@@ -70,11 +55,9 @@ func GetTorrent(infoHash string) (t *Torrent, err error) {
 }
 
 // GetWhitelistedTorrent Retrieves a single whitelisted torrent by its infoHash
-func GetWhitelistedTorrent(infoHash string) (t *White_Torrent, err error) {
-	db, err := OpenConnection()
-	if err != nil {
-		err = err
-	}
+func GetWhitelistedTorrent(db *gorm.DB, infoHash string) (t *White_Torrent, err error) {
+	db = assertOpenConnection(db)
+
 	t = &White_Torrent{}
 
 	x := db.Where("info_hash = ?", infoHash).First(&t)
@@ -86,11 +69,8 @@ func GetWhitelistedTorrent(infoHash string) (t *White_Torrent, err error) {
 }
 
 // UpdateStats Handles updating statistics relevant to our tracker.
-func UpdateStats(uploaded uint64, downloaded uint64) {
-	db, err := OpenConnection()
-	if err != nil {
-		err = err
-	}
+func UpdateStats(db *gorm.DB, uploaded uint64, downloaded uint64) {
+	db = assertOpenConnection(db)
 
 	ts := &TrackerStats{}
 	db.First(&ts)
@@ -103,12 +83,9 @@ func UpdateStats(uploaded uint64, downloaded uint64) {
 	return
 }
 
-// UpdateStats Handles updating statistics relevant to our tracker.
-func UpdateTorrentStats(seederDelta int64, leecherDelta int64) {
-	db, err := OpenConnection()
-	if err != nil {
-		err = err
-	}
+// UpdateTorrentStats Handles updating statistics relevant to our tracker.
+func UpdateTorrentStats(db *gorm.DB, seederDelta int64, leecherDelta int64) {
+	db = assertOpenConnection(db)
 
 	t := &Torrent{}
 	db.First(&t)
@@ -121,11 +98,10 @@ func UpdateTorrentStats(seederDelta int64, leecherDelta int64) {
 	return
 }
 
-func UpdatePeerStats(uploaded uint64, downloaded uint64, ip string) {
-	db, err := OpenConnection()
-	if err != nil {
-		err = err
-	}
+// UpdatePeerStats handles updating peer info like hits per ip, downloaded
+// amount, uploaded amounts.
+func UpdatePeerStats(db *gorm.DB, uploaded uint64, downloaded uint64, ip string) {
+	db = assertOpenConnection(db)
 
 	ps := &Peer_Stats{Ip: ip}
 	db.First(&ps)
@@ -137,14 +113,11 @@ func UpdatePeerStats(uploaded uint64, downloaded uint64, ip string) {
 	return
 }
 
-// GetWhitelistedTorrent allows us to retrieve all of the white listed
+// GetWhitelistedTorrents allows us to retrieve all of the white listed
 // torrents. Mostly used for populating the Redis KV storage with all of our
 // whitelisted torrents.
-func GetWhitelistedTorrents() (x *sql.Rows, err error) {
-	db, err := OpenConnection()
-	if err != nil {
-		err = err
-	}
+func GetWhitelistedTorrents(db *gorm.DB) (x *sql.Rows, err error) {
+	db = assertOpenConnection(db)
 
 	x, err = db.Table("white_torrents").Rows()
 	if err != nil {
@@ -156,6 +129,35 @@ func GetWhitelistedTorrents() (x *sql.Rows, err error) {
 
 // ScrapeTorrent supports the Scrape convention
 func ScrapeTorrent(db *gorm.DB, infoHash string) (torrent *Torrent) {
+	db = assertOpenConnection(db)
+
 	db.Where("info_hash = ?", infoHash).First(&torrent)
 	return
 }
+
+// formatConnectStrings concatenates the data from the config file into a
+// usable MySQL connection string.
+func formatConnectString(c config.ConfigStruct) string {
+	return fmt.Sprintf("%v:%v@tcp(%v:%v)/%v?parseTime=true",
+		c.MySQLUser,
+		c.MySQLPass,
+		c.MySQLHost,
+		c.MySQLPort,
+		c.MySQLDB,
+	)
+}
+
+// assertOpenConnection handles asserting a connection passed into a sql
+// function is open, not nil. If nil, we'll create a new connection.
+func assertOpenConnection(db *gorm.DB) *gorm.DB {
+	var err error
+
+	if db == nil {
+		db, err = OpenConnection()
+		if err != nil {
+			err = err
+		}
+	}
+
+	return db
+}

database/database_test.go

@@ -24,7 +24,7 @@ func TestAddWhitelistedTorrent(t *testing.T) {
 		DateAdded: time.Now().Unix(),
 	}
 
-	if !newTorrent.AddWhitelistedTorrent() {
+	if !newTorrent.AddWhitelistedTorrent(nil) {
 		t.Fatalf("Failed to Add a whitelisted torrent")
 	}
 }
@@ -44,10 +44,10 @@ func TestGetWhitelistedTorrents(t *testing.T) {
 		DateAdded: time.Now().Unix(),
 	}
 
-	newTorrent.AddWhitelistedTorrent()
-	newTorrent2.AddWhitelistedTorrent()
+	newTorrent.AddWhitelistedTorrent(nil)
+	newTorrent2.AddWhitelistedTorrent(nil)
 
-	_, err := GetWhitelistedTorrents()
+	_, err := GetWhitelistedTorrents(nil)
 	if err != nil {
 		t.Fatalf("Failed to get all whitelisted torrents: %v", err)
 	}
@@ -61,9 +61,9 @@ func TestGetWhitelistedTorrent(t *testing.T) {
 		DateAdded: time.Now().Unix(),
 	}
 
-	newTorrent.AddWhitelistedTorrent()
+	newTorrent.AddWhitelistedTorrent(nil)
 
-	retval, err := GetWhitelistedTorrent(newTorrent.InfoHash)
+	retval, err := GetWhitelistedTorrent(nil, newTorrent.InfoHash)
 	if err != nil {
 		t.Fatalf("Failed to GetWhitelistedTorrent: %v", err)
 	}
@@ -86,7 +86,7 @@ func TestUpdateStats(t *testing.T) {
 	}
 	DBCONN.Save(&newStats)
 
-	UpdateStats(20, 5)
+	UpdateStats(nil, 20, 5)
 
 	retval := &TrackerStats{}
 	DBCONN.First(&retval)
@@ -118,7 +118,7 @@ func TestUpdatePeerStats(t *testing.T) {
 
 	DBCONN.Save(&newPeer)
 
-	UpdatePeerStats(20, 5, "127.0.0.1")
+	UpdatePeerStats(nil, 20, 5, "127.0.0.1")
 
 	retval := &Peer_Stats{}
 	DBCONN.First(&retval)