stdcopy.go 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190
  1. package stdcopy
  2. import (
  3. "bytes"
  4. "encoding/binary"
  5. "errors"
  6. "fmt"
  7. "io"
  8. "sync"
  9. )
  10. // StdType is the type of standard stream
  11. // a writer can multiplex to.
  12. type StdType byte
  13. const (
  14. // Stdin represents standard input stream type.
  15. Stdin StdType = iota
  16. // Stdout represents standard output stream type.
  17. Stdout
  18. // Stderr represents standard error steam type.
  19. Stderr
  20. // Systemerr represents errors originating from the system that make it
  21. // into the the multiplexed stream.
  22. Systemerr
  23. stdWriterPrefixLen = 8
  24. stdWriterFdIndex = 0
  25. stdWriterSizeIndex = 4
  26. startingBufLen = 32*1024 + stdWriterPrefixLen + 1
  27. )
  28. var bufPool = &sync.Pool{New: func() interface{} { return bytes.NewBuffer(nil) }}
  29. // stdWriter is wrapper of io.Writer with extra customized info.
  30. type stdWriter struct {
  31. io.Writer
  32. prefix byte
  33. }
  34. // Write sends the buffer to the underneath writer.
  35. // It inserts the prefix header before the buffer,
  36. // so stdcopy.StdCopy knows where to multiplex the output.
  37. // It makes stdWriter to implement io.Writer.
  38. func (w *stdWriter) Write(p []byte) (n int, err error) {
  39. if w == nil || w.Writer == nil {
  40. return 0, errors.New("Writer not instantiated")
  41. }
  42. if p == nil {
  43. return 0, nil
  44. }
  45. header := [stdWriterPrefixLen]byte{stdWriterFdIndex: w.prefix}
  46. binary.BigEndian.PutUint32(header[stdWriterSizeIndex:], uint32(len(p)))
  47. buf := bufPool.Get().(*bytes.Buffer)
  48. buf.Write(header[:])
  49. buf.Write(p)
  50. n, err = w.Writer.Write(buf.Bytes())
  51. n -= stdWriterPrefixLen
  52. if n < 0 {
  53. n = 0
  54. }
  55. buf.Reset()
  56. bufPool.Put(buf)
  57. return
  58. }
  59. // NewStdWriter instantiates a new Writer.
  60. // Everything written to it will be encapsulated using a custom format,
  61. // and written to the underlying `w` stream.
  62. // This allows multiple write streams (e.g. stdout and stderr) to be muxed into a single connection.
  63. // `t` indicates the id of the stream to encapsulate.
  64. // It can be stdcopy.Stdin, stdcopy.Stdout, stdcopy.Stderr.
  65. func NewStdWriter(w io.Writer, t StdType) io.Writer {
  66. return &stdWriter{
  67. Writer: w,
  68. prefix: byte(t),
  69. }
  70. }
  71. // StdCopy is a modified version of io.Copy.
  72. //
  73. // StdCopy will demultiplex `src`, assuming that it contains two streams,
  74. // previously multiplexed together using a StdWriter instance.
  75. // As it reads from `src`, StdCopy will write to `dstout` and `dsterr`.
  76. //
  77. // StdCopy will read until it hits EOF on `src`. It will then return a nil error.
  78. // In other words: if `err` is non nil, it indicates a real underlying error.
  79. //
  80. // `written` will hold the total number of bytes written to `dstout` and `dsterr`.
  81. func StdCopy(dstout, dsterr io.Writer, src io.Reader) (written int64, err error) {
  82. var (
  83. buf = make([]byte, startingBufLen)
  84. bufLen = len(buf)
  85. nr, nw int
  86. er, ew error
  87. out io.Writer
  88. frameSize int
  89. )
  90. for {
  91. // Make sure we have at least a full header
  92. for nr < stdWriterPrefixLen {
  93. var nr2 int
  94. nr2, er = src.Read(buf[nr:])
  95. nr += nr2
  96. if er == io.EOF {
  97. if nr < stdWriterPrefixLen {
  98. return written, nil
  99. }
  100. break
  101. }
  102. if er != nil {
  103. return 0, er
  104. }
  105. }
  106. stream := StdType(buf[stdWriterFdIndex])
  107. // Check the first byte to know where to write
  108. switch stream {
  109. case Stdin:
  110. fallthrough
  111. case Stdout:
  112. // Write on stdout
  113. out = dstout
  114. case Stderr:
  115. // Write on stderr
  116. out = dsterr
  117. case Systemerr:
  118. // If we're on Systemerr, we won't write anywhere.
  119. // NB: if this code changes later, make sure you don't try to write
  120. // to outstream if Systemerr is the stream
  121. out = nil
  122. default:
  123. return 0, fmt.Errorf("Unrecognized input header: %d", buf[stdWriterFdIndex])
  124. }
  125. // Retrieve the size of the frame
  126. frameSize = int(binary.BigEndian.Uint32(buf[stdWriterSizeIndex : stdWriterSizeIndex+4]))
  127. // Check if the buffer is big enough to read the frame.
  128. // Extend it if necessary.
  129. if frameSize+stdWriterPrefixLen > bufLen {
  130. buf = append(buf, make([]byte, frameSize+stdWriterPrefixLen-bufLen+1)...)
  131. bufLen = len(buf)
  132. }
  133. // While the amount of bytes read is less than the size of the frame + header, we keep reading
  134. for nr < frameSize+stdWriterPrefixLen {
  135. var nr2 int
  136. nr2, er = src.Read(buf[nr:])
  137. nr += nr2
  138. if er == io.EOF {
  139. if nr < frameSize+stdWriterPrefixLen {
  140. return written, nil
  141. }
  142. break
  143. }
  144. if er != nil {
  145. return 0, er
  146. }
  147. }
  148. // we might have an error from the source mixed up in our multiplexed
  149. // stream. if we do, return it.
  150. if stream == Systemerr {
  151. return written, fmt.Errorf("error from daemon in stream: %s", string(buf[stdWriterPrefixLen:frameSize+stdWriterPrefixLen]))
  152. }
  153. // Write the retrieved frame (without header)
  154. nw, ew = out.Write(buf[stdWriterPrefixLen : frameSize+stdWriterPrefixLen])
  155. if ew != nil {
  156. return 0, ew
  157. }
  158. // If the frame has not been fully written: error
  159. if nw != frameSize {
  160. return 0, io.ErrShortWrite
  161. }
  162. written += int64(nw)
  163. // Move the rest of the buffer to the beginning
  164. copy(buf, buf[frameSize+stdWriterPrefixLen:])
  165. // Move the index
  166. nr -= frameSize + stdWriterPrefixLen
  167. }
  168. }