Initial commit: Null DRM Official
Capture, decrypt, and restream toolkit with compiled-in app modules (RTE, TG4, BBC), on-device MITM proxy, streamd control plane, and www. BBC module.yaml is published (clear streams); other module values stay local.
This commit is contained in:
404errordeveloper
commit
2fa8f2435f
121 changed files with 17802 additions and 0 deletions
350
apps/streamd/internal/ws/hub.go
Normal file
350
apps/streamd/internal/ws/hub.go
Normal file
|
|
@ -0,0 +1,350 @@
|
|||
package ws
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
|
||||
"drmdecryption/apps/streamd/internal/db"
|
||||
)
|
||||
|
||||
const (
|
||||
agentHeartbeatEvery = 20 * time.Second
|
||||
agentStaleAfter = 45 * time.Second
|
||||
claimRenewTTL = 90 * time.Second
|
||||
writeWait = 10 * time.Second
|
||||
pongWait = 60 * time.Second
|
||||
)
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool { return true },
|
||||
}
|
||||
|
||||
// Hub tracks agent WebSocket sessions and UI listeners.
|
||||
type Hub struct {
|
||||
Store *db.Store
|
||||
Token string
|
||||
Log *log.Logger
|
||||
|
||||
mu sync.Mutex
|
||||
agents map[string]*agentConn // agent_id -> conn
|
||||
uis map[*websocket.Conn]struct{}
|
||||
}
|
||||
|
||||
type agentConn struct {
|
||||
conn *websocket.Conn
|
||||
agentID string
|
||||
lastSeen time.Time
|
||||
send chan []byte
|
||||
}
|
||||
|
||||
type inbound struct {
|
||||
Type string `json:"type"`
|
||||
AgentID string `json:"agent_id,omitempty"`
|
||||
}
|
||||
|
||||
type outbound struct {
|
||||
Type string `json:"type"`
|
||||
AgentID string `json:"agent_id,omitempty"`
|
||||
OK bool `json:"ok,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
Agents []any `json:"agents,omitempty"`
|
||||
Message string `json:"message,omitempty"`
|
||||
}
|
||||
|
||||
func NewHub(store *db.Store, token string, lg *log.Logger) *Hub {
|
||||
if lg == nil {
|
||||
lg = log.Default()
|
||||
}
|
||||
return &Hub{
|
||||
Store: store,
|
||||
Token: token,
|
||||
Log: lg,
|
||||
agents: map[string]*agentConn{},
|
||||
uis: map[*websocket.Conn]struct{}{},
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) Mount(mux *http.ServeMux) {
|
||||
mux.HandleFunc("/ws/agent", h.handleAgent)
|
||||
mux.HandleFunc("/ws/ui", h.handleUI)
|
||||
mux.HandleFunc("/api/agents", h.handleAgentsAPI)
|
||||
}
|
||||
|
||||
// RunExpireLoop periodically drops expired claims and dead agent sessions.
|
||||
func (h *Hub) RunExpireLoop(stop <-chan struct{}) {
|
||||
t := time.NewTicker(10 * time.Second)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
case <-t.C:
|
||||
h.tick()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) tick() {
|
||||
n, err := h.Store.ExpireClaims()
|
||||
if err != nil {
|
||||
h.Log.Printf("ws: expire claims: %v", err)
|
||||
} else if n > 0 {
|
||||
h.Log.Printf("ws: expired %d stale claim(s)", n)
|
||||
h.broadcastUI(outbound{Type: "claims_changed", Message: "expired"})
|
||||
}
|
||||
|
||||
now := time.Now().UTC()
|
||||
var stale []string
|
||||
h.mu.Lock()
|
||||
for id, a := range h.agents {
|
||||
if now.Sub(a.lastSeen) > agentStaleAfter {
|
||||
stale = append(stale, id)
|
||||
}
|
||||
}
|
||||
h.mu.Unlock()
|
||||
for _, id := range stale {
|
||||
h.Log.Printf("ws: agent %s heartbeat stale — releasing claims", id)
|
||||
h.dropAgent(id, true)
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) authToken(r *http.Request) bool {
|
||||
if h.Token == "" {
|
||||
return true
|
||||
}
|
||||
tok := r.URL.Query().Get("token")
|
||||
if tok == "" {
|
||||
auth := r.Header.Get("Authorization")
|
||||
tok = strings.TrimPrefix(auth, "Bearer ")
|
||||
}
|
||||
if tok == "" {
|
||||
tok = r.Header.Get("X-Streamd-Token")
|
||||
}
|
||||
return tok == h.Token
|
||||
}
|
||||
|
||||
func (h *Hub) handleAgentsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
http.Error(w, `{"error":"method not allowed"}`, http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]any{"agents": h.snapshotAgents()})
|
||||
}
|
||||
|
||||
func (h *Hub) handleUI(w http.ResponseWriter, r *http.Request) {
|
||||
conn, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
h.uis[conn] = struct{}{}
|
||||
h.mu.Unlock()
|
||||
defer func() {
|
||||
h.mu.Lock()
|
||||
delete(h.uis, conn)
|
||||
h.mu.Unlock()
|
||||
_ = conn.Close()
|
||||
}()
|
||||
|
||||
_ = conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
conn.SetPongHandler(func(string) error {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
return nil
|
||||
})
|
||||
// Push initial agent list.
|
||||
h.sendJSON(conn, outbound{Type: "agents", Agents: h.snapshotAgents()})
|
||||
for {
|
||||
if _, _, err := conn.ReadMessage(); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) handleAgent(w http.ResponseWriter, r *http.Request) {
|
||||
if !h.authToken(r) {
|
||||
http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
agentID := strings.TrimSpace(r.URL.Query().Get("agent_id"))
|
||||
if agentID == "" {
|
||||
http.Error(w, `{"error":"agent_id required"}`, http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
conn, err := upgrader.Upgrade(w, r, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
ac := &agentConn{
|
||||
conn: conn,
|
||||
agentID: agentID,
|
||||
lastSeen: time.Now().UTC(),
|
||||
send: make(chan []byte, 8),
|
||||
}
|
||||
h.mu.Lock()
|
||||
if old, ok := h.agents[agentID]; ok {
|
||||
close(old.send)
|
||||
_ = old.conn.Close()
|
||||
}
|
||||
h.agents[agentID] = ac
|
||||
h.mu.Unlock()
|
||||
h.Log.Printf("ws: agent connected id=%s", agentID)
|
||||
h.broadcastUI(outbound{Type: "agent_online", AgentID: agentID, Agents: h.snapshotAgents()})
|
||||
|
||||
go h.agentWriter(ac)
|
||||
|
||||
defer h.dropAgent(agentID, true)
|
||||
|
||||
_ = conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
conn.SetPongHandler(func(string) error {
|
||||
_ = conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
return nil
|
||||
})
|
||||
|
||||
// Greeting + first renew.
|
||||
h.renew(agentID)
|
||||
h.queue(ac, outbound{Type: "hello", AgentID: agentID, OK: true})
|
||||
|
||||
for {
|
||||
_, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_ = conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
var msg inbound
|
||||
if err := json.Unmarshal(data, &msg); err != nil {
|
||||
continue
|
||||
}
|
||||
switch strings.ToLower(msg.Type) {
|
||||
case "heartbeat", "ping", "hello":
|
||||
h.mu.Lock()
|
||||
if cur, ok := h.agents[agentID]; ok {
|
||||
cur.lastSeen = time.Now().UTC()
|
||||
}
|
||||
h.mu.Unlock()
|
||||
n, err := h.renew(agentID)
|
||||
if err != nil {
|
||||
h.queue(ac, outbound{Type: "heartbeat_ack", OK: false, Error: err.Error()})
|
||||
continue
|
||||
}
|
||||
h.queue(ac, outbound{Type: "heartbeat_ack", OK: true, Message: "renewed", AgentID: agentID})
|
||||
if n > 0 {
|
||||
h.broadcastUI(outbound{Type: "claims_changed", AgentID: agentID, Message: "renewed"})
|
||||
}
|
||||
case "release_all":
|
||||
_, _ = h.Store.ReleaseAgentClaims(agentID)
|
||||
h.broadcastUI(outbound{Type: "claims_changed", AgentID: agentID, Message: "released"})
|
||||
h.queue(ac, outbound{Type: "release_ack", OK: true})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) renew(agentID string) (int64, error) {
|
||||
return h.Store.RenewAgentClaims(agentID, claimRenewTTL)
|
||||
}
|
||||
|
||||
func (h *Hub) dropAgent(agentID string, releaseClaims bool) {
|
||||
h.mu.Lock()
|
||||
ac, ok := h.agents[agentID]
|
||||
if ok {
|
||||
delete(h.agents, agentID)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
close(ac.send)
|
||||
_ = ac.conn.Close()
|
||||
if releaseClaims {
|
||||
n, err := h.Store.ReleaseAgentClaims(agentID)
|
||||
if err != nil {
|
||||
h.Log.Printf("ws: release claims for %s: %v", agentID, err)
|
||||
} else if n > 0 {
|
||||
h.Log.Printf("ws: released %d claim(s) for disconnected agent %s", n, agentID)
|
||||
}
|
||||
}
|
||||
h.Log.Printf("ws: agent disconnected id=%s", agentID)
|
||||
h.broadcastUI(outbound{Type: "agent_offline", AgentID: agentID, Agents: h.snapshotAgents()})
|
||||
}
|
||||
|
||||
func (h *Hub) agentWriter(ac *agentConn) {
|
||||
ticker := time.NewTicker(agentHeartbeatEvery)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case msg, ok := <-ac.send:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
_ = ac.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
if err := ac.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
|
||||
return
|
||||
}
|
||||
case <-ticker.C:
|
||||
_ = ac.conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
if err := ac.conn.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(writeWait)); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) queue(ac *agentConn, msg outbound) {
|
||||
b, err := json.Marshal(msg)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case ac.send <- b:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) sendJSON(conn *websocket.Conn, msg outbound) {
|
||||
b, _ := json.Marshal(msg)
|
||||
_ = conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
_ = conn.WriteMessage(websocket.TextMessage, b)
|
||||
}
|
||||
|
||||
func (h *Hub) broadcastUI(msg outbound) {
|
||||
b, err := json.Marshal(msg)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
for conn := range h.uis {
|
||||
_ = conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
if err := conn.WriteMessage(websocket.TextMessage, b); err != nil {
|
||||
_ = conn.Close()
|
||||
delete(h.uis, conn)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Hub) snapshotAgents() []any {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
out := make([]any, 0, len(h.agents))
|
||||
now := time.Now().UTC()
|
||||
for id, a := range h.agents {
|
||||
out = append(out, map[string]any{
|
||||
"agent_id": id,
|
||||
"last_seen_s": now.Sub(a.lastSeen).Seconds(),
|
||||
"connected": true,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, v any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
enc := json.NewEncoder(w)
|
||||
enc.SetIndent("", " ")
|
||||
_ = enc.Encode(v)
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue