teraWeb/backend/internal/directstream/shared_socket.go

package directstream

import (
	"context"
	"errors"
	"net"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

const DefaultSharedSocketPath = "/run/tera-api/shared-stream.sock"

// SharedAttachment is one authenticated terminal, not an aggregate per service.
type SharedAttachment interface {
	Receive(context.Context) ([]byte, error)
	Write(context.Context, []byte) error
	SendBreak(context.Context) error
	Close() error
}

type sharedFrame struct {
	Type  string `json:"type"`
	Data  []byte `json:"data,omitempty"`
	Error string `json:"error,omitempty"`
}

// ServeSharedSocket must only be exposed on a service-group-owned Unix socket.
// The gateway has authenticated the user before it sends an attachment request.
func ServeSharedSocket(listener net.Listener, broker *Broker) error {
	server := &http.Server{ReadHeaderTimeout: 5 * time.Second, Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upgrader := websocket.Upgrader{}
		connection, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer connection.Close()
		connection.SetReadLimit(96 * 1024)
		_ = connection.SetReadDeadline(time.Now().Add(10 * time.Second))
		var request AttachRequest
		if connection.ReadJSON(&request) != nil {
			return
		}
		request.SharedWrite = request.Access == AccessWriter
		attachment, err := broker.Attach(r.Context(), request)
		if err != nil {
			_ = connection.WriteJSON(sharedFrame{Type: "attached", Error: err.Error()})
			return
		}
		defer attachment.Close()
		_ = connection.SetWriteDeadline(time.Now().Add(10 * time.Second))
		if connection.WriteJSON(sharedFrame{Type: "attached"}) != nil {
			return
		}
		_ = connection.SetReadDeadline(time.Time{})
		ctx, cancel := context.WithCancel(r.Context())
		defer cancel()
		var writeMu sync.Mutex
		write := func(frame sharedFrame) error {
			writeMu.Lock()
			defer writeMu.Unlock()
			_ = connection.SetWriteDeadline(time.Now().Add(10 * time.Second))
			return connection.WriteJSON(frame)
		}
		go func() {
			defer connection.Close()
			for {
				payload, err := attachment.Receive(ctx)
				if err != nil || write(sharedFrame{Type: "data", Data: payload}) != nil {
					return
				}
			}
		}()
		for {
			var frame sharedFrame
			if connection.ReadJSON(&frame) != nil {
				return
			}
			switch frame.Type {
			case "write":
				err = attachment.Write(ctx, frame.Data)
			case "break":
				err = attachment.SendBreak(ctx)
			default:
				return
			}
			ack := sharedFrame{Type: "ack"}
			if err != nil {
				ack.Error = err.Error()
			}
			if write(ack) != nil {
				return
			}
		}
	})}
	return server.Serve(listener)
}

type SharedSocketClient struct{ Path string }

func (c SharedSocketClient) Attach(ctx context.Context, request AttachRequest) (SharedAttachment, error) {
	path := c.Path
	if path == "" {
		path = DefaultSharedSocketPath
	}
	dialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second, NetDialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
		return (&net.Dialer{Timeout: 10 * time.Second}).DialContext(ctx, "unix", path)
	}}
	connection, _, err := dialer.DialContext(ctx, "ws://localhost/stream", nil)
	if err != nil {
		return nil, err
	}
	stop := context.AfterFunc(ctx, func() { _ = connection.Close() })
	defer stop()
	connection.SetReadLimit(96 * 1024)
	_ = connection.SetWriteDeadline(time.Now().Add(10 * time.Second))
	_ = connection.SetReadDeadline(time.Now().Add(10 * time.Second))
	if err = connection.WriteJSON(request); err != nil {
		connection.Close()
		return nil, err
	}
	var response sharedFrame
	if err = connection.ReadJSON(&response); err != nil {
		connection.Close()
		return nil, err
	}
	if response.Error != "" || response.Type != "attached" {
		connection.Close()
		if response.Error == ErrAttachmentLimit.Error() {
			return nil, ErrAttachmentLimit
		}
		return nil, errors.New("shared stream attachment rejected: " + response.Error)
	}
	_ = connection.SetReadDeadline(time.Time{})
	attachment := &remoteAttachment{connection: connection, queue: newByteQueue(defaultAttachmentQueueBytes), acknowledgements: make(chan error, 1), done: make(chan struct{})}
	go attachment.read()
	return attachment, nil
}

type remoteAttachment struct {
	connection       *websocket.Conn
	queue            *byteQueue
	requestMu        sync.Mutex
	closeOnce        sync.Once
	acknowledgements chan error
	done             chan struct{}
}

func (a *remoteAttachment) read() {
	defer a.Close()
	for {
		var frame sharedFrame
		if a.connection.ReadJSON(&frame) != nil {
			return
		}
		switch frame.Type {
		case "data":
			if !a.queue.push(frame.Data) {
				return
			}
		case "ack":
			var err error
			if frame.Error != "" {
				err = errors.New(frame.Error)
			}
			select {
			case a.acknowledgements <- err:
			default:
				return
			}
		default:
			return
		}
	}
}

func (a *remoteAttachment) Receive(ctx context.Context) ([]byte, error) { return a.queue.pop(ctx) }
func (a *remoteAttachment) Write(ctx context.Context, payload []byte) error {
	if len(payload) == 0 || len(payload) > MaxDataFrameBytes {
		return ErrInvalidDataFrame
	}
	return a.command(ctx, sharedFrame{Type: "write", Data: payload})
}
func (a *remoteAttachment) SendBreak(ctx context.Context) error {
	return a.command(ctx, sharedFrame{Type: "break"})
}
func (a *remoteAttachment) command(ctx context.Context, frame sharedFrame) error {
	a.requestMu.Lock()
	defer a.requestMu.Unlock()
	if err := ctx.Err(); err != nil {
		return err
	}
	deadline := time.Now().Add(10 * time.Second)
	if value, ok := ctx.Deadline(); ok && value.Before(deadline) {
		deadline = value
	}
	_ = a.connection.SetWriteDeadline(deadline)
	if err := a.connection.WriteJSON(frame); err != nil {
		a.Close()
		return err
	}
	timer := time.NewTimer(time.Until(deadline))
	defer timer.Stop()
	select {
	case err := <-a.acknowledgements:
		return err
	case <-a.done:
		return ErrChannelClosed
	case <-ctx.Done():
		a.Close()
		return ctx.Err()
	case <-timer.C:
		a.Close()
		return context.DeadlineExceeded
	}
}
func (a *remoteAttachment) Close() error {
	a.closeOnce.Do(func() {
		close(a.done)
		_ = a.connection.Close()
		a.queue.close(ErrChannelClosed)
	})
	return nil
}