main.go raw

   1  // marmot-wasm — WASM module exposing the marmot MLS DM protocol to JS.
   2  // Compiled with: GOOS=js GOARCH=wasm go build -o marmot.wasm ./cmd/marmot-wasm
   3  // Loaded by the marmot service worker.
   4  package main
   5  
   6  import (
   7  	"context"
   8  	"fmt"
   9  	"os"
  10  	"sync"
  11  	"syscall/js"
  12  	"time"
  13  
  14  	"git.smesh.lol/orly/pkg/lol"
  15  	"git.smesh.lol/orly/pkg/nostr/encoders/event"
  16  	"git.smesh.lol/orly/pkg/nostr/encoders/filter"
  17  	"git.smesh.lol/orly/pkg/nostr/encoders/hex"
  18  	"git.smesh.lol/orly/pkg/nostr/protocol/marmot"
  19  	"git.smesh.lol/orly/pkg/version"
  20  )
  21  
  22  var (
  23  	client     *marmot.Client
  24  	crypto     *marmot.ProxyCrypto
  25  	store      *jsGroupStore
  26  	relay      *jsRelay
  27  	mu         sync.Mutex
  28  	statusFn   js.Value // onStatusFn callback — pushes status messages to page
  29  )
  30  
  31  // jsRelay bridges RelayConnection to JS callbacks.
  32  type jsRelay struct {
  33  	publishFn   js.Value // (eventJSON: string) => void
  34  	subscribeFn js.Value // (filterJSON: string) => int (subscription handle)
  35  	eventChs    map[int]chan *event.E
  36  	nextSub     int
  37  	mu          sync.Mutex
  38  	// Fix 2a: buffer events arriving before any subscription is registered.
  39  	preSubBuf    []*preSubEvent
  40  	preSubActive bool
  41  }
  42  
  43  type preSubEvent struct {
  44  	subID int
  45  	ev    *event.E
  46  }
  47  
  48  func (r *jsRelay) Publish(ctx context.Context, ev *event.E) error {
  49  	b, err := ev.MarshalJSON()
  50  	if err != nil {
  51  		return err
  52  	}
  53  	r.publishFn.Invoke(string(b))
  54  	return nil
  55  }
  56  
  57  func (r *jsRelay) Subscribe(ctx context.Context, ff *filter.S) (marmot.EventStream, error) {
  58  	b := ff.Marshal(nil)
  59  	r.mu.Lock()
  60  	id := r.nextSub
  61  	r.nextSub++
  62  	ch := make(chan *event.E, 64)
  63  	r.eventChs[id] = ch
  64  	// Fix 2a: drain pre-subscription buffer into the new channel.
  65  	if r.preSubActive {
  66  		r.preSubActive = false
  67  		for _, pe := range r.preSubBuf {
  68  			select {
  69  			case ch <- pe.ev:
  70  			default:
  71  			}
  72  		}
  73  		r.preSubBuf = nil
  74  	}
  75  	r.mu.Unlock()
  76  
  77  	r.subscribeFn.Invoke(id, string(b))
  78  	return &jsEventStream{id: id, ch: ch, relay: r}, nil
  79  }
  80  
  81  type jsEventStream struct {
  82  	id    int
  83  	ch    chan *event.E
  84  	relay *jsRelay
  85  }
  86  
  87  func (s *jsEventStream) Events() <-chan *event.E { return s.ch }
  88  func (s *jsEventStream) Close() {
  89  	s.relay.mu.Lock()
  90  	delete(s.relay.eventChs, s.id)
  91  	s.relay.mu.Unlock()
  92  }
  93  
  94  // deliverEvent routes an incoming event JSON to the right subscription channel.
  95  // Fix 2a: if no subscription channel exists yet, buffer the event for later delivery.
  96  func deliverEvent(subID int, evJSON string) {
  97  	ev := event.New()
  98  	if err := ev.UnmarshalJSON([]byte(evJSON)); err != nil {
  99  		return
 100  	}
 101  	relay.mu.Lock()
 102  	ch, ok := relay.eventChs[subID]
 103  	if !ok && relay.preSubActive {
 104  		if len(relay.preSubBuf) < 64 {
 105  			relay.preSubBuf = append(relay.preSubBuf, &preSubEvent{subID: subID, ev: ev})
 106  		}
 107  		relay.mu.Unlock()
 108  		return
 109  	}
 110  	relay.mu.Unlock()
 111  	if !ok {
 112  		return
 113  	}
 114  	select {
 115  	case ch <- ev:
 116  	default:
 117  	}
 118  }
 119  
 120  // safeFunc wraps a js.Func callback with recover so panics don't kill the WASM instance.
 121  func safeFunc(name string, fn func(this js.Value, args []js.Value) any) js.Func {
 122  	return js.FuncOf(func(this js.Value, args []js.Value) (ret any) {
 123  		defer func() {
 124  			if r := recover(); r != nil {
 125  				msg := fmt.Sprintf("[marmot-wasm] PANIC in %s: %v", name, r)
 126  				js.Global().Get("console").Call("error", msg)
 127  				ret = "error: panic: " + fmt.Sprint(r)
 128  			}
 129  		}()
 130  		return fn(this, args)
 131  	})
 132  }
 133  
 134  func main() {
 135  	// Suppress error-level logging from hex decoder etc — these fire on every
 136  	// corrupt event and burn CPU with formatted output. Errors still propagate
 137  	// via return values; they just don't print.
 138  	lol.Level.Store(lol.Off)
 139  
 140  	store = newJSGroupStore()
 141  	relay = &jsRelay{eventChs: make(map[int]chan *event.E), preSubActive: true}
 142  
 143  	// Register JS API — all wrapped with panic recovery
 144  	js.Global().Set("_marmot", js.ValueOf(map[string]any{
 145  		"init":            safeFunc("init", jsInit),
 146  		"sendDM":          safeFunc("sendDM", jsSendDM),
 147  		"subscribe":       safeFunc("subscribe", jsSubscribe),
 148  		"publishKP":       safeFunc("publishKP", jsPublishKP),
 149  		"listGroups":      safeFunc("listGroups", jsListGroups),
 150  		"handleEvent":     safeFunc("handleEvent", jsHandleEvent),
 151  		"deliverEvent":    safeFunc("deliverEvent", jsDeliverEvent),
 152  		"cryptoResult":    safeFunc("cryptoResult", jsCryptoResult),
 153  		"storeResult":     safeFunc("storeResult", jsStoreResult),
 154  		"keyPackageEvent": safeFunc("keyPackageEvent", jsKeyPackageEvent),
 155  		"lastEventTS":     safeFunc("lastEventTS", jsLastEventTS),
 156  		"backupGroups":    safeFunc("backupGroups", jsBackupGroups),
 157  		"restoreGroups":   safeFunc("restoreGroups", jsRestoreGroups),
 158  		"ratchetGroup":    safeFunc("ratchetGroup", jsRatchetGroup),
 159  		"version":         version.V,
 160  	}))
 161  
 162  
 163  	// Keep WASM alive
 164  	select {}
 165  }
 166  
 167  // jsInit(pubkeyHex, publishFn, subscribeFn, cryptoSendFn, onDMFn, onStatusFn, onReadyFn, lastEventTS, relayURLs[])
 168  // NewClient calls store.ListGroups/LoadGroup which block on async IDB callbacks.
 169  // Running it in a goroutine lets the JS event loop process those callbacks.
 170  func jsInit(this js.Value, args []js.Value) any {
 171  	if len(args) < 8 {
 172  		return "error: need pubkeyHex, publishFn, subscribeFn, cryptoSendFn, onDMFn, onStatusFn, onReadyFn, lastEventTS"
 173  	}
 174  	pubHex := args[0].String()
 175  	pubBytes, err := hex.Dec(pubHex)
 176  	if err != nil {
 177  		return "error: invalid pubkey: " + err.Error()
 178  	}
 179  
 180  	relay.publishFn = args[1]
 181  	relay.subscribeFn = args[2]
 182  	cryptoSendFn := args[3]
 183  	onDMFn := args[4]
 184  	onStatusFn := args[5]
 185  	onReadyFn := args[6]
 186  	lastEventTS := int64(args[7].Int())
 187  
 188  	var relays []string
 189  	if len(args) > 8 {
 190  		for i := 8; i < len(args); i++ {
 191  			relays = append(relays, args[i].String())
 192  		}
 193  	}
 194  
 195  	crypto = marmot.NewProxyCrypto(pubBytes, func(op, peerHex, data string, id int) {
 196  		cryptoSendFn.Invoke(op, peerHex, data, id)
 197  	})
 198  
 199  	// Run NewClient in a goroutine — it calls store.ListGroups() which blocks
 200  	// on a channel waiting for async IDB callbacks. Running synchronously would
 201  	// deadlock because Go WASM's single thread can't yield to the JS event loop.
 202  	go func() {
 203  		mu.Lock()
 204  		defer mu.Unlock()
 205  
 206  		c, err := marmot.NewClient(crypto, store, relay, relays...)
 207  		if err != nil {
 208  			onReadyFn.Invoke("error: " + err.Error())
 209  			return
 210  		}
 211  
 212  		if lastEventTS > 0 {
 213  			c.SetLastEventTS(lastEventTS)
 214  		}
 215  
 216  		c.OnDM(func(senderPub []byte, plaintext []byte) {
 217  			onDMFn.Invoke(hex.Enc(senderPub), string(plaintext))
 218  		})
 219  
 220  		statusFn = onStatusFn
 221  		client = c
 222  		onReadyFn.Invoke("ok")
 223  	}()
 224  
 225  	return nil
 226  }
 227  
 228  func sendStatus(msg string) {
 229  	if statusFn.IsUndefined() || statusFn.IsNull() {
 230  		return
 231  	}
 232  	statusFn.Invoke(msg)
 233  }
 234  
 235  func jsSendDM(this js.Value, args []js.Value) any {
 236  	if len(args) < 2 {
 237  		return "error: missing args"
 238  	}
 239  	if client == nil {
 240  		return "error: not initialized"
 241  	}
 242  	recipientHex := args[0].String()
 243  	content := args[1].String()
 244  	sendStatus("sendDM: starting for " + recipientHex[:8] + "...")
 245  
 246  	go func() {
 247  		recipientPub, err := hex.Dec(recipientHex)
 248  		if err != nil {
 249  			sendStatus("sendDM error: invalid recipient: " + err.Error())
 250  			return
 251  		}
 252  		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 253  		defer cancel()
 254  		if err := client.SendDM(ctx, recipientPub, []byte(content)); err != nil {
 255  			sendStatus("sendDM error: " + err.Error())
 256  			return
 257  		}
 258  		sendStatus("sendDM ok: sent to " + recipientHex[:8] + "...")
 259  	}()
 260  	return nil
 261  }
 262  
 263  // Fix 2c: track active subscription so duplicate calls cancel the previous one.
 264  var (
 265  	activeSubCancel context.CancelFunc
 266  	activeSubMu     sync.Mutex
 267  )
 268  
 269  func jsSubscribe(this js.Value, args []js.Value) any {
 270  	if client == nil {
 271  		return nil
 272  	}
 273  	activeSubMu.Lock()
 274  	if activeSubCancel != nil {
 275  		activeSubCancel()
 276  	}
 277  	ctx, cancel := context.WithCancel(context.Background())
 278  	activeSubCancel = cancel
 279  	activeSubMu.Unlock()
 280  
 281  	go func() {
 282  		defer cancel()
 283  		ff := client.SubscriptionFilters()
 284  		stream, err := relay.Subscribe(ctx, ff)
 285  		if err != nil {
 286  			return
 287  		}
 288  		defer stream.Close()
 289  
 290  		for {
 291  			select {
 292  			case <-ctx.Done():
 293  				return
 294  			case ev := <-stream.Events():
 295  				if ev == nil {
 296  					return
 297  				}
 298  				_ = client.HandleEvent(ctx, ev)
 299  			case <-client.GroupsChanged():
 300  				stream.Close()
 301  				ff = client.SubscriptionFilters()
 302  				stream, err = relay.Subscribe(ctx, ff)
 303  				if err != nil {
 304  					return
 305  				}
 306  			}
 307  		}
 308  	}()
 309  	return nil
 310  }
 311  
 312  func jsPublishKP(this js.Value, args []js.Value) any {
 313  	if client == nil {
 314  		return nil
 315  	}
 316  	go func() {
 317  		if err := client.PublishKeyPackage(context.Background()); err != nil {
 318  			fmt.Println("marmot-wasm: publishKP error:", err)
 319  		}
 320  	}()
 321  	return nil
 322  }
 323  
 324  func jsListGroups(this js.Value, args []js.Value) any {
 325  	if client == nil {
 326  		return "[]"
 327  	}
 328  	ids := client.ActiveGroupIDs()
 329  	out := "["
 330  	for i, id := range ids {
 331  		if i > 0 {
 332  			out += ","
 333  		}
 334  		out += "\"" + id + "\""
 335  	}
 336  	out += "]"
 337  	return out
 338  }
 339  
 340  func jsBackupGroups(this js.Value, args []js.Value) any {
 341  	if client == nil {
 342  		return "error: not initialized"
 343  	}
 344  	go func() {
 345  		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 346  		defer cancel()
 347  		if err := client.BackupGroups(ctx); err != nil {
 348  			sendStatus("backup error: " + err.Error())
 349  			return
 350  		}
 351  		sendStatus("backup ok")
 352  	}()
 353  	return nil
 354  }
 355  
 356  func jsRestoreGroups(this js.Value, args []js.Value) any {
 357  	if client == nil {
 358  		return "error: not initialized"
 359  	}
 360  	go func() {
 361  		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
 362  		defer cancel()
 363  		n, err := client.RestoreGroups(ctx)
 364  		if err != nil {
 365  			sendStatus("restore error: " + err.Error())
 366  			return
 367  		}
 368  		sendStatus(fmt.Sprintf("restore ok:%d", n))
 369  	}()
 370  	return nil
 371  }
 372  
 373  func jsRatchetGroup(this js.Value, args []js.Value) any {
 374  	if len(args) < 1 || client == nil {
 375  		return "error: missing args or not initialized"
 376  	}
 377  	peerHex := args[0].String()
 378  	go func() {
 379  		peerPub, err := hex.Dec(peerHex)
 380  		if err != nil {
 381  			sendStatus("ratchet error: invalid peer: " + err.Error())
 382  			return
 383  		}
 384  		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
 385  		defer cancel()
 386  		if err := client.RatchetGroup(ctx, peerPub); err != nil {
 387  			sendStatus("ratchet error: " + err.Error())
 388  			return
 389  		}
 390  		sendStatus("ratchet ok:" + peerHex)
 391  	}()
 392  	return nil
 393  }
 394  
 395  func jsHandleEvent(this js.Value, args []js.Value) any {
 396  	if len(args) < 1 || client == nil {
 397  		return nil
 398  	}
 399  	evJSON := args[0].String()
 400  	go func() {
 401  		ev := event.New()
 402  		if err := ev.UnmarshalJSON([]byte(evJSON)); err != nil {
 403  			return
 404  		}
 405  		_ = client.HandleEvent(context.Background(), ev)
 406  	}()
 407  	return nil
 408  }
 409  
 410  func jsDeliverEvent(this js.Value, args []js.Value) any {
 411  	if len(args) < 2 {
 412  		return nil
 413  	}
 414  	subID := args[0].Int()
 415  	evJSON := args[1].String()
 416  	deliverEvent(subID, evJSON)
 417  	return nil
 418  }
 419  
 420  func jsCryptoResult(this js.Value, args []js.Value) any {
 421  	if len(args) < 3 || crypto == nil {
 422  		return nil
 423  	}
 424  	id := args[0].Int()
 425  	result := args[1].String()
 426  	errMsg := args[2].String()
 427  	crypto.Resolve(id, result, errMsg)
 428  	return nil
 429  }
 430  
 431  func jsLastEventTS(this js.Value, args []js.Value) any {
 432  	if client == nil {
 433  		return 0
 434  	}
 435  	return client.LastEventTS()
 436  }
 437  
 438  func jsKeyPackageEvent(this js.Value, args []js.Value) any {
 439  	if client == nil {
 440  		return ""
 441  	}
 442  	ev, err := client.KeyPackageEvent()
 443  	if err != nil {
 444  		return "error: " + err.Error()
 445  	}
 446  	b, err := ev.MarshalJSON()
 447  	if err != nil {
 448  		return "error: " + err.Error()
 449  	}
 450  	return string(b)
 451  }
 452  
 453  // --- IDB-backed GroupStore via JS callbacks ---
 454  
 455  type storeResult struct {
 456  	data string
 457  	err  string
 458  }
 459  
 460  type jsGroupStore struct {
 461  	mu      sync.Mutex
 462  	pending map[int]chan storeResult
 463  	nextID  int
 464  }
 465  
 466  func newJSGroupStore() *jsGroupStore {
 467  	return &jsGroupStore{pending: make(map[int]chan storeResult)}
 468  }
 469  
 470  func (s *jsGroupStore) newPending() (int, chan storeResult) {
 471  	s.mu.Lock()
 472  	id := s.nextID
 473  	s.nextID++
 474  	ch := make(chan storeResult, 1)
 475  	s.pending[id] = ch
 476  	s.mu.Unlock()
 477  	return id, ch
 478  }
 479  
 480  func (s *jsGroupStore) resolve(id int, data, errMsg string) {
 481  	s.mu.Lock()
 482  	ch, ok := s.pending[id]
 483  	if ok {
 484  		delete(s.pending, id)
 485  	}
 486  	s.mu.Unlock()
 487  	if ok {
 488  		ch <- storeResult{data: data, err: errMsg}
 489  	}
 490  }
 491  
 492  func (s *jsGroupStore) SaveGroup(groupID, state []byte) error {
 493  	id, ch := s.newPending()
 494  	js.Global().Call("_marmot_store_save", id, hex.Enc(groupID), hex.Enc(state))
 495  	r := <-ch
 496  	if r.err != "" {
 497  		return fmt.Errorf("%s", r.err)
 498  	}
 499  	return nil
 500  }
 501  
 502  func (s *jsGroupStore) LoadGroup(groupID []byte) ([]byte, error) {
 503  	id, ch := s.newPending()
 504  	js.Global().Call("_marmot_store_load", id, hex.Enc(groupID))
 505  	r := <-ch
 506  	if r.err != "" {
 507  		return nil, fmt.Errorf("%s", r.err)
 508  	}
 509  	if r.data == "" {
 510  		return nil, os.ErrNotExist
 511  	}
 512  	return hex.Dec(r.data)
 513  }
 514  
 515  func (s *jsGroupStore) ListGroups() ([][]byte, error) {
 516  	id, ch := s.newPending()
 517  	js.Global().Call("_marmot_store_list", id)
 518  	r := <-ch
 519  	if r.err != "" {
 520  		return nil, fmt.Errorf("%s", r.err)
 521  	}
 522  	if r.data == "" {
 523  		return nil, nil
 524  	}
 525  	var result [][]byte
 526  	start := 0
 527  	for i := 0; i <= len(r.data); i++ {
 528  		if i == len(r.data) || r.data[i] == ',' {
 529  			h := r.data[start:i]
 530  			start = i + 1
 531  			if h == "" {
 532  				continue
 533  			}
 534  			b, err := hex.Dec(h)
 535  			if err != nil {
 536  				continue
 537  			}
 538  			result = append(result, b)
 539  		}
 540  	}
 541  	return result, nil
 542  }
 543  
 544  func (s *jsGroupStore) DeleteGroup(groupID []byte) error {
 545  	id, ch := s.newPending()
 546  	js.Global().Call("_marmot_store_delete", id, hex.Enc(groupID))
 547  	r := <-ch
 548  	if r.err != "" {
 549  		return fmt.Errorf("%s", r.err)
 550  	}
 551  	return nil
 552  }
 553  
 554  // SaveKeyPackage is a no-op in WASM — key packages are ephemeral.
 555  func (s *jsGroupStore) SaveKeyPackage([]byte) error { return nil }
 556  
 557  // LoadKeyPackage always returns not-found in WASM — generates fresh each time.
 558  func (s *jsGroupStore) LoadKeyPackage() ([]byte, error) { return nil, os.ErrNotExist }
 559  
 560  func jsStoreResult(this js.Value, args []js.Value) any {
 561  	if len(args) < 3 || store == nil {
 562  		return nil
 563  	}
 564  	store.resolve(args[0].Int(), args[1].String(), args[2].String())
 565  	return nil
 566  }
 567