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