job-search/internal/cli/ws/dial.go

package ws

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/coder/websocket"
	"github.com/spf13/cobra"

	"git.avsdev.uk/Go/job-search/internal/cli"
	"git.avsdev.uk/Go/job-search/internal/config"
	"git.avsdev.uk/Go/job-search/internal/util"
	hub "git.avsdev.uk/Go/job-search/internal/websocket"
)

// pingID is the ID of the ping --ping sends, which the pong it waits for
// carries back.
const pingID = "ping"

// dialSettings are the dial command's flag values.
type dialSettings struct {
	url      string
	origin   string
	messages []string
	ping     bool
	listen   bool
	timeout  time.Duration
}

// commandDial returns the "ws dial" command, which opens a WebSocket on the
// HTTP server's /ws path, sends messages and prints the replies.
func commandDial(opts *cli.Options) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "dial",
		Short: "Open a WebSocket on /ws, send messages and print the replies",
		Long: `Upgrade an HTTP connection to a WebSocket on the server's /ws path and print
every message the server sends, starting with its welcome.

With --ping, it sends a ping and waits for the pong, failing if none arrives
within --timeout.

With --message, each one is sent in turn and the reply to it printed.

With --ping or --message, dial then closes the connection, or with --listen
prints every message the server pushes until the connection closes or Ctrl+C
is pressed.

With neither, each line read from stdin is sent as a message (blank lines are
skipped), until Ctrl+C is pressed or the server closes the connection. The end
of stdin stops the sending, but the connection stays open.

The URL defaults to the configured server.host and server.port (wss:// if
server.tls.enabled is set). Set --origin to test the Origin check: an upgrade
from an origin that isn't the server's own or one of
server.csrf.cross_origin_protection.trusted_origins is refused with 403.`,
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			slog.Debug("command persistent pre-run", "command", "ws.dial")
		},
		PersistentPostRun: func(cmd *cobra.Command, args []string) {
			slog.Debug("command persistent post-run", "command", "ws.dial")
		},
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			slog.Debug("command run", "command", "ws.dial")

			settings, err := dialFlags(cmd)
			if err != nil {
				return err
			}

			if settings.url == "" {
				conf, err := opts.RunSetup(cmd, true, true)
				if err != nil {
					return fmt.Errorf("load configuration: %w", err)
				}
				settings.url = defaultURL(conf.Server)
			}

			// ctx is cancelled on the first SIGINT (Ctrl+C) or SIGTERM.
			ctx, stop := signal.NotifyContext(
				context.Background(), os.Interrupt, syscall.SIGTERM,
			)
			defer stop()

			return dial(ctx, cmd.InOrStdin(), cmd.OutOrStdout(), settings)
		},
	}

	flags := cmd.Flags()
	flags.String(
		"url",
		"",
		"WebSocket URL (default from config, ws://127.0.0.1:9393/ws)",
	)
	flags.String(
		"origin",
		"",
		"Origin header to send, as a browser would (default none)",
	)
	flags.StringArray(
		"message",
		nil,
		"message to send, as JSON; repeat to send several (default: stdin)",
	)
	flags.Bool(
		"ping",
		false,
		"send a ping and wait for the pong (before any --message)",
	)
	flags.Bool(
		"listen",
		false,
		"with --ping or --message, keep printing pushed messages until the "+
			"connection closes or Ctrl+C",
	)
	flags.Duration(
		"timeout",
		10*time.Second,
		"how long to wait for the upgrade and for each reply; 0s for no limit",
	)

	return cmd
}

// dialFlags reads the dial command's flags. The flags are all defined by
// commandDial, so getting them can't fail.
func dialFlags(cmd *cobra.Command) (dialSettings, error) {
	flags := cmd.Flags()
	url, urlErr := flags.GetString("url")
	origin, originErr := flags.GetString("origin")
	messages, messagesErr := flags.GetStringArray("message")
	ping, pingErr := flags.GetBool("ping")
	listen, listenErr := flags.GetBool("listen")
	timeout, timeoutErr := flags.GetDuration("timeout")

	err := errors.Join(
		urlErr,
		originErr,
		messagesErr,
		pingErr,
		listenErr,
		timeoutErr,
	)
	if err != nil {
		return dialSettings{}, err
	}

	return dialSettings{
		url:      url,
		origin:   origin,
		messages: messages,
		ping:     ping,
		listen:   listen,
		timeout:  timeout,
	}, nil
}

// defaultURL is the /ws URL of the server conf configures: ws://, or wss://
// with TLS, on its host and port. An empty host (every interface) is dialled
// as 127.0.0.1.
func defaultURL(conf config.JSCServer) string {
	scheme := "ws"
	if conf.TLS.Enabled {
		scheme = "wss"
	}

	host := conf.Host
	if host == "" {
		host = "127.0.0.1"
	}

	address := net.JoinHostPort(host, strconv.Itoa(conf.Port))

	return scheme + "://" + address + "/ws"
}

// dial opens a WebSocket at settings.url and prints the messages received
// to out, each after "<". It sends each of settings.messages, printed after
// ">", or without them each line of input (see interact), which isn't
// printed again.
func dial(
	ctx context.Context,
	input io.Reader,
	out io.Writer,
	settings dialSettings,
) error {
	header := http.Header{}
	if settings.origin != "" {
		header.Set("Origin", settings.origin)
	}

	dialCtx, cancel := util.WithTimeout(ctx, settings.timeout)
	defer cancel()

	conn, res, err := websocket.Dial(
		dialCtx,
		settings.url,
		&websocket.DialOptions{
			HTTPHeader: header,
		},
	)
	if err != nil {
		if res != nil {
			return fmt.Errorf("upgrade refused: %s", res.Status)
		}
		return fmt.Errorf("dialling %s: %w", settings.url, err)
	}
	defer conn.CloseNow() //nolint:errcheck // finished with; it may already be closed

	fmt.Fprintf(out, "connected (%s)\n", res.Status) //nolint:errcheck // printing only

	if !settings.ping && len(settings.messages) == 0 {
		return interact(ctx, input, out, conn, settings.timeout)
	}

	// The server sends a welcome first.
	err = receive(ctx, out, conn, settings.timeout)
	if err != nil {
		return err
	}

	if settings.ping {
		err = pingPong(ctx, out, conn, settings.timeout)
		if err != nil {
			return err
		}
	}

	for _, message := range settings.messages {
		fmt.Fprintf(out, "> %s\n", message) //nolint:errcheck // printing only

		err = send(ctx, conn, message, settings.timeout)
		if err != nil {
			return err
		}

		err = receive(ctx, out, conn, settings.timeout)
		if err != nil {
			return err
		}
	}

	if settings.listen {
		return listen(ctx, out, conn)
	}

	return conn.Close(websocket.StatusNormalClosure, "")
}

// pingPong sends a ping on conn and waits at most timeout for its pong,
// printing each message received meanwhile to out.
func pingPong(
	ctx context.Context,
	out io.Writer,
	conn *websocket.Conn,
	timeout time.Duration,
) error {
	ping, err := json.Marshal(hub.Envelope{
		Type: hub.TypePing,
		ID:   pingID,
	})
	if err != nil {
		return fmt.Errorf("encoding the ping: %w", err)
	}

	fmt.Fprintf(out, "> %s\n", ping) //nolint:errcheck // printing only
	err = send(ctx, conn, string(ping), timeout)
	if err != nil {
		return err
	}

	waitCtx, cancel := util.WithTimeout(ctx, timeout)
	defer cancel()

	for {
		_, data, readErr := conn.Read(waitCtx)
		if readErr != nil {
			if ctx.Err() == nil && waitCtx.Err() != nil {
				return fmt.Errorf("no pong within %s", timeout)
			}
			return fmt.Errorf("waiting for the pong: %w", readErr)
		}
		fmt.Fprintf(out, "< %s\n", data) //nolint:errcheck // printing only

		// Anything that isn't the pong, such as a push, is only printed.
		var reply hub.Envelope
		if json.Unmarshal(data, &reply) == nil &&
			reply.Type == hub.TypePong && reply.ID == pingID {
			return nil
		}
	}
}

// interact prints each message received on conn to out, and sends each line
// of input as a message, skipping blank lines, until ctx is cancelled
// (Ctrl+C) or the server closes the connection. The end of input stops the
// sending; the connection stays open.
func interact(
	ctx context.Context,
	input io.Reader,
	out io.Writer,
	conn *websocket.Conn,
	timeout time.Duration,
) error {
	listening := make(chan error, 1)
	go func() {
		listening <- listen(ctx, out, conn)
	}()

	// Reading input can't be interrupted, so it has a goroutine of its own,
	// left blocked if ctx is cancelled first.
	lines := make(chan string)
	go func() {
		defer close(lines)
		scanner := bufio.NewScanner(input)
		for scanner.Scan() {
			lines <- scanner.Text()
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return <-listening
		case err := <-listening:
			return err
		case line, ok := <-lines:
			if !ok {
				// The end of input: stop sending, keep listening.
				lines = nil
				continue
			}
			if strings.TrimSpace(line) == "" {
				continue
			}

			err := send(ctx, conn, line, timeout)
			if err != nil {
				return err
			}
		}
	}
}

// send writes message to conn as a text frame, waiting at most timeout.
func send(
	ctx context.Context,
	conn *websocket.Conn,
	message string,
	timeout time.Duration,
) error {
	writeCtx, cancel := util.WithTimeout(ctx, timeout)
	defer cancel()

	err := conn.Write(writeCtx, websocket.MessageText, []byte(message))
	if err != nil {
		return fmt.Errorf("sending: %w", err)
	}

	return nil
}

// receive reads one message from conn, waiting at most timeout, and prints
// it to out.
func receive(
	ctx context.Context,
	out io.Writer,
	conn *websocket.Conn,
	timeout time.Duration,
) error {
	readCtx, cancel := util.WithTimeout(ctx, timeout)
	defer cancel()

	_, data, err := conn.Read(readCtx)
	if err != nil {
		return fmt.Errorf("waiting for a reply: %w", err)
	}
	fmt.Fprintf(out, "< %s\n", data) //nolint:errcheck // printing only

	return nil
}

// listen prints each message received on conn to out until ctx is cancelled
// (Ctrl+C) or the server closes the connection, either of which is a normal
// end.
func listen(ctx context.Context, out io.Writer, conn *websocket.Conn) error {
	for {
		_, data, err := conn.Read(ctx)
		status := websocket.CloseStatus(err)
		switch {
		case err == nil:
			fmt.Fprintf(out, "< %s\n", data) //nolint:errcheck // printing only
		case ctx.Err() != nil:
			// Ctrl+C: Read has closed the connection already.
			return nil //nolint:nilerr // a normal end
		case status != -1:
			fmt.Fprintf(out, "closed by the server (%d)\n", status) //nolint:errcheck // printing only
			return nil
		default:
			return fmt.Errorf("reading: %w", err)
		}
	}
}