Print memory segments in current process with wx or rwx permissions using procstat(1) on FreeBSD

go.mod

module printwx

go 1.16

require github.com/mattn/go-isatty v0.0.13

printwx.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"sort"
	"strconv"
	"strings"

	"github.com/mattn/go-isatty"
)

type ColorOption int
const (
	ColorNo ColorOption = iota
	ColorYes
	ColorTTYOnly
)

func (c *ColorOption) IsBoolFlag() bool {
	return true
}
func (c *ColorOption) Set(s string) error {
	if c == nil {
		return nil
	}
	ls := strings.ToLower(s)

	no := []string{"0", "f", "false", "n", "no", "none", "off"}
	i := sort.SearchStrings(no, ls)
	if i < len(no) && no[i] == ls {
		*c = ColorNo
		return nil
	}

	yes := []string{"1", "always", "on", "t", "true", "y", "yes"}
	i = sort.SearchStrings(yes, ls)
	if i < len(yes) && yes[i] == ls {
		*c = ColorYes
		return nil
	}

	auto := []string{"auto", "maybe"}
	i = sort.SearchStrings(auto, ls)
	if i < len(auto) && auto[i] == ls {
		*c = ColorTTYOnly
		return nil
	}
	return fmt.Errorf("invalid arg to -color: %s", s)
}
func (c *ColorOption) String() string {
	switch *c {
	case ColorNo:
		return "no"
	case ColorYes:
		return "yes"
	case ColorTTYOnly:
		return "auto"
	}
	panic(fmt.Errorf("unknown color"))
}

var color ColorOption = ColorTTYOnly
var quiet bool

func usage() {
	op := flag.CommandLine.Output()
	fmt.Fprintf(op, "Usage of %s:\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Fprintln(op, "Return values:")
	fmt.Fprintln(op, "  0    No WX segments detected")
	fmt.Fprintln(op, "  1    At least 1 WX segment detected")
	fmt.Fprintln(op, "  >=2  Error")
}

func main() {
	flag.Usage = usage
	flag.Var(&color, "color", "Highlight using color")
	flag.BoolVar(&quiet, "q", false, "")
	flag.BoolVar(&quiet, "quiet", false, "Print nothing")
	flag.Parse()

	if color == ColorTTYOnly {
		if isatty.IsTerminal(os.Stdout.Fd()) {
			color = ColorYes
		} else {
			color = ColorNo
		}
	}

	self := os.Getpid()
	pidArg := strconv.FormatInt(int64(self), 10)

	r, w := io.Pipe()

	// Start searching for any lines where the protection field contains
	// "wx".  The read will block until data is actually written to the
	// write end of the pipe.
	//
	// Printing is finished once a value is sent over the 'found' channel.
	found := printWX(r)

	exitCode := 0
	cmd := exec.Command("procstat", "-h", "-v", pidArg)
	cmd.Stdout = w
	cmd.Stderr = os.Stderr

	err := cmd.Start()

	if err == nil {
		// Execution has begun, wait for procstat to exit and set this
		// program's exit code to the exit code of procstat.
		err = cmd.Wait()
		exitCode = cmd.ProcessState.ExitCode()
	} else {
		// Else procstat could not be started, probably because of
		// $PATH.  EX_OSERR is defined as 71 in sysexits(3).
		exitCode = 71
	}

	// Close the write end of the pipe.
	// If err == nil, the read end will see EOF after all data has been
	// read.  Otherwise, it will see the error.
	w.CloseWithError(err)

	// Block until printWX has completed its work.
	//
	// With the write end of the pipe closed, all data has been sent, and
	// once the read end of the pipe sees EOF or an error, that's when the
	// 'done' channel will finally be able to receive a value.  After that,
	// the read end of the pipe may be closed.
	for wxFound := range found {
		r.Close()
		if wxFound && exitCode == 0 {
			exitCode++
		}
	}

	os.Exit(exitCode)
}

const (
	VM_PID = iota
	VM_START
	VM_END
	VM_PROT
	VM_PAGES
	VM_PRIV_PAGES
	VM_REF_COUNT
	VM_SHADOW_COUNT
	VM_FLAGS
	VM_OBJTYPE
	VM_PATH
)

const (
	// Remove bold, green foreground
	ATTR_LINE_MATCH = "\033[22;32m"
	// Normal
	ATTR_NORMAL = "\033[0m"
	// Add bold, red foreground
	ATTR_STR_MATCH = "\033[1;31m"
)

func printWX(r io.Reader) <-chan bool {
	// Create a channel to send a "finished printing" message/signal.
	found := make(chan bool)

	// Start a goroutine that reads lines from the Reader, printing any
	// with a 3-byte protection field ending with "wx".
	go func() {
		wxFound := false
		scn := bufio.NewScanner(r)
		// Read lines until EOF or error and lines matching the
		// specified criteria ("wx" in the protection field in this
		// case).
		for scn.Scan() {
			line := scn.Text()
			fields := strings.Fields(line)

			if fields[VM_PROT][1:3] != "wx" {
				// No match, print normal if not quiet and
				// continue reading.
				if !quiet {
					fmt.Println(line)
				}
				continue
			}

			wxFound = true
			if !quiet {
				i := strings.Index(line, fields[VM_PROT])
				if color == ColorYes {
					fmt.Print(ATTR_LINE_MATCH, line[:i])
					fmt.Print(ATTR_STR_MATCH, line[i:i+3])
					fmt.Print(ATTR_LINE_MATCH, line[i+3:])
					fmt.Println(ATTR_NORMAL)
				} else {
					fmt.Print("***")
					fmt.Println(line)
				}
			}
		}

		// Print the error, if there was one, and send the message that
		// printing has finished.
		if err := scn.Err(); err != nil {
			fmt.Fprintln(os.Stderr, "error reading input:", err)
		}
		found <- wxFound

		// Signal that printing is complete.
		close(found)
	}()

	// Return the channel, so the main function can receive the
	// message/signal that the printing goroutine has completed.
	return found
}