154 lines
3.2 KiB
Go
154 lines
3.2 KiB
Go
package main
|
|
|
|
import (
|
|
"database/sql"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"strconv"
|
|
"sync"
|
|
|
|
"gossip/pkg/protocol"
|
|
|
|
"github.com/gorilla/websocket"
|
|
_ "github.com/jackc/pgx/v5/stdlib"
|
|
"github.com/joho/godotenv"
|
|
)
|
|
|
|
type Server struct {
|
|
db *sql.DB
|
|
clients map[string]*websocket.Conn
|
|
mu sync.Mutex
|
|
upgrader websocket.Upgrader
|
|
}
|
|
|
|
func main() {
|
|
if err := godotenv.Load(); err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
dbuser := os.Getenv("DBUSER")
|
|
dbname := os.Getenv("DBNAME")
|
|
dbport := os.Getenv("DBPORT")
|
|
dbhost := os.Getenv("DBHOST")
|
|
dbpass := os.Getenv("DBPASS")
|
|
|
|
connStr := fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=require", dbuser, dbpass, dbhost, dbport, dbname)
|
|
|
|
db, err := sql.Open("pgx", connStr)
|
|
if err != nil {
|
|
log.Fatal("Failed to open DB:", err)
|
|
}
|
|
defer db.Close()
|
|
|
|
if err := db.Ping(); err != nil {
|
|
log.Fatal("Failed to connect to DB:", err)
|
|
}
|
|
|
|
srv := &Server{
|
|
db: db,
|
|
clients: make(map[string]*websocket.Conn),
|
|
upgrader: websocket.Upgrader{
|
|
CheckOrigin: func(r *http.Request) bool { return true },
|
|
},
|
|
}
|
|
|
|
http.HandleFunc("/ws", srv.handleWS)
|
|
log.Println("Gossip Relay Server listening on :8080")
|
|
log.Fatal(http.ListenAndServe(":8080", nil))
|
|
}
|
|
|
|
func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
|
|
conn, err := s.upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
log.Println("Upgrade error:", err)
|
|
return
|
|
}
|
|
defer conn.Close()
|
|
|
|
var myPubKey string
|
|
|
|
for {
|
|
_, data, err := conn.ReadMessage()
|
|
if err != nil {
|
|
break
|
|
}
|
|
|
|
var msg protocol.Message
|
|
if err := json.Unmarshal(data, &msg); err != nil {
|
|
continue
|
|
}
|
|
|
|
switch msg.Type {
|
|
case protocol.TypeLogin:
|
|
myPubKey = msg.Sender
|
|
|
|
s.mu.Lock()
|
|
s.clients[myPubKey] = conn
|
|
s.mu.Unlock()
|
|
|
|
accNum, err := s.getOrCreateUser(myPubKey)
|
|
if err != nil {
|
|
log.Println("DB Error:", err)
|
|
continue
|
|
}
|
|
|
|
resp := protocol.Message{
|
|
Type: protocol.TypeIdentity,
|
|
Target: myPubKey,
|
|
Content: strconv.Itoa(accNum),
|
|
}
|
|
conn.WriteJSON(resp)
|
|
log.Printf("User %d connected (%s...)", accNum, myPubKey[:8])
|
|
|
|
case protocol.TypeLookup:
|
|
targetAcc, _ := strconv.Atoi(msg.Content)
|
|
foundKey, ok := s.lookupKey(targetAcc)
|
|
|
|
resp := protocol.Message{
|
|
Type: protocol.TypeLookupResponse,
|
|
Target: msg.Content,
|
|
}
|
|
if ok {
|
|
resp.Content = foundKey
|
|
}
|
|
conn.WriteJSON(resp)
|
|
|
|
case protocol.TypeMsg:
|
|
s.mu.Lock()
|
|
targetConn, ok := s.clients[msg.Target]
|
|
s.mu.Unlock()
|
|
|
|
if ok {
|
|
targetConn.WriteMessage(websocket.TextMessage, data)
|
|
}
|
|
}
|
|
}
|
|
|
|
if myPubKey != "" {
|
|
s.mu.Lock()
|
|
delete(s.clients, myPubKey)
|
|
s.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
func (s *Server) getOrCreateUser(pubKey string) (int, error) {
|
|
var accNum int
|
|
err := s.db.QueryRow("SELECT account_number FROM users WHERE public_key=$1", pubKey).Scan(&accNum)
|
|
if err == nil {
|
|
return accNum, nil
|
|
}
|
|
err = s.db.QueryRow("INSERT INTO users (public_key) VALUES ($1) RETURNING account_number", pubKey).Scan(&accNum)
|
|
return accNum, err
|
|
}
|
|
|
|
func (s *Server) lookupKey(accNum int) (string, bool) {
|
|
var pubKey string
|
|
err := s.db.QueryRow("SELECT public_key FROM users WHERE account_number=$1", accNum).Scan(&pubKey)
|
|
if err != nil {
|
|
return "", false
|
|
}
|
|
return pubKey, true
|
|
}
|