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