enki 7c92aa3ded
Some checks are pending
CI Pipeline / Run Tests (push) Waiting to run
CI Pipeline / Lint Code (push) Waiting to run
CI Pipeline / Security Scan (push) Waiting to run
CI Pipeline / E2E Tests (push) Blocked by required conditions
Major DHT and Torrent fixes.
2025-08-29 21:18:36 -07:00

707 lines
20 KiB
Go

package tracker
import (
"fmt"
"log"
"net/http"
"sync"
"time"
"github.com/gorilla/websocket"
)
type WebSocketTracker struct {
upgrader websocket.Upgrader
swarms map[string]*Swarm
mu sync.RWMutex
tracker *Tracker // Reference to main tracker for HTTP fallback
statsTracker *StatsTracker
}
// StatsTracker collects WebRTC statistics
type StatsTracker struct {
mu sync.RWMutex
totalConnections int64
activeConnections int64
totalBytesUploaded int64
totalBytesDownloaded int64
connectionFailures int64
iceFailures int64
lastReported time.Time
}
type Swarm struct {
peers map[string]*WebRTCPeer
mu sync.RWMutex
}
// PeerConnectionState represents the connection state of a peer
type PeerConnectionState string
const (
StateConnecting PeerConnectionState = "connecting"
StateConnected PeerConnectionState = "connected"
StateDisconnected PeerConnectionState = "disconnected"
StateFailed PeerConnectionState = "failed"
)
// PeerStats tracks statistics for a peer
type PeerStats struct {
BytesUploaded int64 `json:"bytes_uploaded"`
BytesDownloaded int64 `json:"bytes_downloaded"`
ConnectionTime time.Time `json:"connection_time"`
LastActivity time.Time `json:"last_activity"`
RTT int `json:"rtt_ms"`
ConnectionQuality string `json:"connection_quality"`
}
type WebRTCPeer struct {
ID string `json:"peer_id"`
Conn *websocket.Conn `json:"-"`
LastSeen time.Time `json:"last_seen"`
InfoHashes []string `json:"info_hashes"`
State PeerConnectionState `json:"state"`
IsSeeder bool `json:"is_seeder"`
Stats *PeerStats `json:"stats"`
UserAgent string `json:"user_agent"`
WebRTCPeers map[string]time.Time `json:"-"` // Track connections to other peers
SupportsHTTP bool `json:"supports_http"`
Endpoint string `json:"endpoint,omitempty"`
}
type WebTorrentMessage struct {
Action string `json:"action"`
InfoHash string `json:"info_hash,omitempty"`
PeerID string `json:"peer_id,omitempty"`
Answer map[string]interface{} `json:"answer,omitempty"`
Offer map[string]interface{} `json:"offer,omitempty"`
ToPeerID string `json:"to_peer_id,omitempty"`
FromPeerID string `json:"from_peer_id,omitempty"`
NumWant int `json:"numwant,omitempty"`
// ICE candidate exchange
Candidate map[string]interface{} `json:"candidate,omitempty"`
// Connection state tracking
ConnectionState string `json:"connection_state,omitempty"`
// Statistics
Uploaded int64 `json:"uploaded,omitempty"`
Downloaded int64 `json:"downloaded,omitempty"`
Left int64 `json:"left,omitempty"`
Event string `json:"event,omitempty"`
// Client information
UserAgent string `json:"user_agent,omitempty"`
SupportsHTTP bool `json:"supports_http,omitempty"`
Port int `json:"port,omitempty"`
// HTTP fallback
RequestHTTP bool `json:"request_http,omitempty"`
}
func NewWebSocketTracker(tracker *Tracker) *WebSocketTracker {
return &WebSocketTracker{
upgrader: websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true // Allow all origins for WebTorrent compatibility
},
},
swarms: make(map[string]*Swarm),
tracker: tracker,
statsTracker: &StatsTracker{
lastReported: time.Now(),
},
}
}
func (wt *WebSocketTracker) HandleWS(w http.ResponseWriter, r *http.Request) {
conn, err := wt.upgrader.Upgrade(w, r, nil)
if err != nil {
log.Printf("WebSocket upgrade failed: %v", err)
return
}
defer conn.Close()
log.Printf("WebTorrent client connected from %s", r.RemoteAddr)
// Handle WebTorrent protocol
for {
var msg WebTorrentMessage
if err := conn.ReadJSON(&msg); err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
log.Printf("WebSocket error: %v", err)
}
break
}
switch msg.Action {
case "announce":
wt.handleAnnounce(conn, msg)
case "scrape":
wt.handleScrape(conn, msg)
case "offer":
wt.handleOffer(conn, msg)
case "answer":
wt.handleAnswer(conn, msg)
case "ice":
wt.handleICE(conn, msg)
case "connection_state":
wt.handleConnectionState(conn, msg)
case "stats":
wt.handleStats(conn, msg)
}
}
}
func (wt *WebSocketTracker) handleAnnounce(conn *websocket.Conn, msg WebTorrentMessage) {
wt.mu.Lock()
defer wt.mu.Unlock()
// Get or create swarm
if wt.swarms[msg.InfoHash] == nil {
wt.swarms[msg.InfoHash] = &Swarm{
peers: make(map[string]*WebRTCPeer),
}
}
swarm := wt.swarms[msg.InfoHash]
swarm.mu.Lock()
defer swarm.mu.Unlock()
now := time.Now()
// Determine if peer is seeder
isSeeder := msg.Left == 0 || msg.Event == "completed"
// Create or update peer with enhanced state tracking
existingPeer := swarm.peers[msg.PeerID]
var stats *PeerStats
if existingPeer != nil {
stats = existingPeer.Stats
// Update existing stats
if msg.Uploaded > 0 {
stats.BytesUploaded = msg.Uploaded
}
if msg.Downloaded > 0 {
stats.BytesDownloaded = msg.Downloaded
}
stats.LastActivity = now
} else {
// New peer
stats = &PeerStats{
BytesUploaded: msg.Uploaded,
BytesDownloaded: msg.Downloaded,
ConnectionTime: now,
LastActivity: now,
ConnectionQuality: "unknown",
}
wt.statsTracker.mu.Lock()
wt.statsTracker.totalConnections++
wt.statsTracker.activeConnections++
wt.statsTracker.mu.Unlock()
}
peer := &WebRTCPeer{
ID: msg.PeerID,
Conn: conn,
LastSeen: now,
InfoHashes: []string{msg.InfoHash},
State: StateConnecting,
IsSeeder: isSeeder,
Stats: stats,
UserAgent: msg.UserAgent,
WebRTCPeers: make(map[string]time.Time),
SupportsHTTP: msg.SupportsHTTP,
Endpoint: fmt.Sprintf("%s:%d", conn.RemoteAddr().String(), msg.Port),
}
swarm.peers[msg.PeerID] = peer
// Handle different events
switch msg.Event {
case "stopped":
wt.removePeer(swarm, msg.PeerID)
return
case "completed":
log.Printf("Peer %s completed torrent %s", msg.PeerID, msg.InfoHash[:8])
}
// Count active seeders and leechers
var seeders, leechers int
var activePeers []map[string]interface{}
numWant := msg.NumWant
if numWant == 0 {
numWant = 30 // Default
}
count := 0
for peerID, p := range swarm.peers {
if p.State != StateDisconnected && p.State != StateFailed {
if p.IsSeeder {
seeders++
} else {
leechers++
}
if peerID != msg.PeerID && count < numWant {
peerData := map[string]interface{}{
"id": peerID,
}
// Include HTTP endpoint if available and WebRTC not working
if p.SupportsHTTP && msg.RequestHTTP {
peerData["endpoint"] = p.Endpoint
peerData["protocol"] = "http"
}
activePeers = append(activePeers, peerData)
count++
}
}
}
response := map[string]interface{}{
"action": "announce",
"interval": 300, // 5 minutes for WebTorrent
"info_hash": msg.InfoHash,
"complete": seeders,
"incomplete": leechers,
"peers": activePeers,
}
// Add HTTP fallback information if requested or if WebRTC is failing
if msg.RequestHTTP || wt.shouldProvideHTTPFallback(msg.InfoHash) {
wt.addHTTPFallback(response, msg.InfoHash)
}
if err := conn.WriteJSON(response); err != nil {
log.Printf("Failed to send announce response: %v", err)
// Mark peer as disconnected if we can't send to them
peer.State = StateDisconnected
}
}
func (wt *WebSocketTracker) handleScrape(conn *websocket.Conn, msg WebTorrentMessage) {
wt.mu.RLock()
defer wt.mu.RUnlock()
files := make(map[string]map[string]int)
if swarm := wt.swarms[msg.InfoHash]; swarm != nil {
swarm.mu.RLock()
files[msg.InfoHash] = map[string]int{
"complete": len(swarm.peers), // Simplified
"incomplete": 0,
"downloaded": len(swarm.peers),
}
swarm.mu.RUnlock()
}
response := map[string]interface{}{
"action": "scrape",
"files": files,
}
if err := conn.WriteJSON(response); err != nil {
log.Printf("Failed to send scrape response: %v", err)
}
}
func (wt *WebSocketTracker) handleOffer(conn *websocket.Conn, msg WebTorrentMessage) {
wt.mu.RLock()
defer wt.mu.RUnlock()
if swarm := wt.swarms[msg.InfoHash]; swarm != nil {
swarm.mu.RLock()
if targetPeer := swarm.peers[msg.ToPeerID]; targetPeer != nil {
// Forward offer to target peer
offerMsg := map[string]interface{}{
"action": "offer",
"info_hash": msg.InfoHash,
"peer_id": msg.FromPeerID,
"offer": msg.Offer,
"from_peer_id": msg.FromPeerID,
"to_peer_id": msg.ToPeerID,
}
if err := targetPeer.Conn.WriteJSON(offerMsg); err != nil {
log.Printf("Failed to forward offer: %v", err)
targetPeer.State = StateDisconnected
} else {
// Track connection attempt
if fromPeer := swarm.peers[msg.FromPeerID]; fromPeer != nil {
fromPeer.WebRTCPeers[msg.ToPeerID] = time.Now()
log.Printf("Forwarded offer from %s to %s for %s",
msg.FromPeerID, msg.ToPeerID, msg.InfoHash[:8])
}
}
} else {
log.Printf("Target peer %s not found for offer", msg.ToPeerID)
}
swarm.mu.RUnlock()
}
}
func (wt *WebSocketTracker) handleAnswer(conn *websocket.Conn, msg WebTorrentMessage) {
wt.mu.RLock()
defer wt.mu.RUnlock()
if swarm := wt.swarms[msg.InfoHash]; swarm != nil {
swarm.mu.RLock()
if targetPeer := swarm.peers[msg.ToPeerID]; targetPeer != nil {
// Forward answer to target peer
answerMsg := map[string]interface{}{
"action": "answer",
"info_hash": msg.InfoHash,
"peer_id": msg.FromPeerID,
"answer": msg.Answer,
"from_peer_id": msg.FromPeerID,
"to_peer_id": msg.ToPeerID,
}
if err := targetPeer.Conn.WriteJSON(answerMsg); err != nil {
log.Printf("Failed to forward answer: %v", err)
targetPeer.State = StateDisconnected
} else {
// Track connection completion
if fromPeer := swarm.peers[msg.FromPeerID]; fromPeer != nil {
fromPeer.WebRTCPeers[msg.ToPeerID] = time.Now()
log.Printf("Forwarded answer from %s to %s for %s",
msg.FromPeerID, msg.ToPeerID, msg.InfoHash[:8])
}
}
} else {
log.Printf("Target peer %s not found for answer", msg.ToPeerID)
}
swarm.mu.RUnlock()
}
}
// Cleanup expired peers
func (wt *WebSocketTracker) StartCleanup() {
ticker := time.NewTicker(5 * time.Minute)
go func() {
defer ticker.Stop()
for range ticker.C {
wt.cleanupExpiredPeers()
}
}()
}
func (wt *WebSocketTracker) cleanupExpiredPeers() {
wt.mu.Lock()
defer wt.mu.Unlock()
now := time.Now()
expiry := now.Add(-10 * time.Minute) // 10 minute timeout
removedPeers := 0
for infoHash, swarm := range wt.swarms {
swarm.mu.Lock()
for peerID, peer := range swarm.peers {
if peer.LastSeen.Before(expiry) || peer.State == StateDisconnected || peer.State == StateFailed {
if peer.Conn != nil {
peer.Conn.Close()
}
delete(swarm.peers, peerID)
removedPeers++
// Update stats
wt.statsTracker.mu.Lock()
if wt.statsTracker.activeConnections > 0 {
wt.statsTracker.activeConnections--
}
wt.statsTracker.mu.Unlock()
} else {
// Update connection quality for active peers
wt.updateConnectionQuality(peer)
}
}
// Remove empty swarms
if len(swarm.peers) == 0 {
delete(wt.swarms, infoHash)
}
swarm.mu.Unlock()
}
if removedPeers > 0 {
log.Printf("Cleaned up %d expired/disconnected peers", removedPeers)
}
}
// updateConnectionQuality calculates connection quality based on various metrics
func (wt *WebSocketTracker) updateConnectionQuality(peer *WebRTCPeer) {
if peer.Stats == nil {
return
}
now := time.Now()
connectionAge := now.Sub(peer.Stats.ConnectionTime)
timeSinceActivity := now.Sub(peer.Stats.LastActivity)
quality := "good"
// Determine quality based on multiple factors
if timeSinceActivity > 5*time.Minute {
quality = "poor"
} else if timeSinceActivity > 2*time.Minute {
quality = "fair"
} else if connectionAge > 30*time.Minute && peer.Stats.BytesUploaded > 0 {
quality = "excellent"
} else if len(peer.WebRTCPeers) > 0 {
quality = "good"
}
peer.Stats.ConnectionQuality = quality
}
// handleICE forwards ICE candidates between peers
func (wt *WebSocketTracker) handleICE(conn *websocket.Conn, msg WebTorrentMessage) {
wt.mu.RLock()
defer wt.mu.RUnlock()
if swarm := wt.swarms[msg.InfoHash]; swarm != nil {
swarm.mu.RLock()
defer swarm.mu.RUnlock()
if targetPeer := swarm.peers[msg.ToPeerID]; targetPeer != nil {
// Forward ICE candidate to target peer
iceMsg := map[string]interface{}{
"action": "ice",
"info_hash": msg.InfoHash,
"peer_id": msg.FromPeerID,
"candidate": msg.Candidate,
"from_peer_id": msg.FromPeerID,
"to_peer_id": msg.ToPeerID,
}
if err := targetPeer.Conn.WriteJSON(iceMsg); err != nil {
log.Printf("Failed to forward ICE candidate: %v", err)
targetPeer.State = StateDisconnected
// Track ICE failure
wt.statsTracker.mu.Lock()
wt.statsTracker.iceFailures++
wt.statsTracker.mu.Unlock()
} else {
log.Printf("Forwarded ICE candidate from %s to %s for %s",
msg.FromPeerID, msg.ToPeerID, msg.InfoHash[:8])
}
} else {
log.Printf("Target peer %s not found for ICE candidate", msg.ToPeerID)
}
}
}
// handleConnectionState updates peer connection states
func (wt *WebSocketTracker) handleConnectionState(conn *websocket.Conn, msg WebTorrentMessage) {
wt.mu.Lock()
defer wt.mu.Unlock()
if swarm := wt.swarms[msg.InfoHash]; swarm != nil {
swarm.mu.Lock()
defer swarm.mu.Unlock()
if peer := swarm.peers[msg.PeerID]; peer != nil {
oldState := peer.State
newState := PeerConnectionState(msg.ConnectionState)
peer.State = newState
peer.LastSeen = time.Now()
log.Printf("Peer %s connection state changed from %s to %s for %s",
msg.PeerID, oldState, newState, msg.InfoHash[:8])
// Update stats based on state change
wt.statsTracker.mu.Lock()
if oldState != StateConnected && newState == StateConnected {
wt.statsTracker.activeConnections++
} else if oldState == StateConnected && newState != StateConnected {
wt.statsTracker.activeConnections--
if newState == StateFailed {
wt.statsTracker.connectionFailures++
}
}
wt.statsTracker.mu.Unlock()
// If peer disconnected, remove from WebRTC connections
if newState == StateDisconnected || newState == StateFailed {
wt.removePeerConnections(swarm, msg.PeerID)
}
}
}
}
// handleStats processes peer statistics updates
func (wt *WebSocketTracker) handleStats(conn *websocket.Conn, msg WebTorrentMessage) {
wt.mu.Lock()
defer wt.mu.Unlock()
if swarm := wt.swarms[msg.InfoHash]; swarm != nil {
swarm.mu.Lock()
defer swarm.mu.Unlock()
if peer := swarm.peers[msg.PeerID]; peer != nil && peer.Stats != nil {
oldUploaded := peer.Stats.BytesUploaded
oldDownloaded := peer.Stats.BytesDownloaded
// Update peer stats
peer.Stats.BytesUploaded = msg.Uploaded
peer.Stats.BytesDownloaded = msg.Downloaded
peer.Stats.LastActivity = time.Now()
// Update global stats
wt.statsTracker.mu.Lock()
wt.statsTracker.totalBytesUploaded += (msg.Uploaded - oldUploaded)
wt.statsTracker.totalBytesDownloaded += (msg.Downloaded - oldDownloaded)
wt.statsTracker.mu.Unlock()
// Report to main tracker periodically
if wt.tracker != nil && time.Since(wt.statsTracker.lastReported) > 5*time.Minute {
go wt.reportStatsToTracker()
}
log.Printf("Updated stats for peer %s: %d uploaded, %d downloaded",
msg.PeerID, msg.Uploaded, msg.Downloaded)
}
}
}
// Helper methods for peer management
func (wt *WebSocketTracker) removePeer(swarm *Swarm, peerID string) {
if peer := swarm.peers[peerID]; peer != nil {
peer.Conn.Close()
delete(swarm.peers, peerID)
wt.statsTracker.mu.Lock()
wt.statsTracker.activeConnections--
wt.statsTracker.mu.Unlock()
log.Printf("Removed peer %s", peerID)
}
}
func (wt *WebSocketTracker) removePeerConnections(swarm *Swarm, peerID string) {
// Remove peer from other peers' connection maps
for _, peer := range swarm.peers {
delete(peer.WebRTCPeers, peerID)
}
}
// HTTP fallback functionality
func (wt *WebSocketTracker) shouldProvideHTTPFallback(infoHash string) bool {
if wt.tracker == nil {
return false
}
// Check if WebRTC connections are failing for this torrent
wt.statsTracker.mu.RLock()
failureRate := float64(wt.statsTracker.connectionFailures) / float64(wt.statsTracker.totalConnections)
wt.statsTracker.mu.RUnlock()
return failureRate > 0.3 // If more than 30% of connections are failing
}
func (wt *WebSocketTracker) addHTTPFallback(response map[string]interface{}, infoHash string) {
if wt.tracker == nil {
return
}
// Get HTTP peers from main tracker
httpPeers, err := wt.tracker.GetPeersForTorrent(infoHash)
if err != nil {
log.Printf("Failed to get HTTP peers: %v", err)
return
}
// Add HTTP peer endpoints
var fallbackPeers []map[string]interface{}
for _, peer := range httpPeers {
if peer.IsWebSeed {
// Add WebSeed URLs - get URL from the tracker gateway if available
webSeedURL := fmt.Sprintf("http://localhost/webseed/%s", infoHash) // Fallback URL
fallbackPeers = append(fallbackPeers, map[string]interface{}{
"id": peer.PeerID,
"url": webSeedURL,
"protocol": "webseed",
})
} else {
// Add HTTP tracker peers
fallbackPeers = append(fallbackPeers, map[string]interface{}{
"id": peer.PeerID,
"endpoint": fmt.Sprintf("%s:%d", peer.IP, peer.Port),
"protocol": "http",
})
}
}
if len(fallbackPeers) > 0 {
response["http_fallback"] = fallbackPeers
response["supports_hybrid"] = true
log.Printf("Added %d HTTP fallback peers for %s", len(fallbackPeers), infoHash[:8])
}
}
// Stats reporting to main tracker
func (wt *WebSocketTracker) reportStatsToTracker() {
wt.statsTracker.mu.Lock()
stats := map[string]interface{}{
"webrtc_connections": wt.statsTracker.activeConnections,
"total_connections": wt.statsTracker.totalConnections,
"bytes_uploaded": wt.statsTracker.totalBytesUploaded,
"bytes_downloaded": wt.statsTracker.totalBytesDownloaded,
"connection_failures": wt.statsTracker.connectionFailures,
"ice_failures": wt.statsTracker.iceFailures,
"timestamp": time.Now(),
}
wt.statsTracker.lastReported = time.Now()
wt.statsTracker.mu.Unlock()
// Report to main tracker (placeholder - would integrate with tracker's stats system)
log.Printf("WebRTC Stats: %+v", stats)
}
// GetStats returns comprehensive WebSocket tracker statistics
func (wt *WebSocketTracker) GetStats() map[string]interface{} {
wt.mu.RLock()
defer wt.mu.RUnlock()
totalPeers := 0
connectedPeers := 0
seeders := 0
leechers := 0
totalSwarms := len(wt.swarms)
for _, swarm := range wt.swarms {
swarm.mu.RLock()
for _, peer := range swarm.peers {
totalPeers++
if peer.State == StateConnected {
connectedPeers++
if peer.IsSeeder {
seeders++
} else {
leechers++
}
}
}
swarm.mu.RUnlock()
}
wt.statsTracker.mu.RLock()
statsData := map[string]interface{}{
"total_swarms": totalSwarms,
"total_peers": totalPeers,
"connected_peers": connectedPeers,
"seeders": seeders,
"leechers": leechers,
"total_connections": wt.statsTracker.totalConnections,
"active_connections": wt.statsTracker.activeConnections,
"connection_failures": wt.statsTracker.connectionFailures,
"ice_failures": wt.statsTracker.iceFailures,
"total_bytes_uploaded": wt.statsTracker.totalBytesUploaded,
"total_bytes_downloaded": wt.statsTracker.totalBytesDownloaded,
"last_stats_report": wt.statsTracker.lastReported,
"status": "active",
}
wt.statsTracker.mu.RUnlock()
return statsData
}