123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190 |
- package stdcopy
- import (
- "bytes"
- "encoding/binary"
- "errors"
- "fmt"
- "io"
- "sync"
- )
- // 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
- // Systemerr represents errors originating from the system that make it
- // into the the multiplexed stream.
- Systemerr
- stdWriterPrefixLen = 8
- stdWriterFdIndex = 0
- stdWriterSizeIndex = 4
- startingBufLen = 32*1024 + stdWriterPrefixLen + 1
- )
- var bufPool = &sync.Pool{New: func() interface{} { return bytes.NewBuffer(nil) }}
- // 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 inserts 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(p []byte) (n int, err error) {
- if w == nil || w.Writer == nil {
- return 0, errors.New("Writer not instantiated")
- }
- if p == nil {
- return 0, nil
- }
- header := [stdWriterPrefixLen]byte{stdWriterFdIndex: w.prefix}
- binary.BigEndian.PutUint32(header[stdWriterSizeIndex:], uint32(len(p)))
- buf := bufPool.Get().(*bytes.Buffer)
- buf.Write(header[:])
- buf.Write(p)
- n, err = w.Writer.Write(buf.Bytes())
- n -= stdWriterPrefixLen
- if n < 0 {
- n = 0
- }
- buf.Reset()
- bufPool.Put(buf)
- 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 {
- return written, nil
- }
- break
- }
- if er != nil {
- return 0, er
- }
- }
- stream := StdType(buf[stdWriterFdIndex])
- // Check the first byte to know where to write
- switch stream {
- case Stdin:
- fallthrough
- case Stdout:
- // Write on stdout
- out = dstout
- case Stderr:
- // Write on stderr
- out = dsterr
- case Systemerr:
- // If we're on Systemerr, we won't write anywhere.
- // NB: if this code changes later, make sure you don't try to write
- // to outstream if Systemerr is the stream
- out = nil
- default:
- 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]))
- // Check if the buffer is big enough to read the frame.
- // Extend it if necessary.
- if frameSize+stdWriterPrefixLen > bufLen {
- 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 {
- return written, nil
- }
- break
- }
- if er != nil {
- return 0, er
- }
- }
- // we might have an error from the source mixed up in our multiplexed
- // stream. if we do, return it.
- if stream == Systemerr {
- return written, fmt.Errorf("error from daemon in stream: %s", string(buf[stdWriterPrefixLen:frameSize+stdWriterPrefixLen]))
- }
- // Write the retrieved frame (without header)
- nw, ew = out.Write(buf[stdWriterPrefixLen : frameSize+stdWriterPrefixLen])
- if ew != nil {
- return 0, ew
- }
- // If the frame has not been fully written: error
- if 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
- }
- }
|