github.com/qichengzx/mattermost-server@v4.5.1-0.20180604164826-2c75247c97d0+incompatible/plugin/rpcplugin/io.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See License.txt for license information.
     3  
     4  package rpcplugin
     5  
     6  import (
     7  	"bufio"
     8  	"encoding/binary"
     9  	"io"
    10  )
    11  
    12  type rwc struct {
    13  	io.ReadCloser
    14  	io.WriteCloser
    15  }
    16  
    17  func (rwc *rwc) Close() (err error) {
    18  	err = rwc.WriteCloser.Close()
    19  	if rerr := rwc.ReadCloser.Close(); err == nil {
    20  		err = rerr
    21  	}
    22  	return
    23  }
    24  
    25  func NewReadWriteCloser(r io.ReadCloser, w io.WriteCloser) io.ReadWriteCloser {
    26  	return &rwc{r, w}
    27  }
    28  
    29  type RemoteIOReader struct {
    30  	conn io.ReadWriteCloser
    31  }
    32  
    33  func (r *RemoteIOReader) Read(b []byte) (int, error) {
    34  	var buf [10]byte
    35  	n := binary.PutVarint(buf[:], int64(len(b)))
    36  	if _, err := r.conn.Write(buf[:n]); err != nil {
    37  		return 0, err
    38  	}
    39  	return r.conn.Read(b)
    40  }
    41  
    42  func (r *RemoteIOReader) Close() error {
    43  	return r.conn.Close()
    44  }
    45  
    46  func ConnectIOReader(conn io.ReadWriteCloser) io.ReadCloser {
    47  	return &RemoteIOReader{conn}
    48  }
    49  
    50  func ServeIOReader(r io.Reader, conn io.ReadWriteCloser) {
    51  	cr := bufio.NewReader(conn)
    52  	defer conn.Close()
    53  	buf := make([]byte, 32*1024)
    54  	for {
    55  		n, err := binary.ReadVarint(cr)
    56  		if err != nil {
    57  			break
    58  		}
    59  		if written, err := io.CopyBuffer(conn, io.LimitReader(r, n), buf); err != nil || written < n {
    60  			break
    61  		}
    62  	}
    63  }