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. direct chooses a plain
// TCP connection to addr (--direct) over an HTTP upgrade at url
// (--http-upgrade).
type dialSettings struct {
	direct   bool
	url      string
	addr     string
	origin   string
	messages []string
	ping     bool
	listen   bool
	timeout  time.Duration
}

// connection is what dial talks to: a *websocket.Conn with --http-upgrade,
// or a *directConn with --direct.
type connection interface {
	Read(ctx context.Context) (websocket.MessageType, []byte, error)
	Write(
		ctx context.Context,
		messageType websocket.MessageType,
		payload []byte,
	) error
	Close(code websocket.StatusCode, reason string) error
	CloseNow() error
}

// commandDial returns the "dial" command, which connects to the WebSocket
// hub of job-search serve (--http-upgrade) or ws-test serve (--direct),
// sends messages and prints the replies.
func commandDial(opts *cli.Options) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "dial",
		Short: "Connect to a WebSocket hub, send messages and print the replies",
		Long: `Connect to a server's WebSocket hub and print every message the server sends,
starting with its welcome. One of these is required:

  --http-upgrade  upgrade an HTTP connection to a WebSocket on the /ws path,
                  as job-search serve expects
  --direct        connect straight to ws-test serve over plain TCP, with one
                  JSON message per line in each direction

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.

With --http-upgrade, 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.

With --direct, --addr defaults to the configured server.host and server.port.
There is no HTTP request, so --url and --origin don't apply.`,
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			slog.Debug("command persistent pre-run", "command", "dial")
		},
		PersistentPostRun: func(cmd *cobra.Command, args []string) {
			slog.Debug("command persistent post-run", "command", "dial")
		},
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			slog.Debug("command run", "command", "dial")

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

			// Only one of url and addr can be set, for the chosen mode.
			if settings.url == "" && settings.addr == "" {
				conf, err := opts.RunSetup(cmd, true, true)
				if err != nil {
					return fmt.Errorf("load configuration: %w", err)
				}
				if settings.direct {
					settings.addr = defaultAddr(conf.Server)
				} else {
					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.Bool(
		"http-upgrade",
		false,
		"upgrade an HTTP connection to a WebSocket on /ws (job-search serve)",
	)
	flags.Bool(
		"direct",
		false,
		"connect straight over plain TCP, one JSON message per line "+
			"(ws-test serve)",
	)
	flags.String(
		"url",
		"",
		"with --http-upgrade, the WebSocket URL "+
			"(default from config, ws://127.0.0.1:9393/ws)",
	)
	flags.String(
		"addr",
		"",
		"with --direct, the address to connect to "+
			"(default from config, 127.0.0.1:9393)",
	)
	flags.String(
		"origin",
		"",
		"with --http-upgrade, the 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 to connect and for each reply; 0s for no limit",
	)
	cmd.MarkFlagsOneRequired("http-upgrade", "direct")
	cmd.MarkFlagsMutuallyExclusive("http-upgrade", "direct")
	cmd.MarkFlagsMutuallyExclusive("direct", "url")
	cmd.MarkFlagsMutuallyExclusive("direct", "origin")
	cmd.MarkFlagsMutuallyExclusive("http-upgrade", "addr")

	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()
	direct, directErr := flags.GetBool("direct")
	url, urlErr := flags.GetString("url")
	addr, addrErr := flags.GetString("addr")
	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(
		directErr,
		urlErr,
		addrErr,
		originErr,
		messagesErr,
		pingErr,
		listenErr,
		timeoutErr,
	)
	if err != nil {
		return dialSettings{}, err
	}

	return dialSettings{
		direct:   direct,
		url:      url,
		addr:     addr,
		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, at defaultAddr.
func defaultURL(conf config.JSCServer) string {
	scheme := "ws"
	if conf.TLS.Enabled {
		scheme = "wss"
	}

	return scheme + "://" + defaultAddr(conf) + "/ws"
}

// defaultAddr is the host:port of the server conf configures. An empty host
// (every interface) is dialled as 127.0.0.1.
func defaultAddr(conf config.JSCServer) string {
	host := conf.Host
	if host == "" {
		host = "127.0.0.1"
	}

	return net.JoinHostPort(host, strconv.Itoa(conf.Port))
}

// dial connects as settings choose (see connect) 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 {
	connectCtx, cancel := util.WithTimeout(ctx, settings.timeout)
	defer cancel()

	conn, err := connect(connectCtx, out, settings)
	if err != nil {
		return err
	}
	defer conn.CloseNow() //nolint:errcheck // finished with; it may already be closed

	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, "")
}

// connect opens a WebSocket at settings.url, or with settings.direct a
// plain TCP connection to settings.addr, giving up when ctx is done. It
// prints how it connected to out.
func connect(
	ctx context.Context,
	out io.Writer,
	settings dialSettings,
) (connection, error) {
	if settings.direct {
		conn, err := dialDirect(ctx, settings.addr)
		if err != nil {
			return nil, fmt.Errorf("connecting to %s: %w", settings.addr, err)
		}
		fmt.Fprintf(out, "connected directly to %s\n", settings.addr) //nolint:errcheck // printing only

		return conn, nil
	}

	header := http.Header{}
	if settings.origin != "" {
		header.Set("Origin", settings.origin)
	}

	conn, res, err := websocket.Dial(
		ctx,
		settings.url,
		&websocket.DialOptions{
			HTTPHeader: header,
		},
	)
	if err != nil {
		if res != nil {
			return nil, fmt.Errorf("upgrade refused: %s", res.Status)
		}
		return nil, fmt.Errorf("dialling %s: %w", settings.url, err)
	}
	fmt.Fprintf(out, "connected (%s)\n", res.Status) //nolint:errcheck // printing only

	return conn, nil
}

// 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 connection,
	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 connection,
	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. Its error, if any, is sent
	// before lines is closed.
	lines := make(chan string)
	inputErr := make(chan error, 1)
	go func() {
		defer close(lines)
		scanner := bufio.NewScanner(input)
		for scanner.Scan() {
			lines <- scanner.Text()
		}
		inputErr <- scanner.Err()
	}()

	for {
		select {
		case <-ctx.Done():
			return <-listening
		case err := <-listening:
			return err
		case line, ok := <-lines:
			if !ok {
				err := <-inputErr
				if err != nil {
					return fmt.Errorf("reading input: %w", err)
				}

				// 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 connection,
	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 connection,
	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 connection) 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: dial closes the connection.
			return nil //nolint:nilerr // a normal end
		case errors.Is(err, io.EOF):
			// A direct connection has no close status.
			fmt.Fprintln(out, "closed by the server") //nolint:errcheck // printing only
			return nil
		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)
		}
	}
}