privtracker/storage.go
2022-03-21 21:56:11 +01:00

167 lines
3.8 KiB
Go

package main
import (
"bytes"
"crypto/sha1"
"encoding/binary"
"net"
"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)))
}
func PutPeer(room, infoHash, ip string, port uint16, seeding bool) {
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),
}
}
client := serialize(ip, port)
if seeding {
shard.swarms[h].seeders[client] = time.Now().Unix()
} else {
shard.swarms[h].leechers[client] = time.Now().Unix()
}
shard.Unlock()
}
func DeletePeer(room, infoHash, ip string, port uint16) {
h := sha1.Sum([]byte(room + infoHash))
shard := shards[shardIndex(h)]
shard.Lock()
if _, ok := shard.swarms[h]; !ok {
return
}
client := serialize(ip, port)
delete(shard.swarms[h].seeders, client)
delete(shard.swarms[h].leechers, client)
shard.Unlock()
}
func GraduateLeecher(room, infoHash, ip string, port uint16) {
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),
}
}
client := serialize(ip, port)
shard.swarms[h].seeders[client] = time.Now().Unix()
delete(shard.swarms[h].leechers, client)
shard.Unlock()
}
func GetPeers(room, infoHash, ip string, port uint16, seeding bool, numWant uint) (peersIPv4, peersIPv6 []byte, numSeeders, numLeechers int) {
h := sha1.Sum([]byte(room + infoHash))
shard := shards[shardIndex(h)]
shard.RLock()
client := serialize(ip, port)
// seeders don't need other seeders
if !seeding {
for peer := range shard.swarms[h].seeders {
if numWant == 0 {
break
}
if bytes.HasPrefix([]byte(peer), v4InV6Prefix) {
peersIPv4 = append(peersIPv4, peer[12:]...)
} else {
peersIPv6 = append(peersIPv6, peer...)
}
numWant--
}
}
for peer := range shard.swarms[h].leechers {
if peer == client {
continue
}
if numWant == 0 {
break
}
if bytes.HasPrefix([]byte(peer), v4InV6Prefix) {
peersIPv4 = append(peersIPv4, peer[12:]...)
} else {
peersIPv6 = append(peersIPv6, peer...)
}
numWant--
}
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 {
time.Sleep(time.Minute * 3)
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()
}
}
}