Initial commit

vendor/github.com/hyperhq/hypercli/pkg/stdcopy/stdcopy.go

@@ -0,0 +1,178 @@
+package stdcopy
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"io"
+
+	"github.com/Sirupsen/logrus"
+)
+
+// StdType is the type of standard stream
+// a writer can multiplex to.
+type StdType byte
+
+const (
+	// Stdin represents standard input stream type.
+	Stdin StdType = iota
+	// Stdout represents standard output stream type.
+	Stdout
+	// Stderr represents standard error steam type.
+	Stderr
+
+	stdWriterPrefixLen = 8
+	stdWriterFdIndex   = 0
+	stdWriterSizeIndex = 4
+
+	startingBufLen = 32*1024 + stdWriterPrefixLen + 1
+)
+
+// stdWriter is wrapper of io.Writer with extra customized info.
+type stdWriter struct {
+	io.Writer
+	prefix byte
+}
+
+// Write sends the buffer to the underneath writer.
+// It insert the prefix header before the buffer,
+// so stdcopy.StdCopy knows where to multiplex the output.
+// It makes stdWriter to implement io.Writer.
+func (w *stdWriter) Write(buf []byte) (n int, err error) {
+	if w == nil || w.Writer == nil {
+		return 0, errors.New("Writer not instantiated")
+	}
+	if buf == nil {
+		return 0, nil
+	}
+
+	header := [stdWriterPrefixLen]byte{stdWriterFdIndex: w.prefix}
+	binary.BigEndian.PutUint32(header[stdWriterSizeIndex:], uint32(len(buf)))
+
+	line := append(header[:], buf...)
+
+	n, err = w.Writer.Write(line)
+	n -= stdWriterPrefixLen
+
+	if n < 0 {
+		n = 0
+	}
+	return
+}
+
+// NewStdWriter instantiates a new Writer.
+// Everything written to it will be encapsulated using a custom format,
+// and written to the underlying `w` stream.
+// This allows multiple write streams (e.g. stdout and stderr) to be muxed into a single connection.
+// `t` indicates the id of the stream to encapsulate.
+// It can be stdcopy.Stdin, stdcopy.Stdout, stdcopy.Stderr.
+func NewStdWriter(w io.Writer, t StdType) io.Writer {
+	return &stdWriter{
+		Writer: w,
+		prefix: byte(t),
+	}
+}
+
+// StdCopy is a modified version of io.Copy.
+//
+// StdCopy will demultiplex `src`, assuming that it contains two streams,
+// previously multiplexed together using a StdWriter instance.
+// As it reads from `src`, StdCopy will write to `dstout` and `dsterr`.
+//
+// StdCopy will read until it hits EOF on `src`. It will then return a nil error.
+// In other words: if `err` is non nil, it indicates a real underlying error.
+//
+// `written` will hold the total number of bytes written to `dstout` and `dsterr`.
+func StdCopy(dstout, dsterr io.Writer, src io.Reader) (written int64, err error) {
+	var (
+		buf       = make([]byte, startingBufLen)
+		bufLen    = len(buf)
+		nr, nw    int
+		er, ew    error
+		out       io.Writer
+		frameSize int
+	)
+
+	for {
+		// Make sure we have at least a full header
+		for nr < stdWriterPrefixLen {
+			var nr2 int
+			nr2, er = src.Read(buf[nr:])
+			nr += nr2
+			if er == io.EOF {
+				if nr < stdWriterPrefixLen {
+					logrus.Debugf("Corrupted prefix: %v", buf[:nr])
+					return written, nil
+				}
+				break
+			}
+			if er != nil {
+				logrus.Debugf("Error reading header: %s", er)
+				return 0, er
+			}
+		}
+
+		// Check the first byte to know where to write
+		switch StdType(buf[stdWriterFdIndex]) {
+		case Stdin:
+			fallthrough
+		case Stdout:
+			// Write on stdout
+			out = dstout
+		case Stderr:
+			// Write on stderr
+			out = dsterr
+		default:
+			logrus.Debugf("Error selecting output fd: (%d)", buf[stdWriterFdIndex])
+			return 0, fmt.Errorf("Unrecognized input header: %d", buf[stdWriterFdIndex])
+		}
+
+		// Retrieve the size of the frame
+		frameSize = int(binary.BigEndian.Uint32(buf[stdWriterSizeIndex : stdWriterSizeIndex+4]))
+		logrus.Debugf("framesize: %d", frameSize)
+
+		// Check if the buffer is big enough to read the frame.
+		// Extend it if necessary.
+		if frameSize+stdWriterPrefixLen > bufLen {
+			logrus.Debugf("Extending buffer cap by %d (was %d)", frameSize+stdWriterPrefixLen-bufLen+1, len(buf))
+			buf = append(buf, make([]byte, frameSize+stdWriterPrefixLen-bufLen+1)...)
+			bufLen = len(buf)
+		}
+
+		// While the amount of bytes read is less than the size of the frame + header, we keep reading
+		for nr < frameSize+stdWriterPrefixLen {
+			var nr2 int
+			nr2, er = src.Read(buf[nr:])
+			nr += nr2
+			if er == io.EOF {
+				if nr < frameSize+stdWriterPrefixLen {
+					logrus.Debugf("Corrupted frame: %v", buf[stdWriterPrefixLen:nr])
+					return written, nil
+				}
+				break
+			}
+			if er != nil {
+				logrus.Debugf("Error reading frame: %s", er)
+				return 0, er
+			}
+		}
+
+		// Write the retrieved frame (without header)
+		nw, ew = out.Write(buf[stdWriterPrefixLen : frameSize+stdWriterPrefixLen])
+		if ew != nil {
+			logrus.Debugf("Error writing frame: %s", ew)
+			return 0, ew
+		}
+		// If the frame has not been fully written: error
+		if nw != frameSize {
+			logrus.Debugf("Error Short Write: (%d on %d)", nw, frameSize)
+			return 0, io.ErrShortWrite
+		}
+		written += int64(nw)
+
+		// Move the rest of the buffer to the beginning
+		copy(buf, buf[frameSize+stdWriterPrefixLen:])
+		// Move the index
+		nr -= frameSize + stdWriterPrefixLen
+	}
+}

vendor/github.com/hyperhq/hypercli/pkg/stdcopy/stdcopy_test.go

@@ -0,0 +1,260 @@
+package stdcopy
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"io/ioutil"
+	"strings"
+	"testing"
+)
+
+func TestNewStdWriter(t *testing.T) {
+	writer := NewStdWriter(ioutil.Discard, Stdout)
+	if writer == nil {
+		t.Fatalf("NewStdWriter with an invalid StdType should not return nil.")
+	}
+}
+
+func TestWriteWithUnitializedStdWriter(t *testing.T) {
+	writer := stdWriter{
+		Writer: nil,
+		prefix: byte(Stdout),
+	}
+	n, err := writer.Write([]byte("Something here"))
+	if n != 0 || err == nil {
+		t.Fatalf("Should fail when given an uncomplete or uninitialized StdWriter")
+	}
+}
+
+func TestWriteWithNilBytes(t *testing.T) {
+	writer := NewStdWriter(ioutil.Discard, Stdout)
+	n, err := writer.Write(nil)
+	if err != nil {
+		t.Fatalf("Shouldn't have fail when given no data")
+	}
+	if n > 0 {
+		t.Fatalf("Write should have written 0 byte, but has written %d", n)
+	}
+}
+
+func TestWrite(t *testing.T) {
+	writer := NewStdWriter(ioutil.Discard, Stdout)
+	data := []byte("Test StdWrite.Write")
+	n, err := writer.Write(data)
+	if err != nil {
+		t.Fatalf("Error while writing with StdWrite")
+	}
+	if n != len(data) {
+		t.Fatalf("Write should have written %d byte but wrote %d.", len(data), n)
+	}
+}
+
+type errWriter struct {
+	n   int
+	err error
+}
+
+func (f *errWriter) Write(buf []byte) (int, error) {
+	return f.n, f.err
+}
+
+func TestWriteWithWriterError(t *testing.T) {
+	expectedError := errors.New("expected")
+	expectedReturnedBytes := 10
+	writer := NewStdWriter(&errWriter{
+		n:   stdWriterPrefixLen + expectedReturnedBytes,
+		err: expectedError}, Stdout)
+	data := []byte("This won't get written, sigh")
+	n, err := writer.Write(data)
+	if err != expectedError {
+		t.Fatalf("Didn't get expected error.")
+	}
+	if n != expectedReturnedBytes {
+		t.Fatalf("Didn't get expected writen bytes %d, got %d.",
+			expectedReturnedBytes, n)
+	}
+}
+
+func TestWriteDoesNotReturnNegativeWrittenBytes(t *testing.T) {
+	writer := NewStdWriter(&errWriter{n: -1}, Stdout)
+	data := []byte("This won't get written, sigh")
+	actual, _ := writer.Write(data)
+	if actual != 0 {
+		t.Fatalf("Expected returned written bytes equal to 0, got %d", actual)
+	}
+}
+
+func getSrcBuffer(stdOutBytes, stdErrBytes []byte) (buffer *bytes.Buffer, err error) {
+	buffer = new(bytes.Buffer)
+	dstOut := NewStdWriter(buffer, Stdout)
+	_, err = dstOut.Write(stdOutBytes)
+	if err != nil {
+		return
+	}
+	dstErr := NewStdWriter(buffer, Stderr)
+	_, err = dstErr.Write(stdErrBytes)
+	return
+}
+
+func TestStdCopyWriteAndRead(t *testing.T) {
+	stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
+	stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
+	buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+	written, err := StdCopy(ioutil.Discard, ioutil.Discard, buffer)
+	if err != nil {
+		t.Fatal(err)
+	}
+	expectedTotalWritten := len(stdOutBytes) + len(stdErrBytes)
+	if written != int64(expectedTotalWritten) {
+		t.Fatalf("Expected to have total of %d bytes written, got %d", expectedTotalWritten, written)
+	}
+}
+
+type customReader struct {
+	n            int
+	err          error
+	totalCalls   int
+	correctCalls int
+	src          *bytes.Buffer
+}
+
+func (f *customReader) Read(buf []byte) (int, error) {
+	f.totalCalls++
+	if f.totalCalls <= f.correctCalls {
+		return f.src.Read(buf)
+	}
+	return f.n, f.err
+}
+
+func TestStdCopyReturnsErrorReadingHeader(t *testing.T) {
+	expectedError := errors.New("error")
+	reader := &customReader{
+		err: expectedError}
+	written, err := StdCopy(ioutil.Discard, ioutil.Discard, reader)
+	if written != 0 {
+		t.Fatalf("Expected 0 bytes read, got %d", written)
+	}
+	if err != expectedError {
+		t.Fatalf("Didn't get expected error")
+	}
+}
+
+func TestStdCopyReturnsErrorReadingFrame(t *testing.T) {
+	expectedError := errors.New("error")
+	stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
+	stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
+	buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+	reader := &customReader{
+		correctCalls: 1,
+		n:            stdWriterPrefixLen + 1,
+		err:          expectedError,
+		src:          buffer}
+	written, err := StdCopy(ioutil.Discard, ioutil.Discard, reader)
+	if written != 0 {
+		t.Fatalf("Expected 0 bytes read, got %d", written)
+	}
+	if err != expectedError {
+		t.Fatalf("Didn't get expected error")
+	}
+}
+
+func TestStdCopyDetectsCorruptedFrame(t *testing.T) {
+	stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
+	stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
+	buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+	reader := &customReader{
+		correctCalls: 1,
+		n:            stdWriterPrefixLen + 1,
+		err:          io.EOF,
+		src:          buffer}
+	written, err := StdCopy(ioutil.Discard, ioutil.Discard, reader)
+	if written != startingBufLen {
+		t.Fatalf("Expected %d bytes read, got %d", startingBufLen, written)
+	}
+	if err != nil {
+		t.Fatal("Didn't get nil error")
+	}
+}
+
+func TestStdCopyWithInvalidInputHeader(t *testing.T) {
+	dstOut := NewStdWriter(ioutil.Discard, Stdout)
+	dstErr := NewStdWriter(ioutil.Discard, Stderr)
+	src := strings.NewReader("Invalid input")
+	_, err := StdCopy(dstOut, dstErr, src)
+	if err == nil {
+		t.Fatal("StdCopy with invalid input header should fail.")
+	}
+}
+
+func TestStdCopyWithCorruptedPrefix(t *testing.T) {
+	data := []byte{0x01, 0x02, 0x03}
+	src := bytes.NewReader(data)
+	written, err := StdCopy(nil, nil, src)
+	if err != nil {
+		t.Fatalf("StdCopy should not return an error with corrupted prefix.")
+	}
+	if written != 0 {
+		t.Fatalf("StdCopy should have written 0, but has written %d", written)
+	}
+}
+
+func TestStdCopyReturnsWriteErrors(t *testing.T) {
+	stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
+	stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
+	buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+	expectedError := errors.New("expected")
+
+	dstOut := &errWriter{err: expectedError}
+
+	written, err := StdCopy(dstOut, ioutil.Discard, buffer)
+	if written != 0 {
+		t.Fatalf("StdCopy should have written 0, but has written %d", written)
+	}
+	if err != expectedError {
+		t.Fatalf("Didn't get expected error, got %v", err)
+	}
+}
+
+func TestStdCopyDetectsNotFullyWrittenFrames(t *testing.T) {
+	stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
+	stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
+	buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+	dstOut := &errWriter{n: startingBufLen - 10}
+
+	written, err := StdCopy(dstOut, ioutil.Discard, buffer)
+	if written != 0 {
+		t.Fatalf("StdCopy should have return 0 written bytes, but returned %d", written)
+	}
+	if err != io.ErrShortWrite {
+		t.Fatalf("Didn't get expected io.ErrShortWrite error")
+	}
+}
+
+func BenchmarkWrite(b *testing.B) {
+	w := NewStdWriter(ioutil.Discard, Stdout)
+	data := []byte("Test line for testing stdwriter performance\n")
+	data = bytes.Repeat(data, 100)
+	b.SetBytes(int64(len(data)))
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if _, err := w.Write(data); err != nil {
+			b.Fatal(err)
+		}
+	}
+}