github.com/gogo/protobuf@v1.3.2/io/full.go (about)

     1  // Protocol Buffers for Go with Gadgets
     2  //
     3  // Copyright (c) 2013, The GoGo Authors. All rights reserved.
     4  // http://github.com/gogo/protobuf
     5  //
     6  // Redistribution and use in source and binary forms, with or without
     7  // modification, are permitted provided that the following conditions are
     8  // met:
     9  //
    10  //     * Redistributions of source code must retain the above copyright
    11  // notice, this list of conditions and the following disclaimer.
    12  //     * Redistributions in binary form must reproduce the above
    13  // copyright notice, this list of conditions and the following disclaimer
    14  // in the documentation and/or other materials provided with the
    15  // distribution.
    16  //
    17  // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
    18  // "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
    19  // LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
    20  // A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
    21  // OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
    22  // SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
    23  // LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
    24  // DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
    25  // THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
    26  // (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
    27  // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
    28  
    29  package io
    30  
    31  import (
    32  	"github.com/gogo/protobuf/proto"
    33  	"io"
    34  )
    35  
    36  func NewFullWriter(w io.Writer) WriteCloser {
    37  	return &fullWriter{w, nil}
    38  }
    39  
    40  type fullWriter struct {
    41  	w      io.Writer
    42  	buffer []byte
    43  }
    44  
    45  func (this *fullWriter) WriteMsg(msg proto.Message) (err error) {
    46  	var data []byte
    47  	if m, ok := msg.(marshaler); ok {
    48  		n, ok := getSize(m)
    49  		if !ok {
    50  			data, err = proto.Marshal(msg)
    51  			if err != nil {
    52  				return err
    53  			}
    54  		}
    55  		if n >= len(this.buffer) {
    56  			this.buffer = make([]byte, n)
    57  		}
    58  		_, err = m.MarshalTo(this.buffer)
    59  		if err != nil {
    60  			return err
    61  		}
    62  		data = this.buffer[:n]
    63  	} else {
    64  		data, err = proto.Marshal(msg)
    65  		if err != nil {
    66  			return err
    67  		}
    68  	}
    69  	_, err = this.w.Write(data)
    70  	return err
    71  }
    72  
    73  func (this *fullWriter) Close() error {
    74  	if closer, ok := this.w.(io.Closer); ok {
    75  		return closer.Close()
    76  	}
    77  	return nil
    78  }
    79  
    80  type fullReader struct {
    81  	r   io.Reader
    82  	buf []byte
    83  }
    84  
    85  func NewFullReader(r io.Reader, maxSize int) ReadCloser {
    86  	return &fullReader{r, make([]byte, maxSize)}
    87  }
    88  
    89  func (this *fullReader) ReadMsg(msg proto.Message) error {
    90  	length, err := this.r.Read(this.buf)
    91  	if err != nil {
    92  		return err
    93  	}
    94  	return proto.Unmarshal(this.buf[:length], msg)
    95  }
    96  
    97  func (this *fullReader) Close() error {
    98  	if closer, ok := this.r.(io.Closer); ok {
    99  		return closer.Close()
   100  	}
   101  	return nil
   102  }