grain/events/events.go

package events

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"

	"github.com/btcsuite/btcd/btcec/v2"
	"github.com/btcsuite/btcd/btcec/v2/schnorr"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"golang.org/x/net/websocket"
)

type Event struct {
	ID        string     `json:"id"`
	PubKey    string     `json:"pubkey"`
	CreatedAt int64      `json:"created_at"`
	Kind      int        `json:"kind"`
	Tags      [][]string `json:"tags"`
	Content   string     `json:"content"`
	Sig       string     `json:"sig"`
}

var collections = make(map[int]*mongo.Collection)

func InitCollections(client *mongo.Client, kinds ...int) {
	for _, kind := range kinds {
		collectionName := fmt.Sprintf("event-kind%d", kind)
		collections[kind] = client.Database("grain").Collection(collectionName)
		indexModel := mongo.IndexModel{
			Keys:    bson.D{{Key: "id", Value: 1}},
			Options: options.Index().SetUnique(true),
		}
		_, err := collections[kind].Indexes().CreateOne(context.TODO(), indexModel)
		if err != nil {
			fmt.Printf("Failed to create index on %s: %v\n", collectionName, err)
		}
	}
}

func GetCollection(kind int, client *mongo.Client) *mongo.Collection {
	if collection, exists := collections[kind]; exists {
		return collection
	}
	collectionName := fmt.Sprintf("event-kind%d", kind)
	collection := client.Database("grain").Collection(collectionName)
	collections[kind] = collection
	indexModel := mongo.IndexModel{
		Keys:    bson.D{{Key: "id", Value: 1}},
		Options: options.Index().SetUnique(true),
	}
	_, err := collection.Indexes().CreateOne(context.TODO(), indexModel)
	if err != nil {
		fmt.Printf("Failed to create index on %s: %v\n", collectionName, err)
	}
	return collection
}

func HandleEvent(ctx context.Context, evt Event, client *mongo.Client, ws *websocket.Conn) {
	if !ValidateEvent(evt) {
		sendOKResponse(ws, evt.ID, false, "invalid: signature verification failed")
		return
	}

	collection := GetCollection(evt.Kind, client)

	var err error
	switch evt.Kind {
	case 0:
		err = HandleEventKind0(ctx, evt, collection)
	case 1:
		err = HandleEventKind1(ctx, evt, collection)
	default:
		err = HandleDefaultEvent(ctx, evt, collection)
	}

	if err != nil {
		sendOKResponse(ws, evt.ID, false, fmt.Sprintf("error: %v", err))
		return
	}

	sendOKResponse(ws, evt.ID, true, "")
}

func sendOKResponse(ws *websocket.Conn, eventID string, status bool, message string) {
	response := []interface{}{"OK", eventID, status, message}
	responseBytes, _ := json.Marshal(response)
	websocket.Message.Send(ws, string(responseBytes))
}

func ValidateEvent(evt Event) bool {
	serializedEvent := SerializeEvent(evt)
	hash := sha256.Sum256(serializedEvent)
	eventID := hex.EncodeToString(hash[:])
	if eventID != evt.ID {
		return false
	}

	sigBytes, err := hex.DecodeString(evt.Sig)
	if err != nil {
		return false
	}

	sig, err := schnorr.ParseSignature(sigBytes)
	if err != nil {
		return false
	}

	pubKeyBytes, err := hex.DecodeString(evt.PubKey)
	if err != nil {
		return false
	}

	pubKey, err := btcec.ParsePubKey(pubKeyBytes)
	if err != nil {
		return false
	}

	return sig.Verify(hash[:], pubKey)
}

func SerializeEvent(evt Event) []byte {
	eventData := []interface{}{
		0,
		evt.PubKey,
		evt.CreatedAt,
		evt.Kind,
		evt.Tags,
		evt.Content,
	}
	serializedEvent, _ := json.Marshal(eventData)
	return serializedEvent
}

func HandleDefaultEvent(ctx context.Context, evt Event, collection *mongo.Collection) error {
	_, err := collection.InsertOne(ctx, evt)
	if err != nil {
		return fmt.Errorf("Error inserting default event into MongoDB: %v", err)
	}

	fmt.Println("Inserted default event into MongoDB:", evt.ID)
	return nil
}