From f7ce78d7f80a10d7c1489d779e0e1eb0036a2273 Mon Sep 17 00:00:00 2001
From: fiatjaf <fiatjaf@gmail.com>
Date: Tue, 9 May 2023 16:31:10 -0300
Subject: [PATCH] add custom envelope types with json codecs.

---
 envelopes.go      | 190 ++++++++++++++++++++++++++++++++++++++++++++++
 envelopes_test.go | 117 ++++++++++++++++++++++++++++
 go.mod            |   3 +
 go.sum            |   6 ++
 4 files changed, 316 insertions(+)
 create mode 100644 envelopes.go
 create mode 100644 envelopes_test.go

diff --git a/envelopes.go b/envelopes.go
new file mode 100644
index 0000000..d90916c
--- /dev/null
+++ b/envelopes.go
@@ -0,0 +1,190 @@
+package nostr
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/mailru/easyjson"
+	jwriter "github.com/mailru/easyjson/jwriter"
+	"github.com/tidwall/gjson"
+)
+
+type EventEnvelope struct {
+	SubscriptionID *string
+	Event
+}
+
+func (v *EventEnvelope) UnmarshalJSON(data []byte) error {
+	r := gjson.ParseBytes(data)
+	arr := r.Array()
+	switch len(arr) {
+	case 2:
+		return easyjson.Unmarshal([]byte(arr[1].Raw), &v.Event)
+	case 3:
+		v.SubscriptionID = &arr[1].Str
+		return easyjson.Unmarshal([]byte(arr[2].Raw), &v.Event)
+	default:
+		return fmt.Errorf("failed to decode EVENT envelope")
+	}
+}
+
+func (v EventEnvelope) MarshalJSON() ([]byte, error) {
+	w := jwriter.Writer{}
+	w.RawString(`["EVENT",`)
+	if v.SubscriptionID != nil {
+		w.RawString(`"` + *v.SubscriptionID + `",`)
+	}
+	v.MarshalEasyJSON(&w)
+	w.RawString(`]`)
+	return w.BuildBytes()
+}
+
+type ReqEnvelope struct {
+	SubscriptionID string
+	Filters
+}
+
+func (v *ReqEnvelope) UnmarshalJSON(data []byte) error {
+	r := gjson.ParseBytes(data)
+	arr := r.Array()
+	if len(arr) < 3 {
+		return fmt.Errorf("failed to decode REQ envelope: missing filters")
+	}
+	v.SubscriptionID = arr[1].Str
+	v.Filters = make(Filters, len(arr)-2)
+	f := 0
+	for i := 2; i < len(arr); i++ {
+		if err := easyjson.Unmarshal([]byte(arr[i].Raw), &v.Filters[f]); err != nil {
+			return fmt.Errorf("%w -- on filter %d", err, f)
+		}
+		f++
+	}
+
+	return nil
+}
+
+func (v ReqEnvelope) MarshalJSON() ([]byte, error) {
+	w := jwriter.Writer{}
+	w.RawString(`["REQ",`)
+	w.RawString(`"` + v.SubscriptionID + `"`)
+	for _, filter := range v.Filters {
+		w.RawString(`,`)
+		filter.MarshalEasyJSON(&w)
+	}
+	w.RawString(`]`)
+	return w.BuildBytes()
+}
+
+type NoticeEnvelope string
+
+func (v *NoticeEnvelope) UnmarshalJSON(data []byte) error {
+	r := gjson.ParseBytes(data)
+	arr := r.Array()
+	switch len(arr) {
+	case 2:
+		*v = NoticeEnvelope(arr[1].Str)
+		return nil
+	default:
+		return fmt.Errorf("failed to decode NOTICE envelope")
+	}
+}
+
+func (v NoticeEnvelope) MarshalJSON() ([]byte, error) {
+	w := jwriter.Writer{}
+	w.RawString(`["NOTICE",`)
+	w.Raw(json.Marshal(string(v)))
+	w.RawString(`]`)
+	return w.BuildBytes()
+}
+
+type EOSEEnvelope string
+
+func (v *EOSEEnvelope) UnmarshalJSON(data []byte) error {
+	r := gjson.ParseBytes(data)
+	arr := r.Array()
+	switch len(arr) {
+	case 2:
+		*v = EOSEEnvelope(arr[1].Str)
+		return nil
+	default:
+		return fmt.Errorf("failed to decode EOSE envelope")
+	}
+}
+
+func (v EOSEEnvelope) MarshalJSON() ([]byte, error) {
+	w := jwriter.Writer{}
+	w.RawString(`["EOSE",`)
+	w.Raw(json.Marshal(string(v)))
+	w.RawString(`]`)
+	return w.BuildBytes()
+}
+
+type OKEnvelope struct {
+	EventID string
+	OK      bool
+	Reason  *string
+}
+
+func (v *OKEnvelope) UnmarshalJSON(data []byte) error {
+	r := gjson.ParseBytes(data)
+	arr := r.Array()
+	if len(arr) < 3 {
+		return fmt.Errorf("failed to decode OK envelope: missing fields")
+	}
+	v.EventID = arr[1].Str
+	v.OK = arr[2].Raw == "true"
+
+	if len(arr) > 3 {
+		v.Reason = &arr[3].Str
+	}
+
+	return nil
+}
+
+func (v OKEnvelope) MarshalJSON() ([]byte, error) {
+	w := jwriter.Writer{}
+	w.RawString(`["OK",`)
+	w.RawString(`"` + v.EventID + `",`)
+	ok := "false"
+	if v.OK {
+		ok = "true"
+	}
+	w.RawString(ok)
+	if v.Reason != nil {
+		w.RawString(`,`)
+		w.Raw(json.Marshal(v.Reason))
+	}
+	w.RawString(`]`)
+	return w.BuildBytes()
+}
+
+type AuthEnvelope struct {
+	Challenge *string
+	Event     Event
+}
+
+func (v *AuthEnvelope) UnmarshalJSON(data []byte) error {
+	r := gjson.ParseBytes(data)
+	arr := r.Array()
+	if len(arr) < 2 {
+		return fmt.Errorf("failed to decode Auth envelope: missing fields")
+	}
+	if arr[1].IsObject() {
+		return easyjson.Unmarshal([]byte(arr[1].Raw), &v.Event)
+	} else {
+		v.Challenge = &arr[1].Str
+	}
+	return nil
+}
+
+func (v AuthEnvelope) MarshalJSON() ([]byte, error) {
+	w := jwriter.Writer{}
+	w.RawString(`["AUTH",`)
+	if v.Challenge != nil {
+		w.Raw(json.Marshal(*v.Challenge))
+	} else {
+		v.Event.MarshalEasyJSON(&w)
+	}
+	w.RawString(`]`)
+	return w.BuildBytes()
+}
diff --git a/envelopes_test.go b/envelopes_test.go
new file mode 100644
index 0000000..03502a8
--- /dev/null
+++ b/envelopes_test.go
@@ -0,0 +1,117 @@
+package nostr
+
+import (
+	"encoding/json"
+	"testing"
+)
+
+func TestEventEnvelopeEncodingAndDecoding(t *testing.T) {
+	eventEnvelopes := []string{
+		`["EVENT","_",{"id":"dc90c95f09947507c1044e8f48bcf6350aa6bff1507dd4acfc755b9239b5c962","pubkey":"3bf0c63fcb93463407af97a5e5ee64fa883d107ef9e558472c4eb9aaaefa459d","created_at":1644271588,"kind":1,"tags":[],"content":"now that https://blueskyweb.org/blog/2-7-2022-overview was announced we can stop working on nostr?","sig":"230e9d8f0ddaf7eb70b5f7741ccfa37e87a455c9a469282e3464e2052d3192cd63a167e196e381ef9d7e69e9ea43af2443b839974dc85d8aaab9efe1d9296524"}]`,
+		`["EVENT",{"id":"9e662bdd7d8abc40b5b15ee1ff5e9320efc87e9274d8d440c58e6eed2dddfbe2","pubkey":"373ebe3d45ec91977296a178d9f19f326c70631d2a1b0bbba5c5ecc2eb53b9e7","created_at":1644844224,"kind":3,"tags":[["p","3bf0c63fcb93463407af97a5e5ee64fa883d107ef9e558472c4eb9aaaefa459d"],["p","75fc5ac2487363293bd27fb0d14fb966477d0f1dbc6361d37806a6a740eda91e"],["p","46d0dfd3a724a302ca9175163bdf788f3606b3fd1bb12d5fe055d1e418cb60ea"]],"content":"{\"wss://nostr-pub.wellorder.net\":{\"read\":true,\"write\":true},\"wss://nostr.bitcoiner.social\":{\"read\":false,\"write\":true},\"wss://expensive-relay.fiatjaf.com\":{\"read\":true,\"write\":true},\"wss://relayer.fiatjaf.com\":{\"read\":true,\"write\":true},\"wss://relay.bitid.nz\":{\"read\":true,\"write\":true},\"wss://nostr.rocks\":{\"read\":true,\"write\":true}}","sig":"811355d3484d375df47581cb5d66bed05002c2978894098304f20b595e571b7e01b2efd906c5650080ffe49cf1c62b36715698e9d88b9e8be43029a2f3fa66be"}]`,
+	}
+
+	for _, raw := range eventEnvelopes {
+		var env EventEnvelope
+		err := json.Unmarshal([]byte(raw), &env)
+		if err != nil {
+			t.Errorf("failed to parse event envelope json: %v", err)
+		}
+
+		if env.GetID() != env.ID {
+			t.Errorf("error serializing event id: %s != %s", env.GetID(), env.ID)
+		}
+
+		if ok, _ := env.CheckSignature(); !ok {
+			t.Error("signature verification failed when it should have succeeded")
+		}
+
+		asjson, err := json.Marshal(env)
+		if err != nil {
+			t.Errorf("failed to re marshal event as json: %v", err)
+		}
+
+		if string(asjson) != raw {
+			t.Log(string(asjson))
+			t.Error("json serialization broken")
+		}
+	}
+}
+
+func TestNoticeEnvelopeEncodingAndDecoding(t *testing.T) {
+	src := `["NOTICE","kjasbdlasvdluiasvd\"kjasbdksab\\d"]`
+	var env NoticeEnvelope
+	json.Unmarshal([]byte(src), &env)
+	if env != "kjasbdlasvdluiasvd\"kjasbdksab\\d" {
+		t.Error("failed to decode NOTICE")
+	}
+
+	res, _ := json.Marshal(env)
+	if string(res) != src {
+		t.Errorf("failed to encode NOTICE: expected '%s', got '%s'", src, string(res))
+	}
+}
+
+func TestEoseEnvelopeEncodingAndDecoding(t *testing.T) {
+	src := `["EOSE","kjasbdlasvdluiasvd\"kjasbdksab\\d"]`
+	var env EOSEEnvelope
+	json.Unmarshal([]byte(src), &env)
+	if env != "kjasbdlasvdluiasvd\"kjasbdksab\\d" {
+		t.Error("failed to decode EOSE")
+	}
+
+	res, _ := json.Marshal(env)
+	if string(res) != src {
+		t.Errorf("failed to encode EOSE: expected '%s', got '%s'", src, string(res))
+	}
+}
+
+func TestOKEnvelopeEncodingAndDecoding(t *testing.T) {
+	okEnvelopes := []string{
+		`["OK","3bf0c63fcb93463407af97a5e5ee64fa883d107ef9e558472c4eb9aaaefaaaaa",false,"error: could not connect to the database"]`,
+		`["OK","3bf0c63fcb93463407af97a5e5ee64fa883d107ef9e558472c4eb9aaaefaaaaa",true]`,
+	}
+
+	for _, raw := range okEnvelopes {
+		var env OKEnvelope
+		err := json.Unmarshal([]byte(raw), &env)
+		if err != nil {
+			t.Errorf("failed to parse ok envelope json: %v", err)
+		}
+
+		asjson, err := json.Marshal(env)
+		if err != nil {
+			t.Errorf("failed to re marshal ok as json: %v", err)
+		}
+
+		if string(asjson) != raw {
+			t.Log(string(asjson))
+			t.Error("json serialization broken")
+		}
+	}
+}
+
+func TestAuthEnvelopeEncodingAndDecoding(t *testing.T) {
+	authEnvelopes := []string{
+		`["AUTH","kjsabdlasb aslkd kasndkad \"as.kdnbskadb"]`,
+		`["AUTH",{"id":"ae1fc7154296569d87ca4663f6bdf448c217d1590d28c85d158557b8b43b4d69","pubkey":"79be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798","created_at":1683660344,"kind":1,"tags":[],"content":"hello world","sig":"94e10947814b1ebe38af42300ecd90c7642763896c4f69506ae97bfdf54eec3c0c21df96b7d95daa74ff3d414b1d758ee95fc258125deebc31df0c6ba9396a51"}]`,
+	}
+
+	for _, raw := range authEnvelopes {
+		var env AuthEnvelope
+		err := json.Unmarshal([]byte(raw), &env)
+		if err != nil {
+			t.Errorf("failed to parse auth envelope json: %v", err)
+		}
+
+		asjson, err := json.Marshal(env)
+		if err != nil {
+			t.Errorf("failed to re marshal auth as json: %v", err)
+		}
+
+		if string(asjson) != raw {
+			t.Log(string(asjson))
+			t.Error("json serialization broken")
+		}
+	}
+}
diff --git a/go.mod b/go.mod
index 13bb6b5..f9c0694 100644
--- a/go.mod
+++ b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/gobwas/httphead v0.1.0
 	github.com/gobwas/ws v1.2.0
 	github.com/mailru/easyjson v0.7.7
+	github.com/tidwall/gjson v1.14.4
 	github.com/tyler-smith/go-bip32 v1.0.0
 	github.com/tyler-smith/go-bip39 v1.1.0
 	golang.org/x/exp v0.0.0-20221106115401-f9659909a136
@@ -23,6 +24,8 @@ require (
 	github.com/decred/dcrd/dcrec/secp256k1/v4 v4.0.1 // indirect
 	github.com/gobwas/pool v0.2.1 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
+	github.com/tidwall/match v1.1.1 // indirect
+	github.com/tidwall/pretty v1.2.0 // indirect
 	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 // indirect
 	golang.org/x/sys v0.6.0 // indirect
 )
diff --git a/go.sum b/go.sum
index 3340f78..34a8e19 100644
--- a/go.sum
+++ b/go.sum
@@ -83,6 +83,12 @@ github.com/stretchr/testify v1.1.5-0.20170601210322-f6abca593680/go.mod h1:a8OnR
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/syndtr/goleveldb v1.0.1-0.20210819022825-2ae1ddf74ef7/go.mod h1:q4W45IWZaF22tdD+VEXcAWRA037jwmWEB5VWYORlTpc=
+github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM=
+github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
+github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=
+github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
+github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs=
+github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
 github.com/tyler-smith/go-bip32 v1.0.0 h1:sDR9juArbUgX+bO/iblgZnMPeWY1KZMUC2AFUJdv5KE=
 github.com/tyler-smith/go-bip32 v1.0.0/go.mod h1:onot+eHknzV4BVPwrzqY5OoVpyCvnwD7lMawL5aQupE=
 github.com/tyler-smith/go-bip39 v1.1.0 h1:5eUemwrMargf3BSLRRCalXT93Ns6pQJIjYQN2nyfOP8=