| 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 |
| } |
| } |