go_pwner/tube.go

package winpwn

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"sync"
	"syscall"
	"time"
)

// Tube is the core abstraction for talking to a target: a spawned local
// process or a remote TCP connection. It mirrors pwntools' tube class and
// is the type every transport (Spawn, Remote, ...) returns.
//
// Every blocking call returns an error instead of killing the process --
// callers decide what a timeout/EOF/closed-connection means for their
// script, the same way pwntools raises (and lets you catch) EOFError /
// PwnlibException instead of the library calling sys.exit().
type Tube struct {
	cmd    *exec.Cmd
	conn   net.Conn
	stdin  io.WriteCloser
	stdout io.ReadCloser
	reader *bufio.Reader

	// timeout overrides Context.Timeout for this tube specifically; zero
	// means "fall back to Context.Timeout" (which is itself zero/forever
	// by default). Set via SetTimeout.
	timeout time.Duration

	closeOnce sync.Once
	closeErr  error
}

// newTube wraps a writer/reader pair into a Tube with a buffered reader.
func newTube(cmd *exec.Cmd, conn net.Conn, stdin io.WriteCloser, stdout io.ReadCloser) *Tube {
	return &Tube{
		cmd:    cmd,
		conn:   conn,
		stdin:  stdin,
		stdout: stdout,
		reader: bufio.NewReader(stdout),
	}
}

// SetTimeout overrides Context.Timeout for this tube's Recv*/Send* calls.
// Zero means block forever (the default), matching pwntools' per-tube
// timeout= override of the global context.timeout.
func (t *Tube) SetTimeout(d time.Duration) {
	t.timeout = d
}

func (t *Tube) effectiveTimeout() time.Duration {
	if t.timeout > 0 {
		return t.timeout
	}
	return Context.Timeout
}

// PID returns the OS process ID for a locally spawned process, or 0 for
// remote tubes. Use it to feed winpwn.NewProcessSymbols or
// winpwn.ListProcessHeaps without having to track the PID separately.
func (t *Tube) PID() uint32 {
	if t.cmd != nil && t.cmd.Process != nil {
		return uint32(t.cmd.Process.Pid)
	}
	return 0
}

// Close tears down the underlying process/connection. Safe to call more
// than once (idempotent) -- Interactive relies on this to make a pending
// stdin write fail fast instead of dangling.
func (t *Tube) Close() error {
	t.closeOnce.Do(func() {
		if t.stdin != nil {
			_ = t.stdin.Close()
		}
		if t.stdout != nil {
			_ = t.stdout.Close()
		}
		if t.conn != nil {
			t.closeErr = t.conn.Close()
			return
		}
		if t.cmd != nil && t.cmd.Process != nil {
			t.closeErr = t.cmd.Process.Kill()
		}
	})
	return t.closeErr
}

// withTimeout runs fn on its own goroutine and races it against this tube's
// effective timeout. If fn doesn't return in time, withTimeout returns a
// timeout error immediately -- but fn's goroutine is *not* killed (the
// underlying pipe/socket reader has no native per-call deadline), so it
// keeps running in the background until the blocking I/O it's stuck in
// eventually completes or errors. That's the standard, and only portable,
// way to bolt a deadline onto an arbitrary io.Reader/Writer in Go.
func withTimeout[T any](t *Tube, fn func() (T, error)) (T, error) {
	timeout := t.effectiveTimeout()
	if timeout <= 0 {
		return fn()
	}
	type result struct {
		v   T
		err error
	}
	ch := make(chan result, 1)
	go func() {
		v, err := fn()
		ch <- result{v, err}
	}()
	select {
	case r := <-ch:
		return r.v, r.err
	case <-time.After(timeout):
		var zero T
		return zero, fmt.Errorf("winpwn: operation timed out after %s", timeout)
	}
}

// Recv reads up to n bytes from the tube, blocking until n bytes have
// arrived. On a clean EOF after at least one byte it returns the partial
// read with a nil error (mirroring the old behavior); on EOF with nothing
// read yet, it returns io.EOF.
func (t *Tube) Recv(n int) ([]byte, error) {
	return withTimeout(t, func() ([]byte, error) {
		buf := make([]byte, n)
		read, err := io.ReadFull(t.reader, buf)
		if err != nil {
			if err == io.ErrUnexpectedEOF || err == io.EOF {
				if read > 0 {
					return buf[:read], nil
				}
				return nil, io.EOF
			}
			return buf[:read], err
		}
		return buf, nil
	})
}

// RecvUntil reads from the tube until delim is seen (inclusive of delim).
func (t *Tube) RecvUntil(delim []byte) ([]byte, error) {
	return withTimeout(t, func() ([]byte, error) {
		var out []byte
		for {
			b, err := t.reader.ReadByte()
			if err != nil {
				return out, err
			}
			out = append(out, b)
			if bytes.HasSuffix(out, delim) {
				return out, nil
			}
		}
	})
}

// RecvLine reads a single line, including the trailing newline
// (Context.Newline, "\n" by default).
func (t *Tube) RecvLine() ([]byte, error) {
	return t.RecvUntil(Context.Newline)
}

// RecvPred reads one byte at a time until pred(accumulated) reports true,
// the analogue of pwntools' recvpred.
func (t *Tube) RecvPred(pred func([]byte) bool) ([]byte, error) {
	return withTimeout(t, func() ([]byte, error) {
		var out []byte
		for {
			b, err := t.reader.ReadByte()
			if err != nil {
				return out, err
			}
			out = append(out, b)
			if pred(out) {
				return out, nil
			}
		}
	})
}

// RecvRegex reads one byte at a time until the accumulated buffer matches
// re, the analogue of pwntools' recvregex.
func (t *Tube) RecvRegex(re *regexp.Regexp) ([]byte, error) {
	return t.RecvPred(func(buf []byte) bool {
		return re.Match(buf)
	})
}

// Send writes raw bytes to the tube.
func (t *Tube) Send(data []byte) error {
	_, err := withTimeout(t, func() (int, error) {
		return t.stdin.Write(data)
	})
	return err
}

// SendLine writes data followed by Context.Newline.
func (t *Tube) SendLine(data []byte) error {
	return t.Send(append(append([]byte{}, data...), Context.Newline...))
}

// SendAfter waits for delim, then sends data (no trailing newline).
func (t *Tube) SendAfter(delim []byte, data []byte) error {
	if _, err := t.RecvUntil(delim); err != nil {
		return err
	}
	return t.Send(data)
}

// SendLineAfter waits for delim, then sends data followed by a newline.
func (t *Tube) SendLineAfter(delim []byte, data []byte) error {
	if _, err := t.RecvUntil(delim); err != nil {
		return err
	}
	return t.SendLine(data)
}

// Interactive hands the tube's stdin/stdout over to the user's terminal,
// the Go analogue of pwntools' tube.interactive(). Ctrl+C cleanly tears
// down the local process or remote connection.
//
// The stdin-forwarding goroutine below has no portable way to be cancelled
// in Go (os.Stdin.Read blocks with no deadline support), so it keeps
// running until the next keystroke/EOF even after Interactive returns;
// the deferred Close() at least makes its next Write fail fast instead of
// leaving the target side dangling. This is a known, deliberate limitation,
// not an oversight -- don't call Interactive() in a tight loop expecting
// the goroutine to be gone before the next iteration.
func (t *Tube) Interactive() {
	defer t.Close()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
	defer signal.Stop(sigChan)

	stdoutDone := make(chan struct{})
	go func() {
		_, _ = io.Copy(os.Stdout, t.reader)
		close(stdoutDone)
	}()

	go func() {
		_, _ = io.Copy(t.stdin, os.Stdin)
	}()

	Info("Switching to interactive mode")

	if t.cmd != nil {
		// Local process branch.
		waitCh := make(chan error, 1)
		go func() {
			waitCh <- t.cmd.Wait()
		}()

		select {
		case err := <-waitCh:
			if err == nil {
				Info("Process exited normally (code 0)")
			} else {
				var exitErr *exec.ExitError
				if errors.As(err, &exitErr) {
					code := exitErr.ExitCode()
					// 0xC0000005 (Access Violation) is the Windows analogue of SIGSEGV.
					Error("Process crashed/terminated with code: 0x%X", uint32(code))
				} else {
					Error("Process execution error: %v", err)
				}
			}
		case <-sigChan:
			Info("Interrupted by user, killing process...")
		}
	} else {
		// Remote connection branch.
		select {
		case <-stdoutDone:
			Info("Connection closed by foreign host")
		case <-sigChan:
			Info("Interrupted by user, closing connection...")
		}
	}
}