Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
14 kB · 457 lines
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458package iroh_replicator
import ( "bytes" "context" "crypto/rand" "encoding/json" "fmt" "reflect" "strings" "sync" "time"
"github.com/bluesky-social/indigo/util" "golang.org/x/sync/errgroup" "stream.place/streamplace/pkg/bus" "stream.place/streamplace/pkg/config" "stream.place/streamplace/pkg/iroh/generated/iroh_streamplace" "stream.place/streamplace/pkg/log" "stream.place/streamplace/pkg/media" "stream.place/streamplace/pkg/model" "stream.place/streamplace/pkg/placestream" "stream.place/streamplace/pkg/spmetrics")
type IrohSwarm struct { Node *iroh_streamplace.Node DB *iroh_streamplace.Db w *iroh_streamplace.WriteScope mm *media.MediaManager segChan chan *media.NewSegmentNotification NodeID string NodeTicket string bus *bus.Bus originMutex sync.Mutex mod model.Model cli *config.CLI activeSubs map[string]*SwarmOriginInfo handleDataScoped func(pubKey *iroh_streamplace.PublicKey, topic string, data []byte)}
// A message saying "hey I ingested node data at this time"type SwarmOriginInfo struct { Type string `json:"$type"` NodeID string `json:"node_id"` Time string `json:"time"` Streamer string `json:"streamer"`}
type SwarmViewerCount struct { Type string `json:"$type"` Server string `json:"server"` Streamer string `json:"streamer"` Viewers int `json:"viewers"`}
func NewSwarm(ctx context.Context, cli *config.CLI, secret []byte, topic []byte, mm *media.MediaManager, bus *bus.Bus, mod model.Model) (*IrohSwarm, error) { ctx = log.WithLogValues(ctx, "func", "StartKV")
if topic == nil { topic = make([]byte, 32) _, err := rand.Read(topic) if err != nil { return nil, fmt.Errorf("failed to generate random topic: %w", err) } }
log.Log(ctx, "Starting with tickets", "tickets", cli.Tickets) config := iroh_streamplace.Config{ Key: secret, Topic: topic, MaxSendDuration: 1000_000_000, // 1s DisableRelay: cli.DisableIrohRelay, } log.Log(ctx, "Config created", "config", config)
swarm := IrohSwarm{ mm: mm, activeSubs: make(map[string]*SwarmOriginInfo), bus: bus, mod: mod, cli: cli, }
// workaround to get context into the HandleData callback swarm.handleDataScoped = func(_ *iroh_streamplace.PublicKey, topic string, data []byte) { if ctx.Err() != nil { return } err := swarm.mm.ValidateMP4(context.Background(), bytes.NewReader(data), false) if err != nil { log.Error(ctx, "could not validate segment", "error", err, "topic", topic, "data", len(data)) } }
node, err := iroh_streamplace.NodeReceiver(config, &swarm) if err != nil { return nil, fmt.Errorf("failed to create NodeSender: %w", err) }
db := node.Db() w := node.NodeScope()
swarm.DB = db swarm.w = w swarm.Node = node
nodeId, err := node.NodeId() if err != nil { return nil, fmt.Errorf("failed to get NodeId: %w", err) } log.Log(ctx, "Node ID:", "node_id", nodeId) swarm.NodeID = nodeId.String()
ticket, err := node.Ticket() if err != nil { return nil, fmt.Errorf("failed to get Ticket: %w", err) } swarm.NodeTicket = ticket
return &swarm, nil}
func (swarm *IrohSwarm) BuildOriginRecord(origin *placestream.BroadcastOrigin) error { origin.IrohTicket = &swarm.NodeTicket return nil}
func (swarm *IrohSwarm) Start(ctx context.Context, cli *config.CLI) error { if len(cli.Tickets) > 0 { err := swarm.Node.JoinPeers(cli.Tickets) if err != nil { return fmt.Errorf("failed to join peers: %w", err) } } nodeId, err := swarm.Node.NodeId() if err != nil { return fmt.Errorf("failed to get node id: %w", err) } nodeIdStr := nodeId.String() log.Log(ctx, "Node ID:", "node_id", nodeIdStr)
g, ctx := errgroup.WithContext(ctx) g.Go(func() error { return swarm.startKV(ctx) }) g.Go(func() error { return swarm.startSegmentSender(ctx) }) g.Go(func() error { <-ctx.Done() return swarm.Node.Shutdown() }) g.Go(func() error { return swarm.startBusSubscribe(ctx) }) g.Go(func() error { return swarm.startViewerCountSubscribe(ctx) }) return g.Wait()}
func (swarm *IrohSwarm) startKV(ctx context.Context) error { sub := swarm.DB.Subscribe(iroh_streamplace.NewFilter()) for { if ctx.Err() != nil { return ctx.Err() } ev, err := sub.NextRaw() if err != nil { return fmt.Errorf("failed to get next subscription event: %w", err) }
if ev == nil { log.Warn(ctx, "Got empty event from sub.NextRaw(), pausing for a second and continuing") time.Sleep(1 * time.Second) continue }
switch item := (*ev).(type) { case iroh_streamplace.SubscribeItemEntry: err := swarm.handleIrohMessage(ctx, item) if err != nil { log.Error(ctx, "could not handle iroh message", "error", err) continue }
case iroh_streamplace.SubscribeItemCurrentDone: log.Debug(ctx, "SubscribeItemCurrentDone", "currentDone", item) case iroh_streamplace.SubscribeItemExpired: log.Debug(ctx, "SubscribeItemExpired", "expired", item) case iroh_streamplace.SubscribeItemOther: log.Debug(ctx, "SubscribeItemOther", "other", item) } }}
func (swarm *IrohSwarm) handleIrohMessage(ctx context.Context, item iroh_streamplace.SubscribeItemEntry) error { keyStr := string(item.Key) valueStr := string(item.Value) log.Debug(ctx, "SubscribeItemEntry", "key", keyStr, "value", valueStr) if len(valueStr) > 0 && valueStr[0] != '{' { // not JSON, it's one of the rust messages log.Debug(ctx, "not JSON", "key", keyStr, "value", valueStr) return nil } rawMessage, err := decodeIrohMessage(item.Key, item.Value) if err != nil { return fmt.Errorf("could not decode iroh message: %w", err) } switch message := rawMessage.(type) { case SwarmOriginInfo: err = swarm.checkOrigins(ctx, message.Streamer, message.NodeID) if err != nil { return fmt.Errorf("could not check origins: %w", err) } case SwarmViewerCount: log.Log(ctx, "got viewer count", "viewerCount", message) if message.Server == swarm.NodeID { // no infinite loops allowed return nil } swarm.bus.SetViewerCount(message.Streamer, message.Server, message.Viewers) log.Log(ctx, "set viewer count", "viewerCount", message) return nil default: return fmt.Errorf("unknown message type: %s", reflect.TypeOf(rawMessage)) } return nil}
func decodeIrohMessage(key, value []byte) (any, error) { keyStr := string(key) if strings.HasPrefix(keyStr, "origin::") { var originInfo SwarmOriginInfo err := json.Unmarshal(value, &originInfo) if err != nil { return nil, fmt.Errorf("could not unmarshal origin info: %w", err) } return originInfo, nil } if strings.HasPrefix(keyStr, "viewers::") { var viewerCount SwarmViewerCount err := json.Unmarshal(value, &viewerCount) if err != nil { return nil, fmt.Errorf("could not unmarshal viewer count: %w", err) } return viewerCount, nil } return nil, fmt.Errorf("unknown key: %s", keyStr)}
// subscribe to all streamsfunc (swarm *IrohSwarm) startBusSubscribe(ctx context.Context) error { // start subscription first so we're buffering new origins busCh := swarm.bus.Subscribe("") originViews, err := swarm.mod.GetRecentBroadcastOrigins(ctx) if err != nil { return fmt.Errorf("failed to get recent broadcast origins: %w", err) } for _, view := range originViews { err = swarm.handleOriginMessage(ctx, &view) if err != nil { log.Error(ctx, "could not check origin", "error", err) } } log.Log(ctx, "Resumed recent broadcast origins", "count", len(originViews)) for { select { case <-ctx.Done(): return ctx.Err() case msg := <-busCh: if view, ok := msg.(*placestream.BroadcastDefs_BroadcastOriginView); ok { log.Debug(ctx, "got broadcast origin view", "view", view) err = swarm.handleOriginMessage(ctx, view) if err != nil { log.Error(ctx, "could not handle origin message", "error", err) } } } }}
func (swarm *IrohSwarm) startViewerCountSubscribe(ctx context.Context) error { ch := swarm.bus.SubscribeToViewerCount() for { select { case <-ctx.Done(): return ctx.Err() case msg := <-ch: log.Log(ctx, "got viewer count update", "viewerCount", msg) if msg.Origin != "local" { continue } swarmMsg := SwarmViewerCount{ Type: "place.stream.swarm.viewerCount", Server: swarm.NodeID, Streamer: msg.Streamer, Viewers: msg.Count, } bs, err := json.Marshal(swarmMsg) if err != nil { log.Error(ctx, "could not marshal viewer count", "error", err) continue } key := fmt.Sprintf("viewers::%s::%s", swarm.NodeID, msg.Streamer) err = swarm.w.Put(nil, []byte(key), bs) if err != nil { log.Error(ctx, "could not put viewer count to swarm", "error", err) continue } log.Log(ctx, "put viewer count to swarm", "viewerCount", msg)
} }}
func (swarm *IrohSwarm) handleOriginMessage(ctx context.Context, view *placestream.BroadcastDefs_BroadcastOriginView) error { origin, ok := view.Record.Val.(*placestream.BroadcastOrigin) if !ok { return fmt.Errorf("record is not a BroadcastOrigin") } if view.Author.Did != origin.Streamer { // currently, only streamers are allowed to advertise origins return nil } if origin.IrohTicket == nil { return fmt.Errorf("origin has no iroh ticket") } pubKey, err := iroh_streamplace.NodeIdFromTicket(*origin.IrohTicket) if err != nil { return fmt.Errorf("could not get node id from ticket: %w", err) } err = swarm.Node.AddTickets([]string{*origin.IrohTicket}) if err != nil { return fmt.Errorf("could not add tickets: %w", err) } pubKeyStr := pubKey.String() err = swarm.checkOrigins(ctx, origin.Streamer, pubKeyStr) if err != nil { return fmt.Errorf("could not check origin: %w", err) } return nil}
func (swarm *IrohSwarm) checkOrigins(ctx context.Context, streamer string, nodeID string) error { ctx = log.WithLogValues(ctx, "streamer", streamer, "nodeID", nodeID, "func", "checkOrigins") err := swarm.cli.StreamIsAllowed(streamer) if err != nil { return fmt.Errorf("user %s is not allowlisted on this node: %w", streamer, err) } swarm.originMutex.Lock() defer swarm.originMutex.Unlock() oldSub, ok := swarm.activeSubs[streamer] if ok { if oldSub.NodeID == nodeID { log.Debug(ctx, "node hasn't changed", "streamer", streamer) // mmyep. same node still has the stream. great news. return nil } log.Log(ctx, "Stream origin changed, swapping to new node", "old_node", oldSub.NodeID, "new_node", nodeID, "streamer", streamer) pubKey, err := iroh_streamplace.PublicKeyFromString(oldSub.NodeID) if err != nil { log.Error(ctx, "could not create public key", "error", err) return err } // different node has the stream. we need to unsubscribe from the old node. err = swarm.Node.Unsubscribe(streamer, pubKey) if err != nil { log.Error(ctx, "could not unsubscribe from key", "error", err) return err } delete(swarm.activeSubs, streamer) } if nodeID == swarm.NodeID { log.Debug(ctx, "I already have this stream", "streamer", streamer) // oh, i have this stream. cool. do nothing. return nil } log.Log(ctx, "Subscribing to stream start", "new_node", nodeID, "streamer", streamer) pubKey, err := iroh_streamplace.PublicKeyFromString(nodeID) if err != nil { log.Error(ctx, "could not create public key", "error", err) return err } err = swarm.Node.Subscribe(streamer, pubKey) log.Log(ctx, "Subscribing to stream done", "new_node", nodeID, "streamer", streamer, "pubKey", pubKey, "error", err) if err != nil { log.Error(ctx, "could not subscribe to key", "error", err) return err } swarm.activeSubs[streamer] = &SwarmOriginInfo{ Type: "place.stream.swarm.originInfo", NodeID: nodeID, Time: time.Now().Format(util.ISO8601), Streamer: streamer, } return nil}
func (swarm *IrohSwarm) startSegmentSender(ctx context.Context) error { ch := swarm.mm.NewSegment() for { select { case <-ctx.Done(): return ctx.Err() case not := <-ch: err := swarm.SendSegment(ctx, not) if err != nil { log.Error(ctx, "could not send segment to swarm", "error", err) } continue } }}
func (swarm *IrohSwarm) HandleData(pubKey *iroh_streamplace.PublicKey, topic string, data []byte) { swarm.handleDataScoped(pubKey, topic, data)}
func (swarm *IrohSwarm) SendSegment(ctx context.Context, not *media.NewSegmentNotification) error { if !not.Local { return nil } originInfo := SwarmOriginInfo{ Type: "place.stream.swarm.originInfo", NodeID: swarm.NodeID, Time: not.Segment.StartTime.Format(util.ISO8601), Streamer: not.Segment.RepoDID, } bs, err := json.Marshal(originInfo) if err != nil { log.Error(ctx, "could not marshal origin info", "error", err) return err } keyBs := []byte(fmt.Sprintf("origin::%s", not.Segment.RepoDID)) go func() { spmetrics.SwarmPutCalls.WithLabelValues(not.Segment.RepoDID).Inc() defer spmetrics.SwarmPutCalls.WithLabelValues(not.Segment.RepoDID).Dec() err = swarm.w.Put(nil, keyBs, bs) if err != nil { log.Error(ctx, "could not put segment to swarm", "error", err) } }() go func() { spmetrics.SendSegmentCalls.Inc() defer spmetrics.SendSegmentCalls.Dec() // Ship the bare canonical MUXL segment (blindly concatenatable, no // presentation header) — smaller than the flat MP4 and the canonical // form. The receiving node re-validates it verbatim via ValidateMP4. err = swarm.Node.SendSegment(not.Segment.RepoDID, not.Muxl) if err != nil { log.Error(ctx, "could not send segment to swarm", "error", err) } }() return nil}