github.com/nicocha30/gvisor-ligolo@v0.0.0-20230726075806-989fa2c0a413/pkg/tcpip/transport/tcp/segment_queue.go (about)

     1  // Copyright 2018 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 tcp
    16  
    17  import (
    18  	"github.com/nicocha30/gvisor-ligolo/pkg/sync"
    19  )
    20  
    21  // segmentQueue is a bounded, thread-safe queue of TCP segments.
    22  //
    23  // +stateify savable
    24  type segmentQueue struct {
    25  	mu     sync.Mutex  `state:"nosave"`
    26  	list   segmentList `state:"wait"`
    27  	ep     *endpoint
    28  	frozen bool
    29  }
    30  
    31  // emptyLocked determines if the queue is empty.
    32  // Preconditions: q.mu must be held.
    33  func (q *segmentQueue) emptyLocked() bool {
    34  	return q.list.Empty()
    35  }
    36  
    37  // empty determines if the queue is empty.
    38  func (q *segmentQueue) empty() bool {
    39  	q.mu.Lock()
    40  	defer q.mu.Unlock()
    41  	return q.emptyLocked()
    42  }
    43  
    44  // enqueue adds the given segment to the queue.
    45  //
    46  // Returns true when the segment is successfully added to the queue, in which
    47  // case ownership of the reference is transferred to the queue. And returns
    48  // false if the queue is full, in which case ownership is retained by the
    49  // caller.
    50  func (q *segmentQueue) enqueue(s *segment) bool {
    51  	// q.ep.receiveBufferParams() must be called without holding q.mu to
    52  	// avoid lock order inversion.
    53  	bufSz := q.ep.ops.GetReceiveBufferSize()
    54  	used := q.ep.receiveMemUsed()
    55  
    56  	q.mu.Lock()
    57  	defer q.mu.Unlock()
    58  
    59  	// Allow zero sized segments (ACK/FIN/RSTs etc even if the segment queue
    60  	// is currently full).
    61  	allow := (used <= int(bufSz) || s.payloadSize() == 0) && !q.frozen
    62  
    63  	if allow {
    64  		s.IncRef()
    65  		q.list.PushBack(s)
    66  		// Set the owner now that the endpoint owns the segment.
    67  		s.setOwner(q.ep, recvQ)
    68  	}
    69  
    70  	return allow
    71  }
    72  
    73  // dequeue removes and returns the next segment from queue, if one exists.
    74  // Ownership is transferred to the caller, who is responsible for decrementing
    75  // the ref count when done.
    76  func (q *segmentQueue) dequeue() *segment {
    77  	q.mu.Lock()
    78  	defer q.mu.Unlock()
    79  
    80  	s := q.list.Front()
    81  	if s != nil {
    82  		q.list.Remove(s)
    83  	}
    84  
    85  	return s
    86  }
    87  
    88  // freeze prevents any more segments from being added to the queue. i.e all
    89  // future segmentQueue.enqueue will return false and not add the segment to the
    90  // queue till the queue is unfroze with a corresponding segmentQueue.thaw call.
    91  func (q *segmentQueue) freeze() {
    92  	q.mu.Lock()
    93  	defer q.mu.Unlock()
    94  	q.frozen = true
    95  }
    96  
    97  // thaw unfreezes a previously frozen queue using segmentQueue.freeze() and
    98  // allows new segments to be queued again.
    99  func (q *segmentQueue) thaw() {
   100  	q.mu.Lock()
   101  	defer q.mu.Unlock()
   102  	q.frozen = false
   103  }