github.phpd.cn/thought-machine/please@v12.2.0+incompatible/tools/jarcat/unzip/unzip.go (about)

     1  // Package unzip implements unzipping for jarcat.
     2  // We implement this to avoid needing a runtime dependency on unzip,
     3  // which is not a profound package but not installed everywhere by default.
     4  package unzip
     5  
     6  import (
     7  	"io"
     8  	"os"
     9  	"path"
    10  	"strings"
    11  	"sync"
    12  
    13  	"third_party/go/zip"
    14  )
    15  
    16  // concurrency controls the maximum level of concurrency we'll allow.
    17  const concurrency = 4
    18  
    19  // Extract extracts the contents of the given zipfile.
    20  func Extract(in, out, file, prefix string) error {
    21  	e := extractor{
    22  		In:     in,
    23  		Out:    out,
    24  		File:   file,
    25  		Prefix: prefix,
    26  		dirs:   map[string]struct{}{},
    27  	}
    28  	return e.Extract()
    29  }
    30  
    31  // An extractor extracts a single zipfile.
    32  type extractor struct {
    33  	In     string
    34  	Out    string
    35  	File   string
    36  	Prefix string
    37  	dirs   map[string]struct{}
    38  	mutex  sync.Mutex
    39  	wg     sync.WaitGroup
    40  	err    error
    41  }
    42  
    43  func (e *extractor) Extract() error {
    44  	r, err := zip.OpenReader(e.In)
    45  	if err != nil {
    46  		return err
    47  	}
    48  	defer r.Close()
    49  	ch := make(chan *zip.File, 100)
    50  	for i := 0; i < concurrency; i++ {
    51  		go e.consume(ch)
    52  	}
    53  	for _, f := range r.File {
    54  		if e.File != "" && f.Name != e.File {
    55  			continue
    56  		}
    57  		// This will mean that empty directories aren't created. We might need to fix that at some point.
    58  		if f.Mode()&os.ModeDir == 0 {
    59  			e.wg.Add(1)
    60  			ch <- f
    61  		}
    62  	}
    63  	e.wg.Wait()
    64  	close(ch)
    65  	return e.err
    66  }
    67  
    68  func (e *extractor) consume(ch <-chan *zip.File) {
    69  	for f := range ch {
    70  		if err := e.extractFile(f); err != nil {
    71  			e.err = err
    72  		}
    73  		e.wg.Done()
    74  	}
    75  }
    76  
    77  func (e *extractor) extractFile(f *zip.File) error {
    78  	if e.Prefix != "" {
    79  		if !strings.HasPrefix(f.Name, e.Prefix) {
    80  			return nil
    81  		}
    82  		f.Name = strings.TrimLeft(strings.TrimPrefix(f.Name, e.Prefix), "/")
    83  	}
    84  	r, err := f.Open()
    85  	if err != nil {
    86  		return err
    87  	}
    88  	defer r.Close()
    89  	out := path.Join(e.Out, f.Name)
    90  	if e.File != "" {
    91  		out = e.Out
    92  	}
    93  	if err := e.makeDir(out); err != nil {
    94  		return err
    95  	}
    96  	o, err := os.OpenFile(out, os.O_WRONLY|os.O_CREATE, f.Mode())
    97  	if err != nil {
    98  		return err
    99  	}
   100  	defer o.Close()
   101  	_, err = io.Copy(o, r)
   102  	return err
   103  }
   104  
   105  func (e *extractor) makeDir(filename string) error {
   106  	dir := path.Dir(filename)
   107  	e.mutex.Lock()
   108  	defer e.mutex.Unlock()
   109  	if _, present := e.dirs[dir]; !present {
   110  		if err := os.MkdirAll(dir, 0755); err != nil {
   111  			return err
   112  		}
   113  		e.dirs[dir] = struct{}{}
   114  	}
   115  	return nil
   116  }