github.com/SagerNet/gvisor@v0.0.0-20210707092255-7731c139d75c/pkg/shim/proc/io.go (about)

     1  // Copyright 2018 The containerd Authors.
     2  // Copyright 2018 The gVisor Authors.
     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  //     https://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 proc
    17  
    18  import (
    19  	"context"
    20  	"fmt"
    21  	"io"
    22  	"os"
    23  	"sync"
    24  	"sync/atomic"
    25  
    26  	"github.com/containerd/containerd/log"
    27  	"github.com/containerd/fifo"
    28  	runc "github.com/containerd/go-runc"
    29  	"golang.org/x/sys/unix"
    30  )
    31  
    32  // TODO(random-liu): This file can be a util.
    33  
    34  var bufPool = sync.Pool{
    35  	New: func() interface{} {
    36  		buffer := make([]byte, 32<<10)
    37  		return &buffer
    38  	},
    39  }
    40  
    41  func copyPipes(ctx context.Context, rio runc.IO, stdin, stdout, stderr string, wg *sync.WaitGroup) error {
    42  	var sameFile *countingWriteCloser
    43  	for _, i := range []struct {
    44  		name string
    45  		dest func(wc io.WriteCloser, rc io.Closer)
    46  	}{
    47  		{
    48  			name: stdout,
    49  			dest: func(wc io.WriteCloser, rc io.Closer) {
    50  				wg.Add(1)
    51  				go func() {
    52  					p := bufPool.Get().(*[]byte)
    53  					defer bufPool.Put(p)
    54  					if _, err := io.CopyBuffer(wc, rio.Stdout(), *p); err != nil {
    55  						log.G(ctx).Warn("error copying stdout")
    56  					}
    57  					wg.Done()
    58  					wc.Close()
    59  					if rc != nil {
    60  						rc.Close()
    61  					}
    62  				}()
    63  			},
    64  		}, {
    65  			name: stderr,
    66  			dest: func(wc io.WriteCloser, rc io.Closer) {
    67  				wg.Add(1)
    68  				go func() {
    69  					p := bufPool.Get().(*[]byte)
    70  					defer bufPool.Put(p)
    71  					if _, err := io.CopyBuffer(wc, rio.Stderr(), *p); err != nil {
    72  						log.G(ctx).Warn("error copying stderr")
    73  					}
    74  					wg.Done()
    75  					wc.Close()
    76  					if rc != nil {
    77  						rc.Close()
    78  					}
    79  				}()
    80  			},
    81  		},
    82  	} {
    83  		ok, err := isFifo(i.name)
    84  		if err != nil {
    85  			return err
    86  		}
    87  		var (
    88  			fw io.WriteCloser
    89  			fr io.Closer
    90  		)
    91  		if ok {
    92  			if fw, err = fifo.OpenFifo(ctx, i.name, unix.O_WRONLY, 0); err != nil {
    93  				return fmt.Errorf("gvisor-containerd-shim: opening %s failed: %s", i.name, err)
    94  			}
    95  			if fr, err = fifo.OpenFifo(ctx, i.name, unix.O_RDONLY, 0); err != nil {
    96  				return fmt.Errorf("gvisor-containerd-shim: opening %s failed: %s", i.name, err)
    97  			}
    98  		} else {
    99  			if sameFile != nil {
   100  				sameFile.count++
   101  				i.dest(sameFile, nil)
   102  				continue
   103  			}
   104  			if fw, err = os.OpenFile(i.name, unix.O_WRONLY|unix.O_APPEND, 0); err != nil {
   105  				return fmt.Errorf("gvisor-containerd-shim: opening %s failed: %s", i.name, err)
   106  			}
   107  			if stdout == stderr {
   108  				sameFile = &countingWriteCloser{
   109  					WriteCloser: fw,
   110  					count:       1,
   111  				}
   112  			}
   113  		}
   114  		i.dest(fw, fr)
   115  	}
   116  	if stdin == "" {
   117  		return nil
   118  	}
   119  	f, err := fifo.OpenFifo(context.Background(), stdin, unix.O_RDONLY|unix.O_NONBLOCK, 0)
   120  	if err != nil {
   121  		return fmt.Errorf("gvisor-containerd-shim: opening %s failed: %s", stdin, err)
   122  	}
   123  	go func() {
   124  		p := bufPool.Get().(*[]byte)
   125  		defer bufPool.Put(p)
   126  
   127  		io.CopyBuffer(rio.Stdin(), f, *p)
   128  		rio.Stdin().Close()
   129  		f.Close()
   130  	}()
   131  	return nil
   132  }
   133  
   134  // countingWriteCloser masks io.Closer() until close has been invoked a certain number of times.
   135  type countingWriteCloser struct {
   136  	io.WriteCloser
   137  	count int64
   138  }
   139  
   140  func (c *countingWriteCloser) Close() error {
   141  	if atomic.AddInt64(&c.count, -1) > 0 {
   142  		return nil
   143  	}
   144  	return c.WriteCloser.Close()
   145  }
   146  
   147  // isFifo checks if a file is a fifo.
   148  //
   149  // If the file does not exist then it returns false.
   150  func isFifo(path string) (bool, error) {
   151  	stat, err := os.Stat(path)
   152  	if err != nil {
   153  		if os.IsNotExist(err) {
   154  			return false, nil
   155  		}
   156  		return false, err
   157  	}
   158  	if stat.Mode()&os.ModeNamedPipe == os.ModeNamedPipe {
   159  		return true, nil
   160  	}
   161  	return false, nil
   162  }