JustAMassenger/server/internal/federation/federation.go

package federation

import (
	"bytes"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"sync"
	"time"

	jcrypto "github.com/justamessenger/server/internal/crypto"
	"github.com/justamessenger/server/internal/database"
	"github.com/justamessenger/server/internal/models"
	"github.com/justamessenger/server/internal/protocol"
)

type Manager struct {
	db          *database.DB
	domain      string
	privateKey  ed25519.PrivateKey
	publicKey   ed25519.PublicKey
	peers       map[string]*models.FederationPeer
	mu          sync.RWMutex
	httpClient  *http.Client
}

func NewManager(db *database.DB, domain string) (*Manager, error) {
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("failed to generate federation keypair: %w", err)
	}

	m := &Manager{
		db:         db,
		domain:     domain,
		privateKey: priv,
		publicKey:  pub,
		peers:      make(map[string]*models.FederationPeer),
		httpClient: &http.Client{
			Timeout: 30 * time.Second,
			Transport: &http.Transport{
				MaxIdleConns:    20,
				IdleConnTimeout: 90 * time.Second,
			},
		},
	}

	if err := m.loadPeers(); err != nil {
		log.Printf("Warning: failed to load federation peers: %v", err)
	}

	return m, nil
}

func (m *Manager) PublicKey() []byte {
	return m.publicKey
}

func (m *Manager) loadPeers() error {
	rows, err := m.db.Query(
		`SELECT domain, name, public_key, last_seen, is_active
		 FROM federation_peers WHERE is_active = 1`,
	)
	if err != nil {
		return err
	}
	defer rows.Close()

	m.mu.Lock()
	defer m.mu.Unlock()

	for rows.Next() {
		p := &models.FederationPeer{}
		if err := rows.Scan(&p.Domain, &p.Name, &p.PublicKey, &p.LastSeen, &p.IsActive); err != nil {
			return err
		}
		m.peers[p.Domain] = p
	}
	return nil
}

func (m *Manager) AddPeer(domain, name string, publicKey []byte) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	peer := &models.FederationPeer{
		Domain:    domain,
		Name:      name,
		PublicKey: publicKey,
		LastSeen:  time.Now(),
		IsActive:  true,
	}
	m.peers[domain] = peer

	_, err := m.db.Exec(
		`INSERT OR REPLACE INTO federation_peers (domain, name, public_key, last_seen, is_active)
		 VALUES (?, ?, ?, ?, 1)`,
		domain, name, publicKey, time.Now(),
	)
	return err
}

func (m *Manager) Sign(data []byte) []byte {
	return ed25519.Sign(m.privateKey, data)
}

func (m *Manager) Verify(publicKey, data, signature []byte) bool {
	return ed25519.Verify(publicKey, data, signature)
}

func (m *Manager) SendToPeer(peerDomain string, pkt *protocol.FederationData) error {
	m.mu.RLock()
	peer, ok := m.peers[peerDomain]
	m.mu.RUnlock()

	if !ok {
		return fmt.Errorf("unknown peer: %s", peerDomain)
	}

	payload, err := json.Marshal(pkt)
	if err != nil {
		return err
	}

	url := fmt.Sprintf("https://%s/federation/receive", peer.Domain)
	req, err := http.NewRequest("POST", url, bytes.NewReader(payload))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-JAM-Domain", m.domain)
	req.Header.Set("X-JAM-Signature", string(m.Sign(payload)))

	resp, err := m.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("federation request to %s failed: %w", peer.Domain, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("peer %s returned %d: %s", peer.Domain, resp.StatusCode, string(body))
	}

	m.mu.Lock()
	peer.LastSeen = time.Now()
	m.mu.Unlock()

	return nil
}

func (m *Manager) BroadcastToPeers(pkt *protocol.FederationData) {
	m.mu.RLock()
	domains := make([]string, 0, len(m.peers))
	for domain := range m.peers {
		domains = append(domains, domain)
	}
	m.mu.RUnlock()

	for _, domain := range domains {
		if domain == m.domain {
			continue
		}
		if err := m.SendToPeer(domain, pkt); err != nil {
			log.Printf("Federation broadcast to %s failed: %v", domain, err)
		}
	}
}

func (m *Manager) HandleReceive(w http.ResponseWriter, r *http.Request) {
	signature := r.Header.Get("X-JAM-Signature")

	body, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "failed to read body", http.StatusBadRequest)
		return
	}

	var pkt protocol.FederationData
	if err := json.Unmarshal(body, &pkt); err != nil {
		http.Error(w, "invalid packet", http.StatusBadRequest)
		return
	}

	m.mu.RLock()
	peer, ok := m.peers[pkt.FromDomain]
	m.mu.RUnlock()

	if ok && !m.Verify(peer.PublicKey, body, []byte(signature)) {
		http.Error(w, "invalid signature", http.StatusUnauthorized)
		return
	}

	log.Printf("Federation packet from %s: type=%s", pkt.FromDomain, pkt.PacketType)
	w.WriteHeader(http.StatusOK)
	json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}

func (m *Manager) EncryptFederationPayload(payload []byte) ([]byte, []byte, error) {
	key, err := jcrypto.GenerateKey()
	if err != nil {
		return nil, nil, err
	}
	return jcrypto.Encrypt(payload, key)
}