github.com/system-transparency/u-root@v6.0.1-0.20190919065413-ed07a650de4c+incompatible/cmds/core/grep/grep.go (about)

     1  // Copyright 2012-2017 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Concurrent, parallel grep.
     6  //
     7  // Synopsis:
     8  //     grep [-vrlq] [FILE]...
     9  //
    10  // Description:
    11  //     It has to deal with the EMFILE limit. To do so we have one chan that is
    12  //     bounded. From args, we use filepath.Walk to generate a chan of names.
    13  //     From that, we create a chan of grepCommands. From that, we create a chan
    14  //     of grepResults. The grepResults contain matches or not-matches only. If
    15  //     we are in -l mode, the goprocs handling the grep bail out as soon as the
    16  //     condition is met. This grep is about 2x faster than GNU grep for simple
    17  //     non-recursive greps and slower as soon as filepath. Walk enters the
    18  //     picture. Let's fix this.
    19  //
    20  // Options:
    21  //     -v: print only non-matching lines
    22  //     -r: recursive
    23  //     -l: list only files
    24  //     -q: don't print matches; exit on first match
    25  package main
    26  
    27  import (
    28  	"bufio"
    29  	"flag"
    30  	"fmt"
    31  	"log"
    32  	"os"
    33  	"path/filepath"
    34  	"regexp"
    35  	"strings"
    36  )
    37  
    38  type grepResult struct {
    39  	match bool
    40  	c     *grepCommand
    41  	line  *string
    42  }
    43  
    44  type grepCommand struct {
    45  	name string
    46  	*os.File
    47  }
    48  
    49  type oneGrep struct {
    50  	c chan *grepResult
    51  }
    52  
    53  var (
    54  	match           = flag.Bool("v", true, "Print only non-matching lines")
    55  	recursive       = flag.Bool("r", false, "recursive")
    56  	noshowmatch     = flag.Bool("l", false, "list only files")
    57  	quiet           = flag.Bool("q", false, "Don't print matches; exit on first match")
    58  	count           = flag.Bool("c", false, "Just show counts")
    59  	caseinsensitive = flag.Bool("i", false, "case-insensitive matching")
    60  	showname        bool
    61  	allGrep         = make(chan *oneGrep)
    62  	nGrep           int
    63  	matchCount      int
    64  )
    65  
    66  // grep reads data from the os.File embedded in grepCommand.
    67  // It creates a chan of grepResults and pushes a pointer to it into allGrep.
    68  // It matches each line against the re and pushes the matching result
    69  // into the chan.
    70  // Bug: this chan should be created by the caller and passed in
    71  // to preserve file name order. Oops.
    72  // If we are only looking for a match, we exit as soon as the condition is met.
    73  // "match" means result of re.Match == match flag.
    74  func grep(f *grepCommand, re *regexp.Regexp) {
    75  	nGrep++
    76  	r := bufio.NewReader(f)
    77  	res := make(chan *grepResult, 1)
    78  	allGrep <- &oneGrep{res}
    79  	for {
    80  		if i, err := r.ReadString('\n'); err == nil {
    81  			m := re.Match([]byte(i))
    82  			if m == *match {
    83  				res <- &grepResult{re.Match([]byte(i)), f, &i}
    84  				if *noshowmatch {
    85  					break
    86  				}
    87  			}
    88  		} else {
    89  			break
    90  		}
    91  	}
    92  	close(res)
    93  	f.Close()
    94  }
    95  
    96  func printmatch(r *grepResult) {
    97  	var prefix string
    98  	if r.match == *match {
    99  		matchCount++
   100  	}
   101  	if *count {
   102  		return
   103  	}
   104  	if showname {
   105  		fmt.Printf("%v", r.c.name)
   106  		prefix = ":"
   107  	}
   108  	if *noshowmatch {
   109  		return
   110  	}
   111  	if r.match == *match {
   112  		fmt.Printf("%v%v", prefix, *r.line)
   113  	}
   114  }
   115  
   116  func main() {
   117  	r := ".*"
   118  	flag.Parse()
   119  	a := flag.Args()
   120  	if len(a) > 0 {
   121  		r = a[0]
   122  	}
   123  	if *caseinsensitive && !strings.HasPrefix(r, "(?i)") {
   124  		r = "(?i)" + r
   125  	}
   126  	re := regexp.MustCompile(r)
   127  	// very special case, just stdin ...
   128  	if len(a) < 2 {
   129  		go grep(&grepCommand{"<stdin>", os.Stdin}, re)
   130  	} else {
   131  		showname = len(a[1:]) > 1
   132  		// generate a chan of file names, bounded by the size of the chan. This in turn
   133  		// throttles the opens.
   134  		treenames := make(chan string, 128)
   135  		go func() {
   136  			for _, v := range a[1:] {
   137  				// we could parallelize the open part but people might want
   138  				// things to be in order. I don't care but who knows.
   139  				// just ignore the errors. If there is not a single one that works,
   140  				// then all the sizes will be 0 and we'll just fall through.
   141  				filepath.Walk(v, func(name string, fi os.FileInfo, err error) error {
   142  					if err != nil {
   143  						// This is non-fatal because grep searches through
   144  						// all the files it has access to.
   145  						log.Print(err)
   146  						return nil
   147  					}
   148  					if fi.IsDir() && !*recursive {
   149  						fmt.Fprintf(os.Stderr, "grep: %v: Is a directory\n", name)
   150  						return filepath.SkipDir
   151  					}
   152  					if err != nil {
   153  						fmt.Fprintf(os.Stderr, "%v: %v\n", name, err)
   154  						return err
   155  					}
   156  					treenames <- name
   157  					return nil
   158  				})
   159  			}
   160  			close(treenames)
   161  		}()
   162  
   163  		files := make(chan *grepCommand)
   164  		// convert the file names to a stream of os.File
   165  		go func() {
   166  			for i := range treenames {
   167  				fp, err := os.Open(i)
   168  				if err != nil {
   169  					fmt.Fprintf(os.Stderr, "can't open %s: %v\n", i, err)
   170  					continue
   171  				}
   172  				files <- &grepCommand{i, fp}
   173  			}
   174  			close(files)
   175  		}()
   176  		// now kick off the greps
   177  		// bug: file name order is not preserved here. Darn.
   178  
   179  		for f := range files {
   180  			go grep(f, re)
   181  		}
   182  	}
   183  
   184  	for c := range allGrep {
   185  		for r := range c.c {
   186  			// exit on first match.
   187  			if *quiet {
   188  				os.Exit(0)
   189  			}
   190  			printmatch(r)
   191  		}
   192  		nGrep--
   193  		if nGrep == 0 {
   194  			break
   195  		}
   196  	}
   197  	if *quiet {
   198  		os.Exit(1)
   199  	}
   200  	if *count {
   201  		fmt.Printf("%d\n", matchCount)
   202  	}
   203  }