github.com/sagernet/gvisor@v0.0.0-20240428053021-e691de28565f/pkg/flipcall/io.go (about)

     1  // Copyright 2019 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package flipcall
    16  
    17  import (
    18  	"fmt"
    19  	"io"
    20  )
    21  
    22  // DatagramReader implements io.Reader by reading a datagram from an Endpoint's
    23  // packet window. Its use is optional; users that can use Endpoint.Data() more
    24  // efficiently are advised to do so.
    25  type DatagramReader struct {
    26  	ep  *Endpoint
    27  	off uint32
    28  	end uint32
    29  }
    30  
    31  // Init must be called on zero-value DatagramReaders before first use.
    32  //
    33  // Preconditions: dataLen is 0, or was returned by a previous call to
    34  // ep.RecvFirst() or ep.SendRecv().
    35  func (r *DatagramReader) Init(ep *Endpoint, dataLen uint32) {
    36  	r.ep = ep
    37  	r.Reset(dataLen)
    38  }
    39  
    40  // Reset causes r to begin reading a new datagram of the given length from the
    41  // associated Endpoint.
    42  //
    43  // Preconditions: dataLen is 0, or was returned by a previous call to the
    44  // associated Endpoint's RecvFirst() or SendRecv() methods.
    45  func (r *DatagramReader) Reset(dataLen uint32) {
    46  	if dataLen > r.ep.dataCap {
    47  		panic(fmt.Sprintf("invalid dataLen (%d) > ep.dataCap (%d)", dataLen, r.ep.dataCap))
    48  	}
    49  	r.off = 0
    50  	r.end = dataLen
    51  }
    52  
    53  // NewReader is a convenience function that returns an initialized
    54  // DatagramReader allocated on the heap.
    55  //
    56  // Preconditions: dataLen was returned by a previous call to ep.RecvFirst() or
    57  // ep.SendRecv().
    58  func (ep *Endpoint) NewReader(dataLen uint32) *DatagramReader {
    59  	r := &DatagramReader{}
    60  	r.Init(ep, dataLen)
    61  	return r
    62  }
    63  
    64  // Read implements io.Reader.Read.
    65  func (r *DatagramReader) Read(dst []byte) (int, error) {
    66  	n := copy(dst, r.ep.Data()[r.off:r.end])
    67  	r.off += uint32(n)
    68  	if r.off == r.end {
    69  		return n, io.EOF
    70  	}
    71  	return n, nil
    72  }
    73  
    74  // DatagramWriter implements io.Writer by writing a datagram to an Endpoint's
    75  // packet window. Its use is optional; users that can use Endpoint.Data() more
    76  // efficiently are advised to do so.
    77  type DatagramWriter struct {
    78  	ep  *Endpoint
    79  	off uint32
    80  }
    81  
    82  // Init must be called on zero-value DatagramWriters before first use.
    83  func (w *DatagramWriter) Init(ep *Endpoint) {
    84  	w.ep = ep
    85  }
    86  
    87  // Reset causes w to begin writing a new datagram to the associated Endpoint.
    88  func (w *DatagramWriter) Reset() {
    89  	w.off = 0
    90  }
    91  
    92  // NewWriter is a convenience function that returns an initialized
    93  // DatagramWriter allocated on the heap.
    94  func (ep *Endpoint) NewWriter() *DatagramWriter {
    95  	w := &DatagramWriter{}
    96  	w.Init(ep)
    97  	return w
    98  }
    99  
   100  // Write implements io.Writer.Write.
   101  func (w *DatagramWriter) Write(src []byte) (int, error) {
   102  	n := copy(w.ep.Data()[w.off:w.ep.dataCap], src)
   103  	w.off += uint32(n)
   104  	if n != len(src) {
   105  		return n, fmt.Errorf("datagram would exceed maximum size of %d bytes", w.ep.dataCap)
   106  	}
   107  	return n, nil
   108  }
   109  
   110  // Len returns the length of the written datagram.
   111  func (w *DatagramWriter) Len() uint32 {
   112  	return w.off
   113  }