vendor/github.com/tonistiigi/fsutil/diskwriter.go

   1 // +build linux windows
   2
   3 package fsutil
   4
   5 import (
   6         "archive/tar"
   7         "crypto/sha256"
   8         "encoding/hex"
   9         "hash"
  10         "io"
  11         "os"
  12         "path/filepath"
  13         "strconv"
  14         "sync"
  15         "time"
  16
  17         "github.com/docker/docker/pkg/archive"
  18         "github.com/docker/docker/pkg/tarsum"
  19         "github.com/pkg/errors"
  20         "golang.org/x/net/context"
  21         "golang.org/x/sync/errgroup"
  22 )
  23
  24 type WriteToFunc func(context.Context, string, io.WriteCloser) error
  25
  26 type DiskWriterOpt struct {
  27         AsyncDataCb WriteToFunc
  28         SyncDataCb  WriteToFunc
  29         NotifyCb    func(ChangeKind, string, os.FileInfo, error) error
  30 }
  31
  32 type DiskWriter struct {
  33         opt  DiskWriterOpt
  34         dest string
  35
  36         wg     sync.WaitGroup
  37         ctx    context.Context
  38         cancel func()
  39         eg     *errgroup.Group
  40 }
  41
  42 func NewDiskWriter(ctx context.Context, dest string, opt DiskWriterOpt) (*DiskWriter, error) {
  43         if opt.SyncDataCb == nil && opt.AsyncDataCb == nil {
  44                 return nil, errors.New("no data callback specified")
  45         }
  46         if opt.SyncDataCb != nil && opt.AsyncDataCb != nil {
  47                 return nil, errors.New("can't specify both sync and async data callbacks")
  48         }
  49
  50         ctx, cancel := context.WithCancel(ctx)
  51         eg, ctx := errgroup.WithContext(ctx)
  52
  53         return &DiskWriter{
  54                 opt:    opt,
  55                 dest:   dest,
  56                 eg:     eg,
  57                 ctx:    ctx,
  58                 cancel: cancel,
  59         }, nil
  60 }
  61
  62 func (dw *DiskWriter) Wait(ctx context.Context) error {
  63         return dw.eg.Wait()
  64 }
  65
  66 func (dw *DiskWriter) HandleChange(kind ChangeKind, p string, fi os.FileInfo, err error) (retErr error) {
  67         if err != nil {
  68                 return err
  69         }
  70
  71         select {
  72         case <-dw.ctx.Done():
  73                 return dw.ctx.Err()
  74         default:
  75         }
  76
  77         defer func() {
  78                 if retErr != nil {
  79                         dw.cancel()
  80                 }
  81         }()
  82
  83         p = filepath.FromSlash(p)
  84
  85         destPath := filepath.Join(dw.dest, p)
  86
  87         if kind == ChangeKindDelete {
  88                 // todo: no need to validate if diff is trusted but is it always?
  89                 if err := os.RemoveAll(destPath); err != nil {
  90                         return errors.Wrapf(err, "failed to remove: %s", destPath)
  91                 }
  92                 if dw.opt.NotifyCb != nil {
  93                         if err := dw.opt.NotifyCb(kind, p, nil, nil); err != nil {
  94                                 return err
  95                         }
  96                 }
  97                 return nil
  98         }
  99
 100         stat, ok := fi.Sys().(*Stat)
 101         if !ok {
 102                 return errors.Errorf("%s invalid change without stat information", p)
 103         }
 104
 105         rename := true
 106         oldFi, err := os.Lstat(destPath)
 107         if err != nil {
 108                 if os.IsNotExist(err) {
 109                         if kind != ChangeKindAdd {
 110                                 return errors.Wrapf(err, "invalid addition: %s", destPath)
 111                         }
 112                         rename = false
 113                 } else {
 114                         return errors.Wrapf(err, "failed to stat %s", destPath)
 115                 }
 116         }
 117
 118         if oldFi != nil && fi.IsDir() && oldFi.IsDir() {
 119                 if err := rewriteMetadata(destPath, stat); err != nil {
 120                         return errors.Wrapf(err, "error setting dir metadata for %s", destPath)
 121                 }
 122                 return nil
 123         }
 124
 125         newPath := destPath
 126         if rename {
 127                 newPath = filepath.Join(filepath.Dir(destPath), ".tmp."+nextSuffix())
 128         }
 129
 130         isRegularFile := false
 131
 132         switch {
 133         case fi.IsDir():
 134                 if err := os.Mkdir(newPath, fi.Mode()); err != nil {
 135                         return errors.Wrapf(err, "failed to create dir %s", newPath)
 136                 }
 137         case fi.Mode()&os.ModeDevice != 0 || fi.Mode()&os.ModeNamedPipe != 0:
 138                 if err := handleTarTypeBlockCharFifo(newPath, stat); err != nil {
 139                         return errors.Wrapf(err, "failed to create device %s", newPath)
 140                 }
 141         case fi.Mode()&os.ModeSymlink != 0:
 142                 if err := os.Symlink(stat.Linkname, newPath); err != nil {
 143                         return errors.Wrapf(err, "failed to symlink %s", newPath)
 144                 }
 145         case stat.Linkname != "":
 146                 if err := os.Link(filepath.Join(dw.dest, stat.Linkname), newPath); err != nil {
 147                         return errors.Wrapf(err, "failed to link %s to %s", newPath, stat.Linkname)
 148                 }
 149         default:
 150                 isRegularFile = true
 151                 file, err := os.OpenFile(newPath, os.O_CREATE|os.O_WRONLY, fi.Mode()) //todo: windows
 152                 if err != nil {
 153                         return errors.Wrapf(err, "failed to create %s", newPath)
 154                 }
 155                 if dw.opt.SyncDataCb != nil {
 156                         if err := dw.processChange(ChangeKindAdd, p, fi, file); err != nil {
 157                                 file.Close()
 158                                 return err
 159                         }
 160                         break
 161                 }
 162                 if err := file.Close(); err != nil {
 163                         return errors.Wrapf(err, "failed to close %s", newPath)
 164                 }
 165         }
 166
 167         if err := rewriteMetadata(newPath, stat); err != nil {
 168                 return errors.Wrapf(err, "error setting metadata for %s", newPath)
 169         }
 170
 171         if rename {
 172                 if err := os.Rename(newPath, destPath); err != nil {
 173                         return errors.Wrapf(err, "failed to rename %s to %s", newPath, destPath)
 174                 }
 175         }
 176
 177         if isRegularFile {
 178                 if dw.opt.AsyncDataCb != nil {
 179                         dw.requestAsyncFileData(p, destPath, fi)
 180                 }
 181         } else {
 182                 return dw.processChange(kind, p, fi, nil)
 183         }
 184
 185         return nil
 186 }
 187
 188 func (dw *DiskWriter) requestAsyncFileData(p, dest string, fi os.FileInfo) {
 189         // todo: limit worker threads
 190         dw.eg.Go(func() error {
 191                 if err := dw.processChange(ChangeKindAdd, p, fi, &lazyFileWriter{
 192                         dest: dest,
 193                 }); err != nil {
 194                         return err
 195                 }
 196                 return chtimes(dest, fi.ModTime().UnixNano()) // TODO: parent dirs
 197         })
 198 }
 199
 200 func (dw *DiskWriter) processChange(kind ChangeKind, p string, fi os.FileInfo, w io.WriteCloser) error {
 201         origw := w
 202         var hw *hashedWriter
 203         if dw.opt.NotifyCb != nil {
 204                 var err error
 205                 if hw, err = newHashWriter(p, fi, w); err != nil {
 206                         return err
 207                 }
 208                 w = hw
 209         }
 210         if origw != nil {
 211                 fn := dw.opt.SyncDataCb
 212                 if fn == nil && dw.opt.AsyncDataCb != nil {
 213                         fn = dw.opt.AsyncDataCb
 214                 }
 215                 if err := fn(dw.ctx, p, w); err != nil {
 216                         return err
 217                 }
 218         } else {
 219                 if hw != nil {
 220                         hw.Close()
 221                 }
 222         }
 223         if hw != nil {
 224                 return dw.opt.NotifyCb(kind, p, hw, nil)
 225         }
 226         return nil
 227 }
 228
 229 type hashedWriter struct {
 230         os.FileInfo
 231         io.Writer
 232         h   hash.Hash
 233         w   io.WriteCloser
 234         sum string
 235 }
 236
 237 func newHashWriter(p string, fi os.FileInfo, w io.WriteCloser) (*hashedWriter, error) {
 238         h, err := NewTarsumHash(p, fi)
 239         if err != nil {
 240                 return nil, err
 241         }
 242         hw := &hashedWriter{
 243                 FileInfo: fi,
 244                 Writer:   io.MultiWriter(w, h),
 245                 h:        h,
 246                 w:        w,
 247         }
 248         return hw, nil
 249 }
 250
 251 func (hw *hashedWriter) Close() error {
 252         hw.sum = string(hex.EncodeToString(hw.h.Sum(nil)))
 253         if hw.w != nil {
 254                 return hw.w.Close()
 255         }
 256         return nil
 257 }
 258
 259 func (hw *hashedWriter) Hash() string {
 260         return hw.sum
 261 }
 262
 263 type lazyFileWriter struct {
 264         dest string
 265         ctx  context.Context
 266         f    *os.File
 267 }
 268
 269 func (lfw *lazyFileWriter) Write(dt []byte) (int, error) {
 270         if lfw.f == nil {
 271                 file, err := os.OpenFile(lfw.dest, os.O_WRONLY, 0) //todo: windows
 272                 if err != nil {
 273                         return 0, errors.Wrapf(err, "failed to open %s", lfw.dest)
 274                 }
 275                 lfw.f = file
 276         }
 277         return lfw.f.Write(dt)
 278 }
 279
 280 func (lfw *lazyFileWriter) Close() error {
 281         if lfw.f != nil {
 282                 return lfw.f.Close()
 283         }
 284         return nil
 285 }
 286
 287 func mkdev(major int64, minor int64) uint32 {
 288         return uint32(((minor & 0xfff00) << 12) | ((major & 0xfff) << 8) | (minor & 0xff))
 289 }
 290
 291 // Random number state.
 292 // We generate random temporary file names so that there's a good
 293 // chance the file doesn't exist yet - keeps the number of tries in
 294 // TempFile to a minimum.
 295 var rand uint32
 296 var randmu sync.Mutex
 297
 298 func reseed() uint32 {
 299         return uint32(time.Now().UnixNano() + int64(os.Getpid()))
 300 }
 301
 302 func nextSuffix() string {
 303         randmu.Lock()
 304         r := rand
 305         if r == 0 {
 306                 r = reseed()
 307         }
 308         r = r*1664525 + 1013904223 // constants from Numerical Recipes
 309         rand = r
 310         randmu.Unlock()
 311         return strconv.Itoa(int(1e9 + r%1e9))[1:]
 312 }
 313
 314 func NewTarsumHash(p string, fi os.FileInfo) (hash.Hash, error) {
 315         stat, ok := fi.Sys().(*Stat)
 316         link := ""
 317         if ok {
 318                 link = stat.Linkname
 319         }
 320         if fi.IsDir() {
 321                 p += string(os.PathSeparator)
 322         }
 323         h, err := archive.FileInfoHeader(p, fi, link)
 324         if err != nil {
 325                 return nil, err
 326         }
 327         h.Name = p
 328         if ok {
 329                 h.Uid = int(stat.Uid)
 330                 h.Gid = int(stat.Gid)
 331                 h.Linkname = stat.Linkname
 332                 if stat.Xattrs != nil {
 333                         h.Xattrs = make(map[string]string)
 334                         for k, v := range stat.Xattrs {
 335                                 h.Xattrs[k] = string(v)
 336                         }
 337                 }
 338         }
 339         tsh := &tarsumHash{h: h, Hash: sha256.New()}
 340         tsh.Reset()
 341         return tsh, nil
 342 }
 343
 344 // Reset resets the Hash to its initial state.
 345 func (tsh *tarsumHash) Reset() {
 346         tsh.Hash.Reset()
 347         tarsum.WriteV1Header(tsh.h, tsh.Hash)
 348 }
 349
 350 type tarsumHash struct {
 351         hash.Hash
 352         h *tar.Header
 353 }