lchat/server/main.go

package main

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"os"
	"time"

	"database/sql"
	"encoding/json"
	"log/slog"
	"net/http"

	"github.com/coder/websocket"
	"github.com/go-chi/chi/v5"
	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"

	_ "github.com/mattn/go-sqlite3"
)

func verifyJwt(token string, publicKey ed25519.PublicKey) (jwt.MapClaims, bool) {
	parsedToken, err := jwt.Parse(token, func(token *jwt.Token) (interface{}, error) {
		return publicKey, nil
	}, jwt.WithValidMethods([]string{"EdDSA"}))
	if err != nil {
		return nil, false
	}

	if !parsedToken.Valid {
		return nil, false
	}

	claims, ok := parsedToken.Claims.(jwt.MapClaims)
	if !ok {
		return nil, false
	}

	return claims, true
}

func renderJSON(data interface{}, w http.ResponseWriter) {
	w.Header().Set("Content-Type", "application/json")
	err := json.NewEncoder(w).Encode(data)
	if err != nil {
		http.Error(w, "Internal server error", 500)
		slog.Error("Failed to encode JSON: " + err.Error())
		return
	}
}

//goland:noinspection SqlNoDataSourceInspection,GoDfaErrorMayBeNotNil,GoDfaNilDereference
func main() {
	var publicKey ed25519.PublicKey
	var privateKey ed25519.PrivateKey
	privateKey, err := os.ReadFile("key.pem")
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			publicKey, privateKey, err = ed25519.GenerateKey(nil)
			if err != nil {
				slog.Error("Failed to generate private key: " + err.Error())
				os.Exit(1)
			}

			err = os.WriteFile("key.pem", privateKey, 0600)
			if err != nil {
				slog.Error("Failed to write private key: " + err.Error())
				os.Exit(1)
			}
		} else {
			slog.Error("Failed to read private key: " + err.Error())
			os.Exit(1)
		}
	} else {
		publicKey = privateKey.Public().(ed25519.PublicKey)
	}

	conn, err := sql.Open("sqlite3", "database.db")
	if err != nil || conn == nil {
		slog.Error("Failed to open database: " + err.Error())
		os.Exit(1)
	}

	_, err = conn.Exec("CREATE TABLE IF NOT EXISTS users (id BLOB NOT NULL PRIMARY KEY, username TEXT NOT NULL UNIQUE, publicKey BLOB NOT NULL, administrator INTEGER NOT NULL DEFAULT 0 CHECK(administrator IN (0, 1)))")
	if err != nil {
		slog.Error("Failed to create users table: " + err.Error())
		os.Exit(1)
	}

	_, err = conn.Exec("CREATE TABLE IF NOT EXISTS invites (code BLOB NOT NULL UNIQUE, uses INTEGER NOT NULL)")
	if err != nil {
		slog.Error("Failed to create invites table: " + err.Error())
		os.Exit(1)
	}

	_, err = conn.Exec("CREATE TABLE IF NOT EXISTS messages (sender BLOB NOT NULL, senderName TEXT NOT NULL, message TEXT NOT NULL, sent INTEGER NOT NULL, id BLOB NOT NULL PRIMARY KEY)")
	if err != nil {
		slog.Error("Failed to create messages table: " + err.Error())
		os.Exit(1)
	}

	router := chi.NewRouter()
	subscriptions := make(map[time.Time]chan bool)

	router.Post("/api/signup", func(w http.ResponseWriter, r *http.Request) {
		var data struct {
			Username  string `json:"username"`
			PublicKey string `json:"publicKey"`
			Invite    string `json:"invite"`
		}
		err := json.NewDecoder(r.Body).Decode(&data)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid request"}, w)
			return
		}

		if data.Username == "" || data.PublicKey == "" || data.Invite == "" {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid request"}, w)
			return
		}

		inviteBytes, err := hex.DecodeString(data.Invite)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid invite"}, w)
			return
		}

		var uses int
		err = conn.QueryRow("SELECT uses FROM invites WHERE code = ?", inviteBytes).Scan(&uses)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid invite"}, w)
			return
		}

		_, err = conn.Exec("UPDATE invites SET uses = uses - 1 WHERE code = ?", inviteBytes)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to update invite: " + err.Error())
			return
		}

		if uses-1 <= 0 {
			_, err := conn.Exec("DELETE FROM invites WHERE uses <= 0")
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
				slog.Error("Failed to delete invite: " + err.Error())
				return
			}
			if uses <= 0 {
				w.WriteHeader(http.StatusBadRequest)
				renderJSON(map[string]interface{}{"error": "Invalid invite"}, w)
				return
			}
		}

		userId, err := uuid.NewRandom()
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to generate UUID: " + err.Error())
			return
		}

		publicKeyBytes, err := base64.StdEncoding.DecodeString(data.PublicKey)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid public key"}, w)
			return
		}

		_, err = conn.Exec("INSERT INTO users (id, username, publicKey) VALUES (?, ?, ?)", userId, data.Username, publicKeyBytes)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Username already taken"}, w)
			return
		}

		var nonce [32]byte
		read, err := rand.Read(nonce[:])
		if err != nil || read != 32 {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to generate nonce: " + err.Error())
			return
		}

		token, err := jwt.NewWithClaims(jwt.SigningMethodEdDSA, jwt.MapClaims{
			"sub":   userId.String(),
			"name":  data.Username,
			"exp":   time.Now().Add(time.Hour * 24).Unix(),
			"nonce": base64.StdEncoding.EncodeToString(nonce[:]),
		}).SignedString(privateKey)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to sign token: " + err.Error())
			return
		}

		w.WriteHeader(http.StatusOK)
		renderJSON(map[string]interface{}{"token": token, "userId": userId.String()}, w)
	})

	router.Post("/api/loginChallenge", func(w http.ResponseWriter, r *http.Request) {
		var data struct {
			Username string `json:"username"`
		}
		err := json.NewDecoder(r.Body).Decode(&data)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid request"}, w)
			return
		}

		var nonce [32]byte
		read, err := rand.Read(nonce[:])
		if err != nil || read != 32 {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to generate nonce: " + err.Error())
			return
		}

		token, err := jwt.NewWithClaims(jwt.SigningMethodEdDSA, jwt.MapClaims{
			"sub":   data.Username,
			"exp":   time.Now().Add(time.Second * 20).Unix(),
			"nonce": base64.StdEncoding.EncodeToString(nonce[:]),
		}).SignedString(privateKey)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to sign token: " + err.Error())
			return
		}

		w.WriteHeader(http.StatusOK)
		renderJSON(map[string]interface{}{"token": token, "nonce": base64.StdEncoding.EncodeToString(nonce[:])}, w)
	})

	router.Post("/api/login", func(w http.ResponseWriter, r *http.Request) {
		var data struct {
			Username  string `json:"username"`
			Signature string `json:"signature"`
			Challenge string `json:"challenge"`
		}
		err := json.NewDecoder(r.Body).Decode(&data)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid request"}, w)
			return
		}

		claims, ok := verifyJwt(data.Challenge, publicKey)
		if !ok {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid challenge"}, w)
			return
		}

		nonce, err := base64.StdEncoding.DecodeString(claims["nonce"].(string))
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid challenge"}, w)
			return
		}

		signature, err := base64.StdEncoding.DecodeString(data.Signature)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid signature"}, w)
			return
		}

		var userKey []byte
		err = conn.QueryRow("SELECT publicKey FROM users WHERE username = ?", data.Username).Scan(&userKey)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			fmt.Println(err.Error())
			renderJSON(map[string]interface{}{"error": "Invalid username"}, w)
			return
		}

		if !ed25519.Verify(userKey, nonce, signature) {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid signature"}, w)
			return
		}

		var userId uuid.UUID
		err = conn.QueryRow("SELECT id FROM users WHERE username = ?", data.Username).Scan(&userId)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid username"}, w)
			return
		}

		var loginNonce [32]byte
		read, err := rand.Read(loginNonce[:])
		if err != nil || read != 32 {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to generate nonce: " + err.Error())
			return
		}

		token, err := jwt.NewWithClaims(jwt.SigningMethodEdDSA, jwt.MapClaims{
			"sub":   userId.String(),
			"name":  data.Username,
			"exp":   time.Now().Add(time.Hour * 24).Unix(),
			"nonce": base64.StdEncoding.EncodeToString(loginNonce[:]),
		}).SignedString(privateKey)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to sign token: " + err.Error())
			return
		}

		w.WriteHeader(http.StatusOK)
		renderJSON(map[string]interface{}{"token": token, "userId": userId.String()}, w)
	})

	router.Post("/api/invite", func(w http.ResponseWriter, r *http.Request) {
		var data struct {
			Token string  `json:"token"`
			Uses  float64 `json:"uses"`
		}

		err := json.NewDecoder(r.Body).Decode(&data)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid request"}, w)
			return
		}

		claims, ok := verifyJwt(data.Token, publicKey)
		if !ok {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid token"}, w)
			return
		}

		userId, err := uuid.Parse(claims["sub"].(string))
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid token"}, w)
			return
		}

		var admin bool
		err = conn.QueryRow("SELECT administrator FROM users WHERE id = ?", userId).Scan(&admin)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid token"}, w)
			return
		}

		if !admin {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid token"}, w)
			return
		}

		var code [8]byte
		read, err := rand.Read(code[:])
		if err != nil || read != 8 {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to generate code: " + err.Error())
			return
		}

		_, err = conn.Exec("INSERT INTO invites (code, uses) VALUES (?, ?)", code[:], data.Uses)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to insert invite: " + err.Error())
			return
		}

		w.WriteHeader(http.StatusOK)
		renderJSON(map[string]interface{}{"code": hex.EncodeToString(code[:])}, w)
	})

	router.Post("/api/messages", func(w http.ResponseWriter, r *http.Request) {
		var data struct {
			Token string `json:"token"`
		}

		err := json.NewDecoder(r.Body).Decode(&data)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid request"}, w)
			return
		}

		_, ok := verifyJwt(data.Token, publicKey)
		if !ok {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid token"}, w)
			return
		}

		messages := make(map[string]interface{})
		rows, err := conn.Query("SELECT sender, senderName, message, sent, id FROM messages ORDER BY sent ASC")
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to select messages: " + err.Error())
			return
		}

		for rows.Next() {
			var sender uuid.UUID
			var senderName string
			var message string
			var sent float64
			var id uuid.UUID
			err = rows.Scan(&sender, &senderName, &message, &sent, &id)
			if err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
				slog.Error("Failed to scan message: " + err.Error())
				return
			}

			messages[id.String()] = map[string]interface{}{
				"sender":  sender.String(),
				"name":    senderName,
				"message": message,
				"sent":    sent,
			}
		}

		w.WriteHeader(http.StatusOK)
		renderJSON(messages, w)
	})

	router.Post("/api/send", func(w http.ResponseWriter, r *http.Request) {
		var data struct {
			Token   string `json:"token"`
			Message string `json:"message"`
		}

		err := json.NewDecoder(r.Body).Decode(&data)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid request"}, w)
			return
		}

		claims, ok := verifyJwt(data.Token, publicKey)
		if !ok {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid token"}, w)
			return
		}

		userId, err := uuid.Parse(claims["sub"].(string))
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid token"}, w)
			return
		}

		messageId, err := uuid.NewRandom()
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to generate UUID: " + err.Error())
			return
		}

		_, err = conn.Exec("INSERT INTO messages (sender, message, sent, id, senderName) VALUES (?, ?, ?, ?, ?)", userId, data.Message, time.Now().Unix(), messageId, claims["name"])
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to insert message: " + err.Error())
			return
		}

		go func() {
			for subscriptionId, subscription := range subscriptions {
				select {
				case subscription <- false:
				default:
					subscriptions[subscriptionId] = nil
				}
			}
		}()

		w.WriteHeader(http.StatusOK)
	})

	router.Post("/api/delete", func(w http.ResponseWriter, r *http.Request) {
		var data struct {
			Token string `json:"token"`
			Id    string `json:"id"`
		}
		err := json.NewDecoder(r.Body).Decode(&data)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid request"}, w)
			return
		}

		claims, ok := verifyJwt(data.Token, publicKey)
		if !ok {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid token"}, w)
			return
		}

		userId, err := uuid.Parse(claims["sub"].(string))
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid token"}, w)
			return
		}

		messageId, err := uuid.Parse(data.Id)
		if err != nil {
			w.WriteHeader(http.StatusBadRequest)
			renderJSON(map[string]interface{}{"error": "Invalid message ID"}, w)
			return
		}

		_, err = conn.Exec("DELETE FROM messages WHERE sender = ? AND id = ?", userId, messageId)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			renderJSON(map[string]interface{}{"error": "Internal server error"}, w)
			slog.Error("Failed to delete message: " + err.Error())
			return
		}

		go func() {
			for subscriptionId, subscription := range subscriptions {
				select {
				case subscription <- false:
				default:
					subscriptions[subscriptionId] = nil
				}
			}
		}()

		w.WriteHeader(http.StatusOK)
	})

	router.HandleFunc("/api/ping", func(w http.ResponseWriter, r *http.Request) {
		slog.Info("New client online at " + r.RemoteAddr)
		r.Host = "example.org"
		r.Header.Set("Origin", "https://example.org")
		connection, err := websocket.Accept(w, r, nil)
		if err != nil {
			slog.Error("Failed to accept WebSocket connection: " + err.Error())
			return
		}

		subscriptionId := time.Now()

		subscription := make(chan bool)
		subscriptions[subscriptionId] = subscription

		go func() {
			for {
				shouldClose := <-subscription
				if shouldClose {
					return
				} else {
					err := connection.Write(context.Background(), websocket.MessageBinary, []byte{0x00})
					if err != nil {
						println(err.Error())
						slog.Info("Client disconnected: " + r.RemoteAddr)
						subscriptions[subscriptionId] = nil
						return
					}
				}
			}
		}()
	})

	err = http.ListenAndServe(":1974", router)
	if err != nil {
		slog.Error("Failed to start server: " + err.Error())
		os.Exit(1)
	}
}