github.com/ttpreport/gvisor-ligolo@v0.0.0-20240123134145-a858404967ba/pkg/tcpip/ports/flags.go (about)

     1  // Copyright 2021 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 ports
    16  
    17  // Flags represents the type of port reservation.
    18  //
    19  // +stateify savable
    20  type Flags struct {
    21  	// MostRecent represents UDP SO_REUSEADDR.
    22  	MostRecent bool
    23  
    24  	// LoadBalanced indicates SO_REUSEPORT.
    25  	//
    26  	// LoadBalanced takes precidence over MostRecent.
    27  	LoadBalanced bool
    28  
    29  	// TupleOnly represents TCP SO_REUSEADDR.
    30  	TupleOnly bool
    31  }
    32  
    33  // Bits converts the Flags to their bitset form.
    34  func (f Flags) Bits() BitFlags {
    35  	var rf BitFlags
    36  	if f.MostRecent {
    37  		rf |= MostRecentFlag
    38  	}
    39  	if f.LoadBalanced {
    40  		rf |= LoadBalancedFlag
    41  	}
    42  	if f.TupleOnly {
    43  		rf |= TupleOnlyFlag
    44  	}
    45  	return rf
    46  }
    47  
    48  // Effective returns the effective behavior of a flag config.
    49  func (f Flags) Effective() Flags {
    50  	e := f
    51  	if e.LoadBalanced && e.MostRecent {
    52  		e.MostRecent = false
    53  	}
    54  	return e
    55  }
    56  
    57  // BitFlags is a bitset representation of Flags.
    58  type BitFlags uint32
    59  
    60  const (
    61  	// MostRecentFlag represents Flags.MostRecent.
    62  	MostRecentFlag BitFlags = 1 << iota
    63  
    64  	// LoadBalancedFlag represents Flags.LoadBalanced.
    65  	LoadBalancedFlag
    66  
    67  	// TupleOnlyFlag represents Flags.TupleOnly.
    68  	TupleOnlyFlag
    69  
    70  	// nextFlag is the value that the next added flag will have.
    71  	//
    72  	// It is used to calculate FlagMask below. It is also the number of
    73  	// valid flag states.
    74  	nextFlag
    75  
    76  	// FlagMask is a bit mask for BitFlags.
    77  	FlagMask = nextFlag - 1
    78  
    79  	// MultiBindFlagMask contains the flags that allow binding the same
    80  	// tuple multiple times.
    81  	MultiBindFlagMask = MostRecentFlag | LoadBalancedFlag
    82  )
    83  
    84  // ToFlags converts the bitset into a Flags struct.
    85  func (f BitFlags) ToFlags() Flags {
    86  	return Flags{
    87  		MostRecent:   f&MostRecentFlag != 0,
    88  		LoadBalanced: f&LoadBalancedFlag != 0,
    89  		TupleOnly:    f&TupleOnlyFlag != 0,
    90  	}
    91  }
    92  
    93  // FlagCounter counts how many references each flag combination has.
    94  type FlagCounter struct {
    95  	// refs stores the count for each possible flag combination, (0 though
    96  	// FlagMask).
    97  	refs [nextFlag]int
    98  }
    99  
   100  // AddRef increases the reference count for a specific flag combination.
   101  func (c *FlagCounter) AddRef(flags BitFlags) {
   102  	c.refs[flags]++
   103  }
   104  
   105  // DropRef decreases the reference count for a specific flag combination.
   106  func (c *FlagCounter) DropRef(flags BitFlags) {
   107  	c.refs[flags]--
   108  }
   109  
   110  // TotalRefs calculates the total number of references for all flag
   111  // combinations.
   112  func (c FlagCounter) TotalRefs() int {
   113  	var total int
   114  	for _, r := range c.refs {
   115  		total += r
   116  	}
   117  	return total
   118  }
   119  
   120  // FlagRefs returns the number of references with all specified flags.
   121  func (c FlagCounter) FlagRefs(flags BitFlags) int {
   122  	var total int
   123  	for i, r := range c.refs {
   124  		if BitFlags(i)&flags == flags {
   125  			total += r
   126  		}
   127  	}
   128  	return total
   129  }
   130  
   131  // AllRefsHave returns if all references have all specified flags.
   132  func (c FlagCounter) AllRefsHave(flags BitFlags) bool {
   133  	for i, r := range c.refs {
   134  		if BitFlags(i)&flags != flags && r > 0 {
   135  			return false
   136  		}
   137  	}
   138  	return true
   139  }
   140  
   141  // SharedFlags returns the set of flags shared by all references.
   142  func (c FlagCounter) SharedFlags() BitFlags {
   143  	intersection := FlagMask
   144  	for i, r := range c.refs {
   145  		if r > 0 {
   146  			intersection &= BitFlags(i)
   147  		}
   148  	}
   149  	return intersection
   150  }