Compare commits

..

No commits in common. "3ea051cb3727d35dd0d3f4c78cbdff3b00a02a8e" and "56a725c5df0a810fab17ed468d35a5dc27caa3e4" have entirely different histories.

10 changed files with 107 additions and 398 deletions

View File

@ -2,10 +2,6 @@ mongodb:
uri: mongodb://localhost:27017/
database: grain
auth:
enabled: false # Enable or disable AUTH handling
relay_url: "wss://relay.example.com/" # Specify the relay URL
server:
port: :8181
read_timeout: 10 # in seconds
@ -14,12 +10,6 @@ server:
max_connections: 100
max_subscriptions_per_client: 10
resource_limits:
cpu_cores: 2 # Limit the number of CPU cores the application can use
memory_mb: 1024 # Cap the maximum amount of RAM in MB the application can use
heap_size_mb: 512 # Set a limit on the Go garbage collector's heap size in MB
max_goroutines: 100 # Limit the maximum number of concurrently running Go routines
rate_limit:
ws_limit: 100 # WebSocket messages per second
ws_burst: 200 # Allowed burst of WebSocket messages

View File

@ -1,6 +0,0 @@
package config
type AuthConfig struct {
Enabled bool `yaml:"enabled"`
RelayURL string `yaml:"relay_url"`
}

View File

@ -1,8 +0,0 @@
package config
type ResourceLimits struct {
CPUCores int `yaml:"cpu_cores"`
MemoryMB int `yaml:"memory_mb"`
HeapSizeMB int `yaml:"heap_size_mb"`
MaxGoroutines int `yaml:"max_goroutines"`
}

View File

@ -18,6 +18,4 @@ type ServerConfig struct {
KindWhitelist KindWhitelistConfig `yaml:"kind_whitelist"`
DomainWhitelist DomainWhitelistConfig `yaml:"domain_whitelist"`
Blacklist BlacklistConfig `yaml:"blacklist"`
ResourceLimits ResourceLimits `yaml:"resource_limits"`
Auth AuthConfig `yaml:"auth"`
}

26
main.go
View File

@ -25,22 +25,19 @@ func main() {
utils.EnsureFileExists("relay_metadata.json", "app/static/examples/relay_metadata.example.json")
restartChan := make(chan struct{})
go utils.WatchConfigFile("config.yml", restartChan) // Critical goroutine
go utils.WatchConfigFile("config.yml", restartChan)
signalChan := make(chan os.Signal, 1)
signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
var wg sync.WaitGroup
for {
wg.Add(1) // Add to WaitGroup for the server goroutine
wg.Add(1)
cfg, err := config.LoadConfig("config.yml")
if err != nil {
log.Fatal("Error loading config: ", err)
}
utils.ApplyResourceLimits(&cfg.ResourceLimits) // Apply limits once before starting the server
client, err := db.InitDB(cfg)
if err != nil {
log.Fatal("Error initializing database: ", err)
@ -57,22 +54,22 @@ func main() {
}
mux := setupRoutes()
// Start the server
server := startServer(cfg, mux, &wg)
select {
case <-restartChan:
log.Println("Restarting server...")
server.Close() // Stop the current server instance
wg.Wait() // Wait for the server goroutine to finish
time.Sleep(3 * time.Second)
// Close server before restart
server.Close()
wg.Wait()
time.Sleep(3 * time.Second)
case <-signalChan:
log.Println("Shutting down server...")
server.Close() // Stop the server
db.DisconnectDB(client) // Disconnect from MongoDB
wg.Wait() // Wait for all goroutines to finish
server.Close()
db.DisconnectDB(client)
wg.Wait()
return
}
}
@ -100,14 +97,13 @@ func startServer(config *configTypes.ServerConfig, mux *http.ServeMux, wg *sync.
}
go func() {
defer wg.Done() // Notify that the server goroutine is done
defer wg.Done() // Notify that the server is done shutting down
fmt.Printf("Server is running on http://localhost%s\n", config.Server.Port)
err := server.ListenAndServe()
if err != nil && err != http.ErrServerClosed {
fmt.Println("Error starting server:", err)
}
}()
return server
}

View File

@ -1,139 +0,0 @@
package handlers
import (
"encoding/json"
"errors"
"fmt"
"grain/config"
"grain/server/handlers/response"
"grain/server/utils"
"time"
relay "grain/server/types"
"golang.org/x/net/websocket"
)
// HandleAuth handles the "AUTH" message type as defined in NIP-42
func HandleAuth(ws *websocket.Conn, message []interface{}) {
if !config.GetConfig().Auth.Enabled {
fmt.Println("AUTH is disabled in the configuration")
response.SendNotice(ws, "", "AUTH is disabled")
return
}
if len(message) != 2 {
fmt.Println("Invalid AUTH message format")
response.SendNotice(ws, "", "Invalid AUTH message format")
return
}
authData, ok := message[1].(map[string]interface{})
if !ok {
fmt.Println("Invalid auth data format")
response.SendNotice(ws, "", "Invalid auth data format")
return
}
authBytes, err := json.Marshal(authData)
if err != nil {
fmt.Println("Error marshaling auth data:", err)
response.SendNotice(ws, "", "Error marshaling auth data")
return
}
var authEvent relay.Event
err = json.Unmarshal(authBytes, &authEvent)
if err != nil {
fmt.Println("Error unmarshaling auth data:", err)
response.SendNotice(ws, "", "Error unmarshaling auth data")
return
}
err = VerifyAuthEvent(authEvent)
if err != nil {
response.SendOK(ws, authEvent.ID, false, err.Error())
return
}
// Mark the session as authenticated after successful verification
SetAuthenticated(ws)
response.SendOK(ws, authEvent.ID, true, "")
}
// VerifyAuthEvent verifies the authentication event according to NIP-42
func VerifyAuthEvent(evt relay.Event) error {
if evt.Kind != 22242 {
return errors.New("invalid: event kind must be 22242")
}
if time.Since(time.Unix(evt.CreatedAt, 0)) > 10*time.Minute {
return errors.New("invalid: event is too old")
}
challenge, err := extractTag(evt.Tags, "challenge")
if err != nil {
return errors.New("invalid: challenge tag missing")
}
relayURL, err := extractTag(evt.Tags, "relay")
if err != nil {
return errors.New("invalid: relay tag missing")
}
expectedChallenge := GetChallengeForConnection(evt.PubKey)
if challenge != expectedChallenge {
return errors.New("invalid: challenge does not match")
}
if relayURL != config.GetConfig().Auth.RelayURL {
return errors.New("invalid: relay URL does not match")
}
if !utils.CheckSignature(evt) {
return errors.New("invalid: signature verification failed")
}
return nil
}
// extractTag extracts a specific tag from an event's tags
func extractTag(tags [][]string, key string) (string, error) {
for _, tag := range tags {
if len(tag) >= 2 && tag[0] == key {
return tag[1], nil
}
}
return "", errors.New("tag not found")
}
// Map to store challenges associated with connections
var challenges = make(map[string]string)
var authSessions = make(map[*websocket.Conn]bool)
// GetChallengeForConnection retrieves the challenge string for a given connection
func GetChallengeForConnection(pubKey string) string {
mu.Lock()
defer mu.Unlock()
return challenges[pubKey]
}
// SetChallengeForConnection sets the challenge string for a given connection
func SetChallengeForConnection(pubKey, challenge string) {
mu.Lock()
defer mu.Unlock()
challenges[pubKey] = challenge
}
// SetAuthenticated marks a connection as authenticated
func SetAuthenticated(ws *websocket.Conn) {
mu.Lock()
defer mu.Unlock()
authSessions[ws] = true
}
// IsAuthenticated checks if a connection is authenticated
func IsAuthenticated(ws *websocket.Conn) bool {
mu.Lock()
defer mu.Unlock()
return authSessions[ws]
}

View File

@ -16,39 +16,37 @@ import (
)
func HandleEvent(ws *websocket.Conn, message []interface{}) {
utils.LimitedGoRoutine(func() {
if len(message) != 2 {
fmt.Println("Invalid EVENT message format")
response.SendNotice(ws, "", "Invalid EVENT message format")
return
}
if len(message) != 2 {
fmt.Println("Invalid EVENT message format")
response.SendNotice(ws, "", "Invalid EVENT message format")
return
}
eventData, ok := message[1].(map[string]interface{})
if !ok {
fmt.Println("Invalid event data format")
response.SendNotice(ws, "", "Invalid event data format")
return
}
eventBytes, err := json.Marshal(eventData)
if err != nil {
fmt.Println("Error marshaling event data:", err)
response.SendNotice(ws, "", "Error marshaling event data")
return
}
eventData, ok := message[1].(map[string]interface{})
if !ok {
fmt.Println("Invalid event data format")
response.SendNotice(ws, "", "Invalid event data format")
return
}
eventBytes, err := json.Marshal(eventData)
if err != nil {
fmt.Println("Error marshaling event data:", err)
response.SendNotice(ws, "", "Error marshaling event data")
return
}
var evt relay.Event
err = json.Unmarshal(eventBytes, &evt)
if err != nil {
fmt.Println("Error unmarshaling event data:", err)
response.SendNotice(ws, "", "Error unmarshaling event data")
return
}
var evt relay.Event
err = json.Unmarshal(eventBytes, &evt)
if err != nil {
fmt.Println("Error unmarshaling event data:", err)
response.SendNotice(ws, "", "Error unmarshaling event data")
return
}
eventSize := len(eventBytes) // Calculate event size
HandleKind(context.TODO(), evt, ws, eventSize)
eventSize := len(eventBytes) // Calculate event size
HandleKind(context.TODO(), evt, ws, eventSize)
fmt.Println("Event processed:", evt.ID)
})
fmt.Println("Event processed:", evt.ID)
}
func HandleKind(ctx context.Context, evt relay.Event, ws *websocket.Conn, eventSize int) {
@ -154,3 +152,4 @@ func determineCategory(kind int) string {
return "unknown"
}
}

View File

@ -21,89 +21,87 @@ var subscriptions = make(map[string]relay.Subscription)
var mu sync.Mutex // Protect concurrent access to subscriptions map
func HandleReq(ws *websocket.Conn, message []interface{}, subscriptions map[string][]relay.Filter) {
utils.LimitedGoRoutine(func() {
if len(message) < 3 {
fmt.Println("Invalid REQ message format")
response.SendClosed(ws, "", "invalid: invalid REQ message format")
return
}
if len(message) < 3 {
fmt.Println("Invalid REQ message format")
response.SendClosed(ws, "", "invalid: invalid REQ message format")
return
}
subID, ok := message[1].(string)
subID, ok := message[1].(string)
if !ok {
fmt.Println("Invalid subscription ID format")
response.SendClosed(ws, "", "invalid: invalid subscription ID format")
return
}
mu.Lock()
defer mu.Unlock()
// Check the current number of subscriptions for the client
if len(subscriptions) >= config.GetConfig().Server.MaxSubscriptionsPerClient {
// Find and remove the oldest subscription (FIFO)
var oldestSubID string
for id := range subscriptions {
oldestSubID = id
break
}
delete(subscriptions, oldestSubID)
fmt.Println("Dropped oldest subscription:", oldestSubID)
}
filters := make([]relay.Filter, len(message)-2)
for i, filter := range message[2:] {
filterData, ok := filter.(map[string]interface{})
if !ok {
fmt.Println("Invalid subscription ID format")
response.SendClosed(ws, "", "invalid: invalid subscription ID format")
fmt.Println("Invalid filter format")
response.SendClosed(ws, subID, "invalid: invalid filter format")
return
}
mu.Lock()
defer mu.Unlock()
var f relay.Filter
f.IDs = utils.ToStringArray(filterData["ids"])
f.Authors = utils.ToStringArray(filterData["authors"])
f.Kinds = utils.ToIntArray(filterData["kinds"])
f.Tags = utils.ToTagsMap(filterData["tags"])
f.Since = utils.ToTime(filterData["since"])
f.Until = utils.ToTime(filterData["until"])
f.Limit = utils.ToInt(filterData["limit"])
// Check the current number of subscriptions for the client
if len(subscriptions) >= config.GetConfig().Server.MaxSubscriptionsPerClient {
// Find and remove the oldest subscription (FIFO)
var oldestSubID string
for id := range subscriptions {
oldestSubID = id
break
}
delete(subscriptions, oldestSubID)
fmt.Println("Dropped oldest subscription:", oldestSubID)
}
filters[i] = f
}
filters := make([]relay.Filter, len(message)-2)
for i, filter := range message[2:] {
filterData, ok := filter.(map[string]interface{})
if !ok {
fmt.Println("Invalid filter format")
response.SendClosed(ws, subID, "invalid: invalid filter format")
return
}
// Add the new subscription or update the existing one
subscriptions[subID] = filters
fmt.Println("Subscription updated:", subID)
var f relay.Filter
f.IDs = utils.ToStringArray(filterData["ids"])
f.Authors = utils.ToStringArray(filterData["authors"])
f.Kinds = utils.ToIntArray(filterData["kinds"])
f.Tags = utils.ToTagsMap(filterData["tags"])
f.Since = utils.ToTime(filterData["since"])
f.Until = utils.ToTime(filterData["until"])
f.Limit = utils.ToInt(filterData["limit"])
// Query the database with filters and send back the results
queriedEvents, err := QueryEvents(filters, db.GetClient(), "grain")
if err != nil {
fmt.Println("Error querying events:", err)
response.SendClosed(ws, subID, "error: could not query events")
return
}
filters[i] = f
}
// Add the new subscription or update the existing one
subscriptions[subID] = filters
fmt.Println("Subscription updated:", subID)
// Query the database with filters and send back the results
queriedEvents, err := QueryEvents(filters, db.GetClient(), "grain")
for _, evt := range queriedEvents {
msg := []interface{}{"EVENT", subID, evt}
msgBytes, _ := json.Marshal(msg)
err = websocket.Message.Send(ws, string(msgBytes))
if err != nil {
fmt.Println("Error querying events:", err)
response.SendClosed(ws, subID, "error: could not query events")
fmt.Println("Error sending event:", err)
response.SendClosed(ws, subID, "error: could not send event")
return
}
}
for _, evt := range queriedEvents {
msg := []interface{}{"EVENT", subID, evt}
msgBytes, _ := json.Marshal(msg)
err = websocket.Message.Send(ws, string(msgBytes))
if err != nil {
fmt.Println("Error sending event:", err)
response.SendClosed(ws, subID, "error: could not send event")
return
}
}
// Indicate end of stored events
eoseMsg := []interface{}{"EOSE", subID}
eoseBytes, _ := json.Marshal(eoseMsg)
err = websocket.Message.Send(ws, string(eoseBytes))
if err != nil {
fmt.Println("Error sending EOSE:", err)
response.SendClosed(ws, subID, "error: could not send EOSE")
return
}
})
// Indicate end of stored events
eoseMsg := []interface{}{"EOSE", subID}
eoseBytes, _ := json.Marshal(eoseMsg)
err = websocket.Message.Send(ws, string(eoseBytes))
if err != nil {
fmt.Println("Error sending EOSE:", err)
response.SendClosed(ws, subID, "error: could not send EOSE")
return
}
}
// QueryEvents queries events from the MongoDB collection based on filters

View File

@ -105,12 +105,6 @@ func WebSocketHandler(ws *websocket.Conn) {
return
}
handlers.HandleReq(ws, message, subscriptions)
case "AUTH":
if config.GetConfig().Auth.Enabled {
handlers.HandleAuth(ws, message)
} else {
fmt.Println("Received AUTH message, but AUTH is disabled")
}
case "CLOSE":
handlers.HandleClose(ws, message)
default:

View File

@ -1,113 +0,0 @@
package utils
import (
"log"
"runtime"
"runtime/debug"
"sync"
"time"
configTypes "grain/config/types"
)
var (
maxGoroutinesChan chan struct{}
wg sync.WaitGroup
goroutineQueue []func()
goroutineQueueMutex sync.Mutex
)
func ApplyResourceLimits(cfg *configTypes.ResourceLimits) {
// Set CPU cores
runtime.GOMAXPROCS(cfg.CPUCores)
// Set maximum heap size
if cfg.HeapSizeMB > 0 {
heapSize := int64(uint64(cfg.HeapSizeMB) * 1024 * 1024)
debug.SetMemoryLimit(heapSize)
log.Printf("Heap size limited to %d MB\n", cfg.HeapSizeMB)
}
// Start monitoring memory usage
if cfg.MemoryMB > 0 {
go monitorMemoryUsage(cfg.MemoryMB)
log.Printf("Max memory usage limited to %d MB\n", cfg.MemoryMB)
}
// Set maximum number of Go routines
if cfg.MaxGoroutines > 0 {
maxGoroutinesChan = make(chan struct{}, cfg.MaxGoroutines)
log.Printf("Max goroutines limited to %d\n", cfg.MaxGoroutines)
}
}
// LimitedGoRoutine starts a goroutine with limit enforcement
func LimitedGoRoutine(f func()) {
// By default, all routines are considered critical
goroutineQueueMutex.Lock()
goroutineQueue = append(goroutineQueue, f)
goroutineQueueMutex.Unlock()
attemptToStartGoroutine()
}
func attemptToStartGoroutine() {
goroutineQueueMutex.Lock()
defer goroutineQueueMutex.Unlock()
if len(goroutineQueue) > 0 {
select {
case maxGoroutinesChan <- struct{}{}:
wg.Add(1)
go func(f func()) {
defer func() {
wg.Done()
<-maxGoroutinesChan
attemptToStartGoroutine()
}()
f()
}(goroutineQueue[0])
// Remove the started goroutine from the queue
goroutineQueue = goroutineQueue[1:]
default:
// If the channel is full, consider dropping the oldest non-critical goroutine
dropOldestNonCriticalGoroutine()
}
}
}
func dropOldestNonCriticalGoroutine() {
goroutineQueueMutex.Lock()
defer goroutineQueueMutex.Unlock()
if len(goroutineQueue) > 0 {
log.Println("Dropping the oldest non-critical goroutine to free resources.")
goroutineQueue = goroutineQueue[1:]
attemptToStartGoroutine()
}
}
func WaitForGoroutines() {
wg.Wait()
}
func monitorMemoryUsage(maxMemoryMB int) {
for {
var memStats runtime.MemStats
runtime.ReadMemStats(&memStats)
usedMemoryMB := int(memStats.Alloc / 1024 / 1024)
if usedMemoryMB > maxMemoryMB {
log.Printf("Memory usage exceeded limit: %d MB used, limit is %d MB\n", usedMemoryMB, maxMemoryMB)
debug.FreeOSMemory() // Attempt to free memory
// If memory usage is still high, attempt to drop non-critical goroutines
if usedMemoryMB > maxMemoryMB {
dropOldestNonCriticalGoroutine()
}
}
time.Sleep(1 * time.Second)
}
}