ioutil.go raw

   1  // Copyright ©2015 The Go Authors
   2  // Copyright ©2015 Steve Francia <spf@spf13.com>
   3  //
   4  // Licensed under the Apache License, Version 2.0 (the "License");
   5  // you may not use this file except in compliance with the License.
   6  // You may obtain a copy of the License at
   7  //
   8  //     http://www.apache.org/licenses/LICENSE-2.0
   9  //
  10  // Unless required by applicable law or agreed to in writing, software
  11  // distributed under the License is distributed on an "AS IS" BASIS,
  12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  // See the License for the specific language governing permissions and
  14  // limitations under the License.
  15  
  16  package afero
  17  
  18  import (
  19  	"bytes"
  20  	"io"
  21  	"os"
  22  	"path/filepath"
  23  	"sort"
  24  	"strconv"
  25  	"strings"
  26  	"sync"
  27  	"time"
  28  )
  29  
  30  // byName implements sort.Interface.
  31  type byName []os.FileInfo
  32  
  33  func (f byName) Len() int           { return len(f) }
  34  func (f byName) Less(i, j int) bool { return f[i].Name() < f[j].Name() }
  35  func (f byName) Swap(i, j int)      { f[i], f[j] = f[j], f[i] }
  36  
  37  // ReadDir reads the directory named by dirname and returns
  38  // a list of sorted directory entries.
  39  func (a Afero) ReadDir(dirname string) ([]os.FileInfo, error) {
  40  	return ReadDir(a.Fs, dirname)
  41  }
  42  
  43  func ReadDir(fs Fs, dirname string) ([]os.FileInfo, error) {
  44  	f, err := fs.Open(dirname)
  45  	if err != nil {
  46  		return nil, err
  47  	}
  48  	list, err := f.Readdir(-1)
  49  	f.Close()
  50  	if err != nil {
  51  		return nil, err
  52  	}
  53  	sort.Sort(byName(list))
  54  	return list, nil
  55  }
  56  
  57  // ReadFile reads the file named by filename and returns the contents.
  58  // A successful call returns err == nil, not err == EOF. Because ReadFile
  59  // reads the whole file, it does not treat an EOF from Read as an error
  60  // to be reported.
  61  func (a Afero) ReadFile(filename string) ([]byte, error) {
  62  	return ReadFile(a.Fs, filename)
  63  }
  64  
  65  func ReadFile(fs Fs, filename string) ([]byte, error) {
  66  	f, err := fs.Open(filename)
  67  	if err != nil {
  68  		return nil, err
  69  	}
  70  	defer f.Close()
  71  	// It's a good but not certain bet that FileInfo will tell us exactly how much to
  72  	// read, so let's try it but be prepared for the answer to be wrong.
  73  	var n int64
  74  
  75  	if fi, err := f.Stat(); err == nil {
  76  		// Don't preallocate a huge buffer, just in case.
  77  		if size := fi.Size(); size < 1e9 {
  78  			n = size
  79  		}
  80  	}
  81  	// As initial capacity for readAll, use n + a little extra in case Size is zero,
  82  	// and to avoid another allocation after Read has filled the buffer.  The readAll
  83  	// call will read into its allocated internal buffer cheaply.  If the size was
  84  	// wrong, we'll either waste some space off the end or reallocate as needed, but
  85  	// in the overwhelmingly common case we'll get it just right.
  86  	return readAll(f, n+bytes.MinRead)
  87  }
  88  
  89  // readAll reads from r until an error or EOF and returns the data it read
  90  // from the internal buffer allocated with a specified capacity.
  91  func readAll(r io.Reader, capacity int64) (b []byte, err error) {
  92  	buf := bytes.NewBuffer(make([]byte, 0, capacity))
  93  	// If the buffer overflows, we will get bytes.ErrTooLarge.
  94  	// Return that as an error. Any other panic remains.
  95  	defer func() {
  96  		e := recover()
  97  		if e == nil {
  98  			return
  99  		}
 100  		if panicErr, ok := e.(error); ok && panicErr == bytes.ErrTooLarge {
 101  			err = panicErr
 102  		} else {
 103  			panic(e)
 104  		}
 105  	}()
 106  	_, err = buf.ReadFrom(r)
 107  	return buf.Bytes(), err
 108  }
 109  
 110  // ReadAll reads from r until an error or EOF and returns the data it read.
 111  // A successful call returns err == nil, not err == EOF. Because ReadAll is
 112  // defined to read from src until EOF, it does not treat an EOF from Read
 113  // as an error to be reported.
 114  func ReadAll(r io.Reader) ([]byte, error) {
 115  	return readAll(r, bytes.MinRead)
 116  }
 117  
 118  // WriteFile writes data to a file named by filename.
 119  // If the file does not exist, WriteFile creates it with permissions perm;
 120  // otherwise WriteFile truncates it before writing.
 121  func (a Afero) WriteFile(filename string, data []byte, perm os.FileMode) error {
 122  	return WriteFile(a.Fs, filename, data, perm)
 123  }
 124  
 125  func WriteFile(fs Fs, filename string, data []byte, perm os.FileMode) error {
 126  	f, err := fs.OpenFile(filename, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, perm)
 127  	if err != nil {
 128  		return err
 129  	}
 130  	n, err := f.Write(data)
 131  	if err == nil && n < len(data) {
 132  		err = io.ErrShortWrite
 133  	}
 134  	if err1 := f.Close(); err == nil {
 135  		err = err1
 136  	}
 137  	return err
 138  }
 139  
 140  // Random number state.
 141  // We generate random temporary file names so that there's a good
 142  // chance the file doesn't exist yet - keeps the number of tries in
 143  // TempFile to a minimum.
 144  var (
 145  	randNum uint32
 146  	randmu  sync.Mutex
 147  )
 148  
 149  func reseed() uint32 {
 150  	return uint32(time.Now().UnixNano() + int64(os.Getpid()))
 151  }
 152  
 153  func nextRandom() string {
 154  	randmu.Lock()
 155  	r := randNum
 156  	if r == 0 {
 157  		r = reseed()
 158  	}
 159  	r = r*1664525 + 1013904223 // constants from Numerical Recipes
 160  	randNum = r
 161  	randmu.Unlock()
 162  	return strconv.Itoa(int(1e9 + r%1e9))[1:]
 163  }
 164  
 165  // TempFile creates a new temporary file in the directory dir,
 166  // opens the file for reading and writing, and returns the resulting *os.File.
 167  // The filename is generated by taking pattern and adding a random
 168  // string to the end. If pattern includes a "*", the random string
 169  // replaces the last "*".
 170  // If dir is the empty string, TempFile uses the default directory
 171  // for temporary files (see os.TempDir).
 172  // Multiple programs calling TempFile simultaneously
 173  // will not choose the same file. The caller can use f.Name()
 174  // to find the pathname of the file. It is the caller's responsibility
 175  // to remove the file when no longer needed.
 176  func (a Afero) TempFile(dir, pattern string) (f File, err error) {
 177  	return TempFile(a.Fs, dir, pattern)
 178  }
 179  
 180  func TempFile(fs Fs, dir, pattern string) (f File, err error) {
 181  	if dir == "" {
 182  		dir = os.TempDir()
 183  	}
 184  
 185  	var prefix, suffix string
 186  	if pos := strings.LastIndex(pattern, "*"); pos != -1 {
 187  		prefix, suffix = pattern[:pos], pattern[pos+1:]
 188  	} else {
 189  		prefix = pattern
 190  	}
 191  
 192  	nconflict := 0
 193  	for i := 0; i < 10000; i++ {
 194  		name := filepath.Join(dir, prefix+nextRandom()+suffix)
 195  		f, err = fs.OpenFile(name, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o600)
 196  		if os.IsExist(err) {
 197  			if nconflict++; nconflict > 10 {
 198  				randmu.Lock()
 199  				randNum = reseed()
 200  				randmu.Unlock()
 201  			}
 202  			continue
 203  		}
 204  		break
 205  	}
 206  	return
 207  }
 208  
 209  // TempDir creates a new temporary directory in the directory dir
 210  // with a name beginning with prefix and returns the path of the
 211  // new directory.  If dir is the empty string, TempDir uses the
 212  // default directory for temporary files (see os.TempDir).
 213  // Multiple programs calling TempDir simultaneously
 214  // will not choose the same directory.  It is the caller's responsibility
 215  // to remove the directory when no longer needed.
 216  func (a Afero) TempDir(dir, prefix string) (name string, err error) {
 217  	return TempDir(a.Fs, dir, prefix)
 218  }
 219  
 220  func TempDir(fs Fs, dir, prefix string) (name string, err error) {
 221  	if dir == "" {
 222  		dir = os.TempDir()
 223  	}
 224  
 225  	nconflict := 0
 226  	for i := 0; i < 10000; i++ {
 227  		try := filepath.Join(dir, prefix+nextRandom())
 228  		err = fs.Mkdir(try, 0o700)
 229  		if os.IsExist(err) {
 230  			if nconflict++; nconflict > 10 {
 231  				randmu.Lock()
 232  				randNum = reseed()
 233  				randmu.Unlock()
 234  			}
 235  			continue
 236  		}
 237  		if err == nil {
 238  			name = try
 239  		}
 240  		break
 241  	}
 242  	return
 243  }
 244