torrent-gateway/internal/tracker/websocket.go

package tracker

import (
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

type WebSocketTracker struct {
	upgrader websocket.Upgrader
	swarms   map[string]*Swarm
	mu       sync.RWMutex
}

type Swarm struct {
	peers map[string]*WebRTCPeer
	mu    sync.RWMutex
}

type WebRTCPeer struct {
	ID         string          `json:"peer_id"`
	Conn       *websocket.Conn `json:"-"`
	LastSeen   time.Time       `json:"last_seen"`
	InfoHashes []string        `json:"info_hashes"`
}

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"`
}

func NewWebSocketTracker() *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),
	}
}

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)
		}
	}
}

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()

	// Add/update peer
	peer := &WebRTCPeer{
		ID:         msg.PeerID,
		Conn:       conn,
		LastSeen:   time.Now(),
		InfoHashes: []string{msg.InfoHash},
	}
	swarm.peers[msg.PeerID] = peer

	// Return peer list (excluding the requesting peer)
	var peers []map[string]interface{}
	numWant := msg.NumWant
	if numWant == 0 {
		numWant = 30 // Default
	}

	count := 0
	for peerID := range swarm.peers {
		if peerID != msg.PeerID && count < numWant {
			peers = append(peers, map[string]interface{}{
				"id": peerID,
			})
			count++
		}
	}

	response := map[string]interface{}{
		"action":    "announce",
		"interval":  300, // 5 minutes for WebTorrent
		"info_hash": msg.InfoHash,
		"complete":  len(swarm.peers), // Simplified
		"incomplete": 0,
		"peers":     peers,
	}

	if err := conn.WriteJSON(response); err != nil {
		log.Printf("Failed to send announce response: %v", err)
	}
}

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)
			}
		}
		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)
			}
		}
		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

	for infoHash, swarm := range wt.swarms {
		swarm.mu.Lock()
		for peerID, peer := range swarm.peers {
			if peer.LastSeen.Before(expiry) {
				peer.Conn.Close()
				delete(swarm.peers, peerID)
			}
		}
		
		// Remove empty swarms
		if len(swarm.peers) == 0 {
			delete(wt.swarms, infoHash)
		}
		swarm.mu.Unlock()
	}
}