github.com/mysteriumnetwork/node@v0.0.0-20240516044423-365054f76801/services/wireguard/endpoint/netstack-provider/shaper.go (about)

     1  /*
     2   * Copyright (C) 2022 The "MysteriumNetwork/node" Authors.
     3   *
     4   * This program is free software: you can redistribute it and/or modify
     5   * it under the terms of the GNU General Public License as published by
     6   * the Free Software Foundation, either version 3 of the License, or
     7   * (at your option) any later version.
     8   *
     9   * This program is distributed in the hope that it will be useful,
    10   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    11   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    12   * GNU General Public License for more details.
    13   *
    14   * You should have received a copy of the GNU General Public License
    15   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    16   */
    17  
    18  package netstack_provider
    19  
    20  import (
    21  	"context"
    22  	"io"
    23  
    24  	"golang.org/x/time/rate"
    25  )
    26  
    27  type Reader struct {
    28  	r       io.Reader
    29  	limiter *rate.Limiter
    30  	ctx     context.Context
    31  }
    32  
    33  // NewReader returns a reader that implements io.Reader with rate limiting.
    34  func NewReader(r io.Reader, limiter *rate.Limiter) *Reader {
    35  	return &Reader{
    36  		r:       r,
    37  		limiter: limiter,
    38  		ctx:     context.Background(),
    39  	}
    40  }
    41  
    42  // NewReaderWithContext returns a reader that implements io.Reader with rate limiting.
    43  func NewReaderWithContext(r io.Reader, ctx context.Context) *Reader {
    44  	return &Reader{
    45  		r:   r,
    46  		ctx: ctx,
    47  	}
    48  }
    49  
    50  // Read reads bytes into p.
    51  func (s *Reader) Read(p []byte) (int, error) {
    52  	if s.limiter == nil {
    53  		return s.r.Read(p)
    54  	}
    55  	n, err := s.r.Read(p)
    56  	if err != nil {
    57  		return n, err
    58  	}
    59  	if err := s.limiter.WaitN(s.ctx, n); err != nil {
    60  		return n, err
    61  	}
    62  	return n, nil
    63  }