Merge pull request #45 from Quaver/test-env

Add testing environment

Author: Swan

cmd/server/server.go

@@ -65,6 +65,9 @@ func (s *Server) Start() {
 				utils.CloseConnection(conn)
 				return
 			}
+		} else {
+			_ = conn.Close()
+			return
 		}
 
 		// Handle various connection events

cmd/tests/client.go

@@ -0,0 +1,110 @@
+package tests
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"example.com/Quaver/Z/config"
+	"example.com/Quaver/Z/db"
+	"example.com/Quaver/Z/handlers"
+	"fmt"
+	"github.com/gorilla/websocket"
+	"log"
+	"net/url"
+	"os"
+	"os/signal"
+)
+
+type Client struct {
+	User          *db.User
+	Conn          *websocket.Conn
+	enableLogging bool
+}
+
+func newClient(user *db.User) *Client {
+	return &Client{User: user}
+}
+
+func (client *Client) host() string {
+	return fmt.Sprintf("localhost:%v", config.Instance.Server.Port)
+}
+
+func (client *Client) loginData() handlers.LoginData {
+	return handlers.LoginData{
+		Id:        client.User.SteamId,
+		PTicket:   "aGVsbG8=",
+		PcbTicket: 0,
+		Client:    "1|2|3|4|5",
+	}
+}
+
+func (client *Client) loginDataEncoded() string {
+	data := client.loginData()
+	dataJson, _ := json.Marshal(data)
+
+	var encoded = make([]byte, base64.StdEncoding.EncodedLen(len(dataJson)))
+	base64.StdEncoding.Encode(encoded, dataJson)
+
+	return string(encoded)
+}
+func (client *Client) login() {
+	u := url.URL{
+		Scheme:   "ws",
+		Host:     client.host(),
+		Path:     "/",
+		RawQuery: fmt.Sprintf("login=%v", client.loginDataEncoded()),
+	}
+
+	var err error
+
+	client.Conn, _, err = websocket.DefaultDialer.Dial(u.String(), nil)
+
+	if err != nil {
+		log.Fatal("dial:", err)
+	}
+
+	defer client.Conn.Close()
+	client.handleInterrupt(client.readMessages())
+}
+
+func (client *Client) readMessages() chan struct{} {
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+
+		for {
+			_, message, err := client.Conn.ReadMessage()
+
+			if err != nil {
+				log.Println("read:", err)
+				return
+			}
+
+			if client.enableLogging {
+				log.Printf("recv: %s", message)
+			}
+		}
+	}()
+
+	return done
+}
+
+func (client *Client) handleInterrupt(doneReadingChan chan struct{}) {
+	interrupt := make(chan os.Signal, 1)
+	signal.Notify(interrupt, os.Interrupt)
+
+	for {
+		select {
+		case <-doneReadingChan:
+			return
+		case <-interrupt:
+			err := client.Conn.WriteMessage(websocket.CloseMessage,
+				websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
+
+			if err != nil {
+				log.Println("write close:", err)
+				return
+			}
+		}
+	}
+}

cmd/tests/stress_test.go

@@ -0,0 +1,42 @@
+package tests
+
+import (
+	"example.com/Quaver/Z/config"
+	"example.com/Quaver/Z/db"
+	"sync"
+	"testing"
+)
+
+// Tests a 100 users connecting to the server at one time
+func TestLogin100Clients(t *testing.T) {
+	if err := config.Load("../../config.json"); err != nil {
+		t.Fatal(err)
+	}
+
+	db.InitializeSQL()
+
+	wg := sync.WaitGroup{}
+
+	for i := 3; i < 103; i++ {
+		wg.Add(1)
+
+		i := i
+
+		go func() {
+			defer wg.Done()
+
+			user, err := db.GetUserById(i)
+
+			if err != nil {
+				t.Error(err)
+				return
+			}
+
+			client := newClient(user)
+			client.login()
+		}()
+	}
+
+	wg.Wait()
+	db.CloseSQLConnection()
+}

db/sql.go

@@ -34,8 +34,13 @@ func InitializeSQL() {
 	if err != nil {
 		log.Fatalln(err)
 	}
-	
+
 	SQL = db
+
+	SQL.DB.SetMaxOpenConns(100)
+	SQL.DB.SetMaxIdleConns(10)
+	SQL.DB.SetConnMaxLifetime(0)
+
 	log.Println("Successfully connected to SQL database")
 }
 

db/users.go

@@ -28,6 +28,20 @@ func (u *User) GetProfileUrl() string {
 	return fmt.Sprintf("https://quavergame.com/user/%v", u.Id)
 }
 
+// GetUserById Retrieves a user from the database by their id
+func GetUserById(id int) (*User, error) {
+	query := "SELECT id, steam_id, username, allowed, privileges, usergroups, mute_endtime, country, avatar_url, twitch_username FROM users WHERE id = ? LIMIT 1"
+
+	var user User
+	err := SQL.Get(&user, query, id)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &user, nil
+}
+
 // GetUserBySteamId Retrieves a user from the database by their Steam id
 func GetUserBySteamId(steamId string) (*User, error) {
 	query := "SELECT id, steam_id, username, allowed, privileges, usergroups, mute_endtime, country, avatar_url, twitch_username FROM users WHERE steam_id = ? LIMIT 1"

go.mod

@@ -22,6 +22,7 @@ require (
 	github.com/disgoorg/snowflake/v2 v2.0.1 // indirect
 	github.com/gobwas/httphead v0.1.0 // indirect
 	github.com/gobwas/pool v0.2.1 // indirect
+	github.com/gorilla/websocket v1.5.3 // indirect
 	github.com/sasha-s/go-csync v0.0.0-20210812194225-61421b77c44b // indirect
 	golang.org/x/net v0.23.0 // indirect
 	golang.org/x/sys v0.18.0 // indirect

go.sum

@@ -31,6 +31,8 @@ github.com/gobwas/pool v0.2.1 h1:xfeeEhW7pwmX8nuLVlqbzVc7udMDrwetjEv+TZIz1og=
 github.com/gobwas/pool v0.2.1/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
 github.com/gobwas/ws v1.2.0 h1:u0p9s3xLYpZCA1z5JgCkMeB34CKCMMQbM+G8Ii7YD0I=
 github.com/gobwas/ws v1.2.0/go.mod h1:hRKAFb8wOxFROYNsT1bqfWnhX+b5MFeJM9r2ZSwg/KY=
+github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
+github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/jmoiron/sqlx v1.3.5 h1:vFFPA71p1o5gAeqtEAwLU4dnX2napprKtHr7PYIcN3g=
 github.com/jmoiron/sqlx v1.3.5/go.mod h1:nRVWtLre0KfCLJvgxzCsLVMogSvQ1zNJtpYr2Ccp0mQ=
 github.com/lib/pq v1.2.0 h1:LXpIM/LZ5xGFhOpXAQUIMM1HdyqzVYM13zNdjCEEcA0=

handlers/login.go

@@ -345,6 +345,10 @@ func canUserUseCustomGameBuild(user *db.User) bool {
 
 // Sends webhook and disconnects a user for invalid client usage. Returns if the user is allowed to login
 func handleCustomGameBuildUsage(conn net.Conn, user *db.User, client string) bool {
+	if config.Instance.BypassSteamLogin {
+		return true
+	}
+
 	clientStr := fmt.Sprintf("```json\n%v```", formatCustomGameBuild(client))
 	webhooks.SendAntiCheat(user.Username, user.Id, user.GetProfileUrl(), user.AvatarUrl.String, "Invalid Game Build", clientStr)