| 1 | package persistentshell |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "crypto/rand" |
| 7 | _ "embed" |
| 8 | "encoding/base64" |
| 9 | "encoding/binary" |
| 10 | "encoding/json" |
| 11 | "errors" |
| 12 | "fmt" |
| 13 | "io" |
| 14 | "net" |
| 15 | "os" |
| 16 | "os/exec" |
| 17 | "strings" |
| 18 | "sync" |
| 19 | "time" |
| 20 | "unicode/utf16" |
| 21 | |
| 22 | "reasonix/internal/proc" |
| 23 | "reasonix/internal/tool" |
| 24 | ) |
| 25 | |
| 26 | //go:embed powershell.ps1 |
| 27 | var powershellBootstrap string |
| 28 | |
| 29 | const maxControlFrame = 4 << 20 |
| 30 | |
| 31 | type shellFrame struct { |
| 32 | Version int `json:"version"` |
| 33 | Kind string `json:"kind"` |
| 34 | ID string `json:"id"` |
| 35 | Command string `json:"command,omitempty"` |
| 36 | Start string `json:"start,omitempty"` |
| 37 | End string `json:"end,omitempty"` |
| 38 | ExitCode *int `json:"exitCode,omitempty"` |
| 39 | } |
| 40 | |
| 41 | func readShellFrame(r io.Reader) (shellFrame, error) { |
| 42 | var size uint32 |
| 43 | if err := binary.Read(r, binary.LittleEndian, &size); err != nil { |
| 44 | return shellFrame{}, err |
| 45 | } |
| 46 | if size == 0 || size > maxControlFrame { |
| 47 | return shellFrame{}, errors.New("invalid shell control frame size") |
| 48 | } |
| 49 | data := make([]byte, size) |
| 50 | if _, err := io.ReadFull(r, data); err != nil { |
| 51 | return shellFrame{}, err |
| 52 | } |
| 53 | var frame shellFrame |
| 54 | if err := json.Unmarshal(data, &frame); err != nil { |
| 55 | return frame, err |
| 56 | } |
| 57 | if frame.Version != 1 { |
| 58 | return frame, errors.New("unsupported shell control protocol") |
| 59 | } |
| 60 | return frame, nil |
| 61 | } |
| 62 | |
| 63 | func writeShellFrame(w io.Writer, frame shellFrame) error { |
| 64 | data, err := json.Marshal(frame) |
| 65 | if err != nil { |
| 66 | return err |
| 67 | } |
| 68 | if len(data) > maxControlFrame { |
| 69 | return errors.New("shell command exceeds control frame limit") |
| 70 | } |
| 71 | if err := binary.Write(w, binary.LittleEndian, uint32(len(data))); err != nil { |
| 72 | return err |
| 73 | } |
| 74 | _, err = io.Copy(w, bytes.NewReader(data)) |
| 75 | return err |
| 76 | } |
| 77 | |
| 78 | // encodedPowerShell uses the encoding required by both Windows PowerShell and pwsh. |
| 79 | func encodedPowerShell(script string) string { |
| 80 | units := utf16.Encode([]rune(script)) |
| 81 | data := make([]byte, len(units)*2) |
| 82 | for i, unit := range units { |
| 83 | binary.LittleEndian.PutUint16(data[i*2:], unit) |
| 84 | } |
| 85 | return base64.StdEncoding.EncodeToString(data) |
| 86 | } |
| 87 | |
| 88 | type powershellProcess struct { |
| 89 | control net.Conn |
| 90 | output *os.File |
| 91 | cmd *exec.Cmd |
| 92 | job uintptr |
| 93 | once sync.Once |
| 94 | stop chan struct{} |
| 95 | } |
| 96 | |
| 97 | func (p *powershellProcess) Read(b []byte) (int, error) { return p.output.Read(b) } |
| 98 | func (p *powershellProcess) Write([]byte) (int, error) { |
| 99 | return 0, errors.New("use shell control channel") |
| 100 | } |
| 101 | func (p *powershellProcess) Close() error { |
| 102 | p.once.Do(func() { |
| 103 | close(p.stop) |
| 104 | if p.control != nil { |
| 105 | _ = p.control.Close() |
| 106 | } |
| 107 | proc.KillTracked(p.cmd, p.job) |
| 108 | _ = p.output.Close() |
| 109 | _ = p.cmd.Wait() |
| 110 | }) |
| 111 | return nil |
| 112 | } |
| 113 | |
| 114 | func startPowerShell(req Request, fp string) (*session, error) { |
| 115 | name := "rx-" + rand.Text() |
| 116 | listener, err := listenShellControl(name) |
| 117 | if err != nil { |
| 118 | return nil, err |
| 119 | } |
| 120 | defer listener.Close() |
| 121 | reader, writer, err := os.Pipe() |
| 122 | if err != nil { |
| 123 | return nil, err |
| 124 | } |
| 125 | cmd := proc.Command(req.Argv[0], req.Argv[1:]...) |
| 126 | cmd.Dir, cmd.Env = req.Dir, append(append([]string(nil), req.Env...), "REASONIX_PWSH_PIPE="+name) |
| 127 | cmd.Stdout, cmd.Stderr = writer, writer |
| 128 | job, err := startShellTracked(cmd) |
| 129 | _ = writer.Close() |
| 130 | if err != nil { |
| 131 | _ = reader.Close() |
| 132 | return nil, err |
| 133 | } |
| 134 | p := &powershellProcess{output: reader, cmd: cmd, job: job, stop: make(chan struct{})} |
| 135 | // Windows PowerShell 5.1 can take more than 10 seconds to initialize on a |
| 136 | // cold or contended host. Bound the complete connect-and-ready handshake by |
| 137 | // one deadline so a late connection cannot silently start a second budget. |
| 138 | deadline := time.Now().Add(powerShellStartupTimeout) |
| 139 | type accepted struct { |
| 140 | conn net.Conn |
| 141 | err error |
| 142 | } |
| 143 | ready := make(chan accepted, 1) |
| 144 | go func() { conn, err := listener.Accept(); ready <- accepted{conn, err} }() |
| 145 | ctx, cancel := context.WithDeadline(context.Background(), deadline) |
| 146 | defer cancel() |
| 147 | select { |
| 148 | case result := <-ready: |
| 149 | if result.err != nil { |
| 150 | _ = p.Close() |
| 151 | return nil, result.err |
| 152 | } |
| 153 | p.control = result.conn |
| 154 | case <-ctx.Done(): |
| 155 | _ = listener.Close() |
| 156 | result := <-ready |
| 157 | if result.conn != nil { |
| 158 | _ = result.conn.Close() |
| 159 | } |
| 160 | _ = p.Close() |
| 161 | return nil, ctx.Err() |
| 162 | } |
| 163 | _ = p.control.SetDeadline(deadline) |
| 164 | frame, err := readShellFrame(p.control) |
| 165 | if err != nil || frame.Kind != "ready" { |
| 166 | _ = p.Close() |
| 167 | if err == nil { |
| 168 | err = errors.New("missing ready frame") |
| 169 | } |
| 170 | return nil, fmt.Errorf("PowerShell handshake failed: %w", err) |
| 171 | } |
| 172 | _ = p.control.SetDeadline(time.Time{}) |
| 173 | s := &session{conn: p, fp: fp, powershell: p} |
| 174 | s.startReader() |
| 175 | return s, nil |
| 176 | } |
| 177 | |
| 178 | func (s *session) runPowerShell(ctx context.Context, req Request) Result { |
| 179 | runCtx := ctx |
| 180 | if req.Timeout > 0 { |
| 181 | var cancel context.CancelFunc |
| 182 | runCtx, cancel = context.WithTimeout(ctx, req.Timeout) |
| 183 | defer cancel() |
| 184 | } |
| 185 | id := newMarkerID() + newMarkerID() |
| 186 | start, end := "RX_PWSH_START_"+id, "RX_PWSH_END_"+id+":" |
| 187 | capt := newCapture(start, end, req.Progress) |
| 188 | result := Result{Started: true} |
| 189 | type response struct { |
| 190 | frame shellFrame |
| 191 | err error |
| 192 | } |
| 193 | control := make(chan response, 1) |
| 194 | go func() { |
| 195 | err := writeShellFrame(s.powershell.control, shellFrame{Version: 1, Kind: "run", ID: id, Command: req.Command, Start: start, End: end}) |
| 196 | if err != nil { |
| 197 | control <- response{err: err} |
| 198 | return |
| 199 | } |
| 200 | for _, kind := range []string{"started", "completed"} { |
| 201 | frame, err := readShellFrame(s.powershell.control) |
| 202 | if err == nil && (frame.Kind != kind || frame.ID != id) { |
| 203 | err = errors.New("unexpected shell command receipt") |
| 204 | } |
| 205 | if err != nil || kind == "completed" { |
| 206 | control <- response{frame, err} |
| 207 | return |
| 208 | } |
| 209 | } |
| 210 | }() |
| 211 | var completion *int |
| 212 | var failure error |
| 213 | for failure == nil && !(capt.done && completion != nil) { |
| 214 | select { |
| 215 | case <-runCtx.Done(): |
| 216 | failure = runCtx.Err() |
| 217 | case receipt := <-control: |
| 218 | failure = receipt.err |
| 219 | completion = receipt.frame.ExitCode |
| 220 | if failure == nil && completion == nil { |
| 221 | failure = errors.New("missing shell exit status") |
| 222 | } |
| 223 | case chunk := <-s.pendingRead: |
| 224 | capt.push(string(chunk.data)) |
| 225 | if chunk.err != nil && !capt.done { |
| 226 | failure = chunk.err |
| 227 | } |
| 228 | } |
| 229 | } |
| 230 | if failure == nil && *completion != capt.exitCode { |
| 231 | failure = errors.New("shell output fence disagrees with completion") |
| 232 | } |
| 233 | if failure == nil { |
| 234 | result.Output, result.ExitCode, result.ExitCodeKnown = capt.body(), *completion, true |
| 235 | result.State = tool.ShellStateCompleted |
| 236 | if *completion != 0 { |
| 237 | result.State, result.FailurePhase, result.Err = tool.ShellStateFailed, tool.ShellPhaseExecution, fmt.Errorf("exit status %d", *completion) |
| 238 | } |
| 239 | return result |
| 240 | } |
| 241 | // A broken pipe may stop halfway through the private output fence. |
| 242 | // Do not release that protocol suffix to progress or model output. |
| 243 | if at := strings.Index(capt.hold, end); at >= 0 { |
| 244 | capt.hold = capt.hold[:at] |
| 245 | } else { |
| 246 | for size := min(len(capt.hold), len(end)-1); size > 0; size-- { |
| 247 | if strings.HasSuffix(capt.hold, end[:size]) { |
| 248 | capt.hold = capt.hold[:len(capt.hold)-size] |
| 249 | break |
| 250 | } |
| 251 | } |
| 252 | } |
| 253 | result.Output, result.Err, result.ShellDied = capt.partial(), failure, true |
| 254 | result.State, result.FailurePhase = tool.ShellStateFailed, tool.ShellPhaseExecution |
| 255 | if errors.Is(failure, context.Canceled) { |
| 256 | result.Canceled = true |
| 257 | result.State, result.FailurePhase = tool.ShellStateCancelled, tool.ShellPhaseCancellation |
| 258 | } |
| 259 | if errors.Is(failure, context.DeadlineExceeded) { |
| 260 | result.TimedOut = true |
| 261 | result.State, result.FailurePhase = tool.ShellStateTimedOut, tool.ShellPhaseTimeout |
| 262 | } |
| 263 | s.markClosed() |
| 264 | return result |
| 265 | } |
| 266 |