github.com/projectdiscovery/nuclei/v2@v2.9.15/pkg/protocols/http/race/syncedreadcloser.go (about)

     1  package race
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"time"
     7  )
     8  
     9  // SyncedReadCloser is compatible with io.ReadSeeker and performs
    10  // gate-based synced writes to enable race condition testing.
    11  type SyncedReadCloser struct {
    12  	data           []byte
    13  	p              int64
    14  	length         int64
    15  	openGate       chan struct{}
    16  	enableBlocking bool
    17  }
    18  
    19  // NewSyncedReadCloser creates a new SyncedReadCloser instance.
    20  func NewSyncedReadCloser(r io.ReadCloser) *SyncedReadCloser {
    21  	var (
    22  		s   SyncedReadCloser
    23  		err error
    24  	)
    25  	s.data, err = io.ReadAll(r)
    26  	if err != nil {
    27  		return nil
    28  	}
    29  	r.Close()
    30  	s.length = int64(len(s.data))
    31  	s.openGate = make(chan struct{})
    32  	s.enableBlocking = true
    33  	return &s
    34  }
    35  
    36  // NewOpenGateWithTimeout creates a new open gate with a timeout
    37  func NewOpenGateWithTimeout(r io.ReadCloser, d time.Duration) *SyncedReadCloser {
    38  	s := NewSyncedReadCloser(r)
    39  	s.OpenGateAfter(d)
    40  	return s
    41  }
    42  
    43  // SetOpenGate sets the status of the blocking gate
    44  func (s *SyncedReadCloser) SetOpenGate(status bool) {
    45  	s.enableBlocking = status
    46  }
    47  
    48  // OpenGate opens the gate allowing all requests to be completed
    49  func (s *SyncedReadCloser) OpenGate() {
    50  	s.openGate <- struct{}{}
    51  }
    52  
    53  // OpenGateAfter schedules gate to be opened after a duration
    54  func (s *SyncedReadCloser) OpenGateAfter(d time.Duration) {
    55  	time.AfterFunc(d, func() {
    56  		s.openGate <- struct{}{}
    57  	})
    58  }
    59  
    60  // Seek implements seek method for io.ReadSeeker
    61  func (s *SyncedReadCloser) Seek(offset int64, whence int) (int64, error) {
    62  	var err error
    63  	switch whence {
    64  	case io.SeekStart:
    65  		s.p = 0
    66  	case io.SeekCurrent:
    67  		if s.p+offset < s.length {
    68  			s.p += offset
    69  			break
    70  		}
    71  		err = fmt.Errorf("offset is too big")
    72  	case io.SeekEnd:
    73  		if s.length-offset >= 0 {
    74  			s.p = s.length - offset
    75  			break
    76  		}
    77  		err = fmt.Errorf("offset is too big")
    78  	}
    79  	return s.p, err
    80  }
    81  
    82  // Read implements read method for io.ReadSeeker
    83  func (s *SyncedReadCloser) Read(p []byte) (n int, err error) {
    84  	// If the data fits in the buffer blocks awaiting the sync instruction
    85  	if s.p+int64(len(p)) >= s.length && s.enableBlocking {
    86  		<-s.openGate
    87  	}
    88  	n = copy(p, s.data[s.p:])
    89  	s.p += int64(n)
    90  	if s.p == s.length {
    91  		err = io.EOF
    92  	}
    93  	return n, err
    94  }
    95  
    96  // Close closes an io.ReadSeeker
    97  func (s *SyncedReadCloser) Close() error {
    98  	return nil
    99  }
   100  
   101  // Len returns the length of data in reader
   102  func (s *SyncedReadCloser) Len() int {
   103  	return int(s.length)
   104  }