diff --git a/raw_linux.go b/raw_linux.go new file mode 100644 index 0000000..f46db4c --- /dev/null +++ b/raw_linux.go @@ -0,0 +1,251 @@ +/* + Copyright The containerd Authors. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +*/ + +package fifo + +import ( + "io" + "sync" + "sync/atomic" + "syscall" + + "golang.org/x/sys/unix" +) + +const ( + spliceMax int64 = 1 << 62 + spliceToEOF int64 = -1 +) + +var ( + spliceSupported int32 = 1 + bufPool = &sync.Pool{ + New: func() interface{} { + buf := make([]byte, 1<<20) + return &buf + }, + } +) + +func (f *fifo) ReadFrom(r io.Reader) (int64, error) { + if f.flag&(syscall.O_WRONLY|syscall.O_RDWR) == 0 { + return 0, ErrWrToRDONLY + } + select { + case <-f.opened: + return f.readFrom(r) + default: + } + select { + case <-f.opened: + return f.readFrom(r) + case <-f.closed: + return 0, ErrWriteClosed + } +} + +func (f *fifo) readFrom(r io.Reader) (int64, error) { + if atomic.LoadInt32(&spliceSupported) != 1 { + return copyBuffer(f.file, r) + } + + remain := spliceToEOF + lr, ok := r.(*io.LimitedReader) + if ok { + remain, r = lr.N, lr.R + if remain <= 0 { + return 0, nil + } + } + + rscI, ok := r.(syscall.Conn) + if !ok { + return copyBuffer(f.file, r) + } + + rsc, err := rscI.SyscallConn() + if err != nil { + if lr != nil { + r = lr + } + return copyBuffer(f.file, r) + } + + wsc, err := f.SyscallConn() + if err != nil { + if lr != nil { + r = lr + } + return copyBuffer(f.file, r) + } + + handled, written, err := doRawCopy(rsc, wsc, remain) + if err != nil { + return written, err + } + if !handled { + if lr != nil { + r = lr + } + return copyBuffer(f.file, r) + } + + return written, nil +} + +func (f *fifo) WriteTo(w io.Writer) (int64, error) { + if f.flag&syscall.O_WRONLY > 0 { + return 0, ErrRdFrmWRONLY + } + select { + case <-f.opened: + return f.writeTo(w) + default: + } + + select { + case <-f.opened: + return f.writeTo(w) + case <-f.closed: + return 0, ErrWriteClosed + } +} + +func (f *fifo) writeTo(w io.Writer) (int64, error) { + if atomic.LoadInt32(&spliceSupported) != 1 { + return copyBuffer(w, f.file) + } + + wscI, ok := w.(syscall.Conn) + if !ok { + return copyBuffer(w, f.file) + } + + wsc, err := wscI.SyscallConn() + if err != nil { + return copyBuffer(w, f.file) + } + + rsc, err := f.SyscallConn() + if err != nil { + return copyBuffer(w, f.file) + } + + handled, written, err := doRawCopy(rsc, wsc, spliceToEOF) + if err != nil { + return written, err + } + if !handled { + return copyBuffer(w, f.file) + } + return written, err +} + +func doRawCopy(rsc, wsc syscall.RawConn, remain int64) (handled bool, written int64, _ error) { + var ( + spliceErr error + writeErr error + copyToEOF = remain == spliceToEOF + ) + + if remain == spliceToEOF { + remain = spliceMax + } + + // Hear the RawConn Read/Write methods allow us to utilize the go runtime + // poller to wait for the file descriptors to be ready for reads or writes. + // + // Read/Write will sleep the goroutine until the file descriptor is ready. + // Once we are inside the function we've passed in, we know the FD is ready. + // + // Read/Write both run the function they are passed repeatedly until the + // function returns true. + err := rsc.Read(func(rfd uintptr) bool { + err := wsc.Write(func(wfd uintptr) bool { + for copyToEOF || remain > 0 { + // We always use NONBLOCK here. If the file descriptor(s) is not + // opened with O_NONBLOCK then splice just blocks like normal. + // If they opened with O_NONBLOCK, then `unix.Splice` returns + // with EAGAIN when either the read or write would block. + n, err := unix.Splice(int(rfd), nil, int(wfd), nil, int(remain), unix.SPLICE_F_MOVE|unix.SPLICE_F_NONBLOCK) + if n > 0 { + written += n + if !copyToEOF { + remain -= n + } + } + + switch err { + case unix.ENOSYS: + // splice not supported on kernel + atomic.StoreInt32(&spliceSupported, 0) + return true + case syscall.EINVAL, syscall.EOPNOTSUPP, syscall.EPERM: + // In all these cases, there is no data transferred + return true + case nil: + handled = true + if n == 0 { + // At EOF + return true + } + case unix.EINTR: + continue + case unix.EAGAIN: + // Normally we'd want to return false here, because this just means + // we need it to wait for the fd to be ready again, however we don't know + // which fd needs to be waited on. + // So, break out of the write func and let `Read` return false so we end up + // waiting for both fd's to be ready. + return true + default: + spliceErr = err + handled = true + return true + } + } + return true + }) + if err != nil { + writeErr = err + return true + } + if spliceErr != nil { + // I don't like this but it made the linter happy. All hail the mighty linter. + // If splice returned EAGAIN we should return false so we can + // wait for the FD to be ready again. Otherwise we just want to exit + // early. + return spliceErr != unix.EAGAIN + } + + return true + }) + + if spliceErr != nil { + return handled, written, spliceErr + } + if writeErr != nil { + return handled, written, writeErr + } + return handled, written, err +} + +func copyBuffer(w io.Writer, r io.Reader) (int64, error) { + buf := bufPool.Get().(*[]byte) + n, err := io.CopyBuffer(w, r, *buf) + bufPool.Put(buf) + return n, err +} diff --git a/raw_linux_test.go b/raw_linux_test.go new file mode 100644 index 0000000..ace8674 --- /dev/null +++ b/raw_linux_test.go @@ -0,0 +1,232 @@ +/* + Copyright The containerd Authors. + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +*/ + +package fifo + +import ( + "bytes" + "context" + "io" + "io/ioutil" + "net" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestReadFrom(t *testing.T) { + dir, err := ioutil.TempDir("", t.Name()) + assert.NoError(t, err) + defer os.RemoveAll(dir) + + ctx := context.Background() + data := strings.Repeat("This is a test, this is only a test.", 1000) + + // For these test cases we only call ReadFrom and validate there is no error and the + // amouont of data it copied is what we put into it. + // The main test runner will validate the data is correct on each case. + cases := map[string]func(*testing.T, io.ReaderFrom){ + "regular file": func(t *testing.T, w io.ReaderFrom) { + f, err := os.OpenFile(filepath.Join(dir, "data"), os.O_RDWR|os.O_CREATE, 0600) + assert.NoError(t, err) + defer func() { + f.Close() + os.RemoveAll(f.Name()) + }() + + _, err = f.WriteString(data) + assert.NoError(t, err) + + _, err = f.Seek(0, io.SeekStart) + assert.NoError(t, err) + + n, err := w.ReadFrom(f) + w.(io.Closer).Close() + assert.NoError(t, err) + assert.Equal(t, n, int64(len(data))) + }, + "fifo": func(t *testing.T, w io.ReaderFrom) { + // Tests fifo<->fifo copy works. + buf := strings.NewReader(data) + + fifoW, err := OpenFifo(ctx, filepath.Join(dir, "fifo2"), os.O_RDWR|os.O_CREATE, 0600) + assert.NoError(t, err) + defer fifoW.Close() + + fifoR, err := OpenFifo(ctx, filepath.Join(dir, "fifo2"), os.O_RDONLY, 0600) + assert.NoError(t, err) + defer fifoR.Close() + + go func() { + io.Copy(fifoW, buf) + fifoW.Close() + }() + + n, err := w.ReadFrom(fifoR) + w.(io.Closer).Close() + assert.NoError(t, err) + assert.Equal(t, n, buf.Size()) + }, + "unix conn": func(t *testing.T, w io.ReaderFrom) { + sock := filepath.Join(dir, t.Name()) + err := os.MkdirAll(filepath.Dir(sock), 0755) + assert.NoError(t, err) + + l, err := net.Listen("unix", sock) + assert.NoError(t, err) + + go func() { + defer l.Close() + conn, err := l.Accept() + if err != nil { + return + } + defer conn.Close() + io.Copy(conn, strings.NewReader(data)) + }() + + conn, err := net.Dial("unix", filepath.Join(dir, t.Name())) + assert.NoError(t, err) + defer conn.Close() + + n, err := w.ReadFrom(conn) + w.(io.Closer).Close() + assert.NoError(t, err) + assert.Equal(t, n, int64(len(data))) + }, + "userspace": func(t *testing.T, w io.ReaderFrom) { + // Tests that copying from userspace explicitly works. + data := strings.NewReader(data) + n, err := w.ReadFrom(data) + w.(io.Closer).Close() + assert.NoError(t, err) + assert.Equal(t, n, data.Size()) + }, + "limited reader": func(t *testing.T, w io.ReaderFrom) { + // Makes sure we don't read too much from a file wrapped in a limited reader + // This is important because normally we'd just try to splice as much data as possible, + // If the user wraps with a LimitedReader, we still want to the benefits of splice but need + // to limit the splice call to the value limited in the LimitedReader. + f, err := os.OpenFile(filepath.Join(dir, "data"), os.O_RDWR|os.O_CREATE, 0600) + assert.NoError(t, err) + defer func() { + f.Close() + os.RemoveAll(f.Name()) + }() + + // Write data twice, will limit the reader to just one. + written, err := f.WriteString(data + data) + assert.NoError(t, err) + assert.Equal(t, written, len(data+data)) + + _, err = f.Seek(0, io.SeekStart) + assert.NoError(t, err) + + n, err := w.ReadFrom(io.LimitReader(f, int64(len(data)))) + w.(io.Closer).Close() + assert.NoError(t, err) + assert.Equal(t, n, int64(len(data))) + }, + } + + buf := bytes.NewBuffer(nil) + for name, testCase := range cases { + name := name + testCase := testCase + + t.Run(name, func(t *testing.T) { + buf.Reset() + pipeW, err := OpenFifo(ctx, filepath.Join(dir, "fifo"), os.O_RDWR|os.O_CREATE, 0600) + assert.NoError(t, err) + defer pipeW.Close() + + pipeR, err := OpenFifo(ctx, filepath.Join(dir, "fifo"), os.O_RDONLY, 0600) + assert.NoError(t, err) + defer pipeR.Close() + + done := make(chan struct{}) + go func() { + io.Copy(buf, pipeR) + close(done) + }() + + testCase(t, pipeW.(io.ReaderFrom)) + + select { + case <-done: + case <-time.After(10 * time.Second): + t.Error("timeout waiting for copy") + // Force copy to end and wait for it + // We have already failed the test above but there's still some value in seeing the results of the copy below. + pipeR.Close() + <-done + } + assert.Equal(t, len(data), buf.Len()) + assert.Equal(t, buf.String(), data) + }) + } +} + +func TestWriteTo(t *testing.T) { + dir, err := ioutil.TempDir("", t.Name()) + assert.NoError(t, err) + defer os.RemoveAll(dir) + + ctx := context.Background() + + pipeW, err := OpenFifo(ctx, filepath.Join(dir, "fifo"), os.O_RDWR|os.O_CREATE, 0600) + assert.NoError(t, err) + defer pipeW.Close() + + pipeR, err := OpenFifo(ctx, filepath.Join(dir, "fifo"), os.O_RDONLY, 0600) + assert.NoError(t, err) + defer pipeR.Close() + + f, err := os.OpenFile(filepath.Join(dir, "data"), os.O_RDWR|os.O_CREATE, 0600) + assert.NoError(t, err) + defer f.Close() + + data := strings.Repeat("This is a test, this is only a test.", 100) + + done := make(chan struct{}) + go func() { + io.Copy(pipeW, strings.NewReader(data)) + pipeW.Close() + close(done) + }() + + _, err = f.Seek(0, io.SeekStart) + assert.NoError(t, err) + + n, err := pipeR.(io.WriterTo).WriteTo(f) + pipeR.Close() + assert.NoError(t, err) + assert.Equal(t, int64(len(data)), n) + + <-done + + _, err = f.Seek(0, io.SeekStart) + assert.NoError(t, err) + + buf := bytes.NewBuffer(nil) + _, err = io.Copy(buf, f) + assert.NoError(t, err) + assert.Equal(t, buf.String(), data) +}