glchat/internal/gateway/session.go

package gateway

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/coder/websocket"
)

// clientSession — одно подключение клиента (AGENT.md 8.3: сессий может быть
// несколько, события доставляются во все).
type clientSession struct {
	id        string
	userID    uint64
	conn      *websocket.Conn
	send      chan outbound
	closeOnce sync.Once
	closed    chan struct{}
	buffer    *resumeBuffer
}

// outbound — кадр в очереди на отправку. done используется, когда отправитель
// ждёт фактической записи в сокет (например, перед закрытием соединения).
type outbound struct {
	payload []byte
	done    chan error
}

func (s *clientSession) write(payload []byte) error {
	return s.enqueue(outbound{payload: payload})
}

// writeSync ставит кадр в очередь и дожидается результата записи.
func (s *clientSession) writeSync(ctx context.Context, payload []byte) error {
	item := outbound{payload: payload, done: make(chan error, 1)}
	if err := s.enqueue(item); err != nil {
		return err
	}
	select {
	case err := <-item.done:
		return err
	case <-ctx.Done():
		return ctx.Err()
	case <-s.closed:
		return errors.New("session closed")
	}
}

func (s *clientSession) enqueue(item outbound) error {
	select {
	case s.send <- item:
		return nil
	case <-s.closed:
		return errors.New("session closed")
	case <-time.After(WriteTimeout):
		// Медленный потребитель: отключаем, чтобы не держать буфер (AGENT.md 8.3).
		s.close()
		return errors.New("slow consumer")
	}
}

func (s *clientSession) close() {
	s.closeOnce.Do(func() {
		close(s.closed)
		_ = s.conn.Close(websocket.StatusNormalClosure, "closed")
	})
}

// ServeHTTP обрабатывает подключение к /gateway.
func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	origin := r.Header.Get("Origin")
	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
		OriginPatterns: s.originPatterns(),
		// Клиенты десктопа обращаются без Origin: проверяем вручную ниже.
		InsecureSkipVerify: origin == "",
		CompressionMode:    websocket.CompressionDisabled,
	})
	if err != nil {
		s.logger.DebugContext(r.Context(), "gateway accept failed", slog.Any("error", err))
		return
	}
	defer func() { _ = conn.CloseNow() }()

	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()

	session := &clientSession{
		id:     newSessionID(),
		conn:   conn,
		send:   make(chan outbound, 128),
		closed: make(chan struct{}),
	}
	conn.SetReadLimit(MaxIncomingFrame)

	// HELLO с интервалом heartbeat и идентификатором сессии.
	hello, err := frame(OpHello, map[string]any{
		"heartbeat_interval_ms": HeartbeatInterval.Milliseconds(),
		"session_id":            session.id,
	})
	if err != nil {
		return
	}
	if err := writeDirect(ctx, conn, hello); err != nil {
		return
	}

	writerDone := make(chan struct{})
	go func() {
		defer close(writerDone)
		s.writeLoop(ctx, session)
	}()

	s.readLoop(ctx, session)
	session.close()
	cancel()
	<-writerDone

	s.unregister(session)
}

func (s *Service) writeLoop(ctx context.Context, session *clientSession) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-session.closed:
			return
		case item := <-session.send:
			err := writeDirect(ctx, session.conn, item.payload)
			if item.done != nil {
				item.done <- err
				close(item.done)
			}
			if err != nil {
				session.close()
				return
			}
		}
	}
}

// readLoop читает кадры клиента: IDENTIFY, RESUME, HEARTBEAT.
func (s *Service) readLoop(ctx context.Context, session *clientSession) {
	identified := false
	lastIdentify := time.Time{}

	for {
		_, data, err := session.conn.Read(ctx)
		if err != nil {
			return
		}
		var incoming serverFrame
		if err := json.Unmarshal(data, &incoming); err != nil {
			s.sendInvalidSession(session, "malformed frame")
			return
		}

		switch incoming.Op {
		case OpHeartbeat:
			ack, err := frame(OpHeartbeatAck, nil)
			if err == nil {
				_ = session.write(ack)
			}
			if !identified {
				// До IDENTIFY сердцебиение не подтверждает сессию: ждём идентификацию.
				continue
			}
		case OpIdentify, OpResume:
			if time.Since(lastIdentify) < IdentifyRateLimit {
				s.sendInvalidSession(session, "identify rate limited")
				return
			}
			lastIdentify = time.Now()

			payload, err := decodeIdentify(incoming.D)
			if err != nil {
				s.sendInvalidSession(session, "malformed identify")
				return
			}
			if err := s.handleIdentify(ctx, session, payload, incoming.Op == OpResume); err != nil {
				s.logger.DebugContext(ctx, "gateway identify failed", slog.Any("error", err))
				return
			}
			identified = true
		default:
			// Неизвестные оп-коды игнорируем: протокол расширяемый.
		}
		if !identified {
			continue
		}
	}
}

// sendInvalidSession сообщает об ошибке идентификации и дожидается отправки
// кадра: соединение закрывается сразу после этого (AGENT.md 8.3).
func (s *Service) sendInvalidSession(session *clientSession, reason string) {
	payload, err := frame(OpInvalidSess, map[string]any{"reason": reason})
	if err != nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), WriteTimeout)
	defer cancel()
	if err := session.writeSync(ctx, payload); err != nil {
		s.logger.DebugContext(ctx, "gateway invalid session not delivered", slog.Any("error", err))
	}
}

// handleIdentify проверяет токен, регистрирует сессию и отправляет READY
// либо догоняет пропущенные события при RESUME.
func (s *Service) handleIdentify(ctx context.Context, session *clientSession, payload identifyPayload, resume bool) error {
	user, authSession, err := s.auth.ResolveSession(ctx, payload.Token)
	if err != nil {
		s.sendInvalidSession(session, "invalid token")
		return err
	}
	session.userID = user.ID
	session.buffer = s.bufferFor(authSession.TokenHash)

	if resume && payload.ResumeSeq > 0 && session.buffer != nil {
		if missed, ok := session.buffer.since(payload.ResumeSeq); ok {
			for _, item := range missed {
				if err := session.write(item); err != nil {
					return err
				}
			}
			// Подтверждаем продолжение сессии: клиент знает, что состояние
			// восстановлено без полного READY (AGENT.md 8.3).
			s.mu.Lock()
			s.seq++
			resumed := Envelope{Op: OpDispatch, T: "RESUMED", S: s.seq}
			s.mu.Unlock()
			encoded, err := json.Marshal(resumed)
			if err != nil {
				return err
			}
			if err := session.write(encoded); err != nil {
				return err
			}
			s.register(session, authSession.TokenHash)
			return nil
		}
		// Буфер не содержит нужную последовательность: полный READY (AGENT.md 8.3).
	}

	snapshot, err := s.readiness.Build(ctx, user)
	if err != nil {
		s.sendInvalidSession(session, "ready failed")
		return err
	}
	snapshot.SessionID = session.id
	snapshot.HeartbeatMS = int(HeartbeatInterval.Milliseconds())

	raw, err := json.Marshal(snapshot)
	if err != nil {
		return err
	}
	s.mu.Lock()
	s.seq++
	envelope := Envelope{Op: OpDispatch, T: "READY", D: raw, S: s.seq}
	encoded, err := json.Marshal(envelope)
	s.mu.Unlock()
	if err != nil {
		return err
	}
	if err := session.write(encoded); err != nil {
		return err
	}
	s.register(session, authSession.TokenHash)
	return nil
}

func (s *Service) register(session *clientSession, tokenHash string) {
	s.mu.Lock()
	s.sessions[session.id] = session
	buffer, ok := s.buffers[tokenHash]
	if !ok {
		buffer = newResumeBuffer(ResumeBufferSize)
		s.buffers[tokenHash] = buffer
	}
	// Владелец буфера нужен, чтобы доставлять адресные события (SendToUser)
	// даже когда все сессии пользователя отключены.
	buffer.owner = session.userID
	session.buffer = buffer
	s.mu.Unlock()
}

func (s *Service) unregister(session *clientSession) {
	s.mu.Lock()
	delete(s.sessions, session.id)
	s.mu.Unlock()
}

func (s *Service) bufferFor(tokenHash string) *resumeBuffer {
	s.mu.Lock()
	defer s.mu.Unlock()
	if buffer, ok := s.buffers[tokenHash]; ok {
		return buffer
	}
	buffer := newResumeBuffer(ResumeBufferSize)
	s.buffers[tokenHash] = buffer
	return buffer
}

// originPatterns разрешает подключения с домена инстанса (AGENT.md 9.7).
// coder/websocket сопоставляет шаблон с хостом из Origin, поэтому из
// настроенных адресов берём только хост; пустой список означает проверку
// на совпадение с Host запроса.
func (s *Service) originPatterns() []string {
	patterns := make([]string, 0, len(s.allowedOrigins))
	for _, origin := range s.allowedOrigins {
		origin = strings.TrimSpace(origin)
		if origin == "" {
			continue
		}
		if origin == "*" {
			return []string{"*"}
		}
		if parsed, err := url.Parse(origin); err == nil && parsed.Host != "" {
			patterns = append(patterns, parsed.Host)
			continue
		}
		patterns = append(patterns, origin)
	}
	return patterns
}

func writeDirect(ctx context.Context, conn *websocket.Conn, payload []byte) error {
	writeCtx, cancel := context.WithTimeout(ctx, WriteTimeout)
	defer cancel()
	return conn.Write(writeCtx, websocket.MessageText, payload)
}

func decodeIdentify(raw json.RawMessage) (identifyPayload, error) {
	var payload identifyPayload
	if len(raw) == 0 {
		return payload, errors.New("empty identify payload")
	}
	if err := json.Unmarshal(raw, &payload); err != nil {
		return payload, err
	}
	if payload.Token == "" {
		return payload, errors.New("token is required")
	}
	return payload, nil
}

func newSessionID() string {
	const alphabet = "0123456789abcdef"
	buf := make([]byte, 16)
	now := time.Now().UnixNano()
	for i := range buf {
		buf[i] = alphabet[(now>>uint(i*4))&0xf]
	}
	return string(buf)
}

// resumeBuffer хранит последние события для RESUME (AGENT.md 8.3).
type resumeBuffer struct {
	mu      sync.Mutex
	owner   uint64
	limit   int
	items   []bufferedEvent
	updated time.Time
}

type bufferedEvent struct {
	seq     int64
	payload []byte
}

func newResumeBuffer(limit int) *resumeBuffer {
	return &resumeBuffer{limit: limit, updated: time.Now()}
}

func (b *resumeBuffer) append(seq int64, payload []byte) {
	if b == nil {
		return
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	b.items = append(b.items, bufferedEvent{seq: seq, payload: payload})
	if len(b.items) > b.limit {
		b.items = b.items[len(b.items)-b.limit:]
	}
	b.updated = time.Now()
}

// since возвращает события после указанной последовательности; ok=false,
// если буфер уже не содержит их (нужен полный READY).
func (b *resumeBuffer) since(seq int64) ([][]byte, bool) {
	if b == nil {
		return nil, false
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	if len(b.items) == 0 {
		return nil, false
	}
	if b.items[0].seq > seq+1 {
		return nil, false
	}
	// Клиент заявляет последовательность впереди сервера: состояние разошлось
	// (например, сервер перезапускался) — нужен полный READY.
	if seq > b.items[len(b.items)-1].seq {
		return nil, false
	}
	result := make([][]byte, 0, len(b.items))
	for _, item := range b.items {
		if item.seq > seq {
			result = append(result, item.payload)
		}
	}
	return result, true
}

func (b *resumeBuffer) lastSeen() time.Time {
	if b == nil {
		return time.Time{}
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.updated
}