ruby-sea-battle/game-server/internal/ws/handler_test.go

package ws_test

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/coder/websocket"
	"github.com/golang-jwt/jwt/v5"

	"github.com/ruby-sea-battle/game-server/internal/auth"
	"github.com/ruby-sea-battle/game-server/internal/game"
	"github.com/ruby-sea-battle/game-server/internal/ws"
)

var secret = []byte("integration-secret")

const publicAddr = "ws://game.example.com/ws"

func gameToken(t *testing.T, matchID string, p1, p2 int64, gameServer string, exp time.Time) string {
	t.Helper()
	token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, &auth.GameClaims{
		MatchID:    matchID,
		Player1ID:  p1,
		Player2ID:  p2,
		GameServer: gameServer,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(exp),
		},
	}).SignedString(secret)
	if err != nil {
		t.Fatalf("sign game token: %v", err)
	}
	return token
}

func userToken(t *testing.T, userID int64) string {
	t.Helper()
	token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, &auth.UserClaims{
		UserID: userID,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(time.Now().Add(time.Hour)),
		},
	}).SignedString(secret)
	if err != nil {
		t.Fatalf("sign user token: %v", err)
	}
	return token
}

func startServer(t *testing.T, joinTimeout time.Duration) (url string) {
	t.Helper()
	manager := game.NewManager(joinTimeout)
	log := slog.New(slog.NewTextHandler(testWriter{t}, nil))
	handler := ws.NewHandler(manager, secret, publicAddr, []string{"*"}, log)
	srv := httptest.NewServer(handler)
	t.Cleanup(srv.Close)
	return "ws" + strings.TrimPrefix(srv.URL, "http")
}

type testWriter struct{ t *testing.T }

func (w testWriter) Write(p []byte) (int, error) {
	w.t.Log(strings.TrimSpace(string(p)))
	return len(p), nil
}

func dial(t *testing.T, url string) *websocket.Conn {
	t.Helper()
	c, _, err := websocket.Dial(context.Background(), url, nil)
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	t.Cleanup(func() { c.Close(websocket.StatusNormalClosure, "") })
	return c
}

func join(t *testing.T, c *websocket.Conn, gt, ut string) {
	t.Helper()
	msg, _ := json.Marshal([]any{1, gt, ut})
	if err := c.Write(context.Background(), websocket.MessageText, msg); err != nil {
		t.Fatalf("write JOIN: %v", err)
	}
}

func readMsg(t *testing.T, c *websocket.Conn) []any {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, data, err := c.Read(ctx)
	if err != nil {
		t.Fatalf("read: %v", err)
	}
	var v []any
	if err := json.Unmarshal(data, &v); err != nil {
		t.Fatalf("unmarshal %s: %v", data, err)
	}
	return v
}

func TestTwoPlayersStartMatch(t *testing.T) {
	url := startServer(t, time.Minute)
	exp := time.Now().Add(time.Hour)
	gt := gameToken(t, "m-1", 1, 2, publicAddr, exp)

	c1, c2 := dial(t, url), dial(t, url)
	join(t, c1, gt, userToken(t, 1))
	join(t, c2, gt, userToken(t, 2))

	for i, c := range []*websocket.Conn{c1, c2} {
		v := readMsg(t, c)
		if v[0] != float64(3) || v[1] != float64(1) || v[2] != float64(1) {
			t.Fatalf("player %d: unexpected message %v", i, v)
		}
		// ACK the state.
		msg, _ := json.Marshal([]any{4, 1})
		if err := c.Write(context.Background(), websocket.MessageText, msg); err != nil {
			t.Fatalf("write ACK: %v", err)
		}
	}
}

func TestWrongServerAddressRejected(t *testing.T) {
	url := startServer(t, time.Minute)
	gt := gameToken(t, "m-2", 1, 2, "ws://other.example.com/ws", time.Now().Add(time.Hour))

	c := dial(t, url)
	join(t, c, gt, userToken(t, 1))

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	_, _, err := c.Read(ctx)
	if err == nil {
		t.Fatal("expected connection to be closed")
	}
	if code := websocket.CloseStatus(err); code != 4003 {
		t.Fatalf("close code = %d, want 4003", code)
	}
}

func TestJoinTimeoutEndsMatchAndReplay(t *testing.T) {
	url := startServer(t, 200*time.Millisecond)
	exp := time.Now().Add(time.Hour)
	gt := gameToken(t, "m-3", 1, 2, publicAddr, exp)

	c1 := dial(t, url)
	join(t, c1, gt, userToken(t, 1))

	v := readMsg(t, c1) // END
	if v[0] != float64(2) || v[1] != float64(1) || v[2] != float64(-1) || v[3] != "join timeout exceeded" {
		t.Fatalf("unexpected END: %v", v)
	}

	// Late join with the same token replays END and closes.
	c2 := dial(t, url)
	join(t, c2, gt, userToken(t, 2))
	v = readMsg(t, c2)
	if v[0] != float64(2) || v[2] != float64(-1) {
		t.Fatalf("unexpected replayed END: %v", v)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if _, _, err := c2.Read(ctx); err == nil {
		t.Fatal("expected connection to be closed after END replay")
	}
}

func TestReconnectDuringGame(t *testing.T) {
	url := startServer(t, time.Minute)
	exp := time.Now().Add(time.Hour)
	gt := gameToken(t, "m-4", 1, 2, publicAddr, exp)

	c1, c2 := dial(t, url), dial(t, url)
	join(t, c1, gt, userToken(t, 1))
	join(t, c2, gt, userToken(t, 2))
	readMsg(t, c1) // STATE, not acknowledged

	// Simulate a dropped connection.
	c1.Close(websocket.StatusAbnormalClosure, "network gone")

	// Reconnect with the same tokens: unacknowledged STATE is replayed.
	c1b := dial(t, url)
	join(t, c1b, gt, userToken(t, 1))
	v := readMsg(t, c1b)
	if v[0] != float64(3) || v[1] != float64(1) || v[2] != float64(1) {
		t.Fatalf("unexpected replayed message: %v", v)
	}
}