privtracker/storage.go

170 lines
3.8 KiB
Go
Raw Normal View History

2020-12-06 15:36:57 -06:00
package main
import (
"bytes"
"crypto/sha1"
"encoding/binary"
"net"
2022-11-09 03:11:53 -06:00
"runtime"
2020-12-06 15:36:57 -06:00
"sync"
"time"
)
type serializedPeer string
type hash [20]byte
type shard struct {
swarms map[hash]swarm
sync.RWMutex
}
type swarm struct {
seeders map[serializedPeer]int64
leechers map[serializedPeer]int64
}
var shards = NewShards(512)
var v4InV6Prefix = []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff}
func shardIndex(hash [20]byte) int {
return int(binary.BigEndian.Uint32(hash[:4])) % len(shards)
}
func NewShards(size int) []*shard {
shards := make([]*shard, size)
for i := 0; i < size; i++ {
shards[i] = &shard{
swarms: make(map[hash]swarm),
}
}
return shards
}
func serialize(ip string, port uint16) serializedPeer {
return serializedPeer(append(net.ParseIP(ip), byte(port>>8), byte(port)))
}
2020-12-08 07:28:17 -06:00
func PutPeer(room, infoHash, ip string, port uint16, seeding bool) {
2020-12-06 15:36:57 -06:00
h := sha1.Sum([]byte(room + infoHash))
shard := shards[shardIndex(h)]
shard.Lock()
if _, ok := shard.swarms[h]; !ok {
shard.swarms[h] = swarm{
seeders: make(map[serializedPeer]int64),
leechers: make(map[serializedPeer]int64),
}
}
2020-12-08 07:28:17 -06:00
client := serialize(ip, port)
2020-12-06 15:36:57 -06:00
if seeding {
2020-12-08 07:28:17 -06:00
shard.swarms[h].seeders[client] = time.Now().Unix()
2020-12-06 15:36:57 -06:00
} else {
2020-12-08 07:28:17 -06:00
shard.swarms[h].leechers[client] = time.Now().Unix()
2020-12-06 15:36:57 -06:00
}
shard.Unlock()
}
2020-12-08 07:28:17 -06:00
func DeletePeer(room, infoHash, ip string, port uint16) {
2020-12-06 15:36:57 -06:00
h := sha1.Sum([]byte(room + infoHash))
shard := shards[shardIndex(h)]
shard.Lock()
if _, ok := shard.swarms[h]; !ok {
return
}
2020-12-08 07:28:17 -06:00
client := serialize(ip, port)
delete(shard.swarms[h].seeders, client)
delete(shard.swarms[h].leechers, client)
2020-12-06 15:36:57 -06:00
shard.Unlock()
}
2020-12-08 07:28:17 -06:00
func GraduateLeecher(room, infoHash, ip string, port uint16) {
2020-12-06 15:36:57 -06:00
h := sha1.Sum([]byte(room + infoHash))
shard := shards[shardIndex(h)]
shard.Lock()
if _, ok := shard.swarms[h]; !ok {
shard.swarms[h] = swarm{
seeders: make(map[serializedPeer]int64),
leechers: make(map[serializedPeer]int64),
}
}
2020-12-08 07:28:17 -06:00
client := serialize(ip, port)
shard.swarms[h].seeders[client] = time.Now().Unix()
delete(shard.swarms[h].leechers, client)
2020-12-06 15:36:57 -06:00
shard.Unlock()
}
2020-12-08 07:29:56 -06:00
func GetPeers(room, infoHash, ip string, port uint16, seeding bool, numWant uint) (peersIPv4, peersIPv6 []byte, numSeeders, numLeechers int) {
2020-12-06 15:36:57 -06:00
h := sha1.Sum([]byte(room + infoHash))
shard := shards[shardIndex(h)]
shard.RLock()
2020-12-08 07:28:17 -06:00
client := serialize(ip, port)
2020-12-08 06:34:02 -06:00
// seeders don't need other seeders
if !seeding {
2020-12-08 07:28:17 -06:00
for peer := range shard.swarms[h].seeders {
2020-12-06 17:52:51 -06:00
if numWant == 0 {
break
}
2020-12-08 07:28:17 -06:00
if bytes.HasPrefix([]byte(peer), v4InV6Prefix) {
2020-12-08 07:29:56 -06:00
peersIPv4 = append(peersIPv4, peer[12:]...)
2020-12-06 17:52:51 -06:00
} else {
2020-12-08 07:29:56 -06:00
peersIPv6 = append(peersIPv6, peer...)
2020-12-06 17:52:51 -06:00
}
numWant--
}
2020-12-08 06:34:02 -06:00
}
2020-12-08 07:28:17 -06:00
for peer := range shard.swarms[h].leechers {
if peer == client {
continue
}
2020-12-08 06:34:02 -06:00
if numWant == 0 {
break
}
2020-12-08 07:28:17 -06:00
if bytes.HasPrefix([]byte(peer), v4InV6Prefix) {
2020-12-08 07:29:56 -06:00
peersIPv4 = append(peersIPv4, peer[12:]...)
2020-12-08 06:34:02 -06:00
} else {
2020-12-08 07:29:56 -06:00
peersIPv6 = append(peersIPv6, peer...)
2020-12-06 17:52:51 -06:00
}
2020-12-08 06:34:02 -06:00
numWant--
2020-12-06 15:36:57 -06:00
}
numSeeders = len(shard.swarms[h].seeders)
numLeechers = len(shard.swarms[h].leechers)
shard.RUnlock()
return
}
func GetStats(room, infoHash string) (numSeeders, numLeechers int) {
h := sha1.Sum([]byte(room + infoHash))
shard := shards[shardIndex(h)]
shard.RLock()
numSeeders = len(shard.swarms[h].seeders)
numLeechers = len(shard.swarms[h].leechers)
shard.RUnlock()
return
}
func Cleanup() {
for {
2022-03-21 15:56:11 -05:00
time.Sleep(time.Minute * 3)
2020-12-06 15:36:57 -06:00
expiration := time.Now().Unix() - 600
for _, shard := range shards {
shard.Lock()
for h, swarm := range shard.swarms {
for peer, lastSeen := range swarm.seeders {
if lastSeen < expiration {
delete(swarm.seeders, peer)
}
}
for peer, lastSeen := range swarm.leechers {
if lastSeen < expiration {
delete(swarm.leechers, peer)
}
}
if len(swarm.leechers) == 0 && len(swarm.seeders) == 0 {
delete(shard.swarms, h)
}
}
shard.Unlock()
}
2022-11-09 03:11:53 -06:00
runtime.GC()
2020-12-06 15:36:57 -06:00
}
}