github.com/sagernet/sing@v0.4.0-beta.19.0.20240518125136-f67a0988a636/common/network/counter.go (about)

     1  package network
     2  
     3  import (
     4  	"io"
     5  )
     6  
     7  type CountFunc func(n int64)
     8  
     9  type ReadCounter interface {
    10  	io.Reader
    11  	UnwrapReader() (io.Reader, []CountFunc)
    12  }
    13  
    14  type WriteCounter interface {
    15  	io.Writer
    16  	UnwrapWriter() (io.Writer, []CountFunc)
    17  }
    18  
    19  type PacketReadCounter interface {
    20  	PacketReader
    21  	UnwrapPacketReader() (PacketReader, []CountFunc)
    22  }
    23  
    24  type PacketWriteCounter interface {
    25  	PacketWriter
    26  	UnwrapPacketWriter() (PacketWriter, []CountFunc)
    27  }
    28  
    29  func UnwrapCountReader(reader io.Reader, countFunc []CountFunc) (io.Reader, []CountFunc) {
    30  	reader = UnwrapReader(reader)
    31  	if counter, isCounter := reader.(ReadCounter); isCounter {
    32  		upstreamReader, upstreamCountFunc := counter.UnwrapReader()
    33  		countFunc = append(countFunc, upstreamCountFunc...)
    34  		return UnwrapCountReader(upstreamReader, countFunc)
    35  	}
    36  	return reader, countFunc
    37  }
    38  
    39  func UnwrapCountWriter(writer io.Writer, countFunc []CountFunc) (io.Writer, []CountFunc) {
    40  	writer = UnwrapWriter(writer)
    41  	if counter, isCounter := writer.(WriteCounter); isCounter {
    42  		upstreamWriter, upstreamCountFunc := counter.UnwrapWriter()
    43  		countFunc = append(countFunc, upstreamCountFunc...)
    44  		return UnwrapCountWriter(upstreamWriter, countFunc)
    45  	}
    46  	return writer, countFunc
    47  }
    48  
    49  func UnwrapCountPacketReader(reader PacketReader, countFunc []CountFunc) (PacketReader, []CountFunc) {
    50  	reader = UnwrapPacketReader(reader)
    51  	if counter, isCounter := reader.(PacketReadCounter); isCounter {
    52  		upstreamReader, upstreamCountFunc := counter.UnwrapPacketReader()
    53  		countFunc = append(countFunc, upstreamCountFunc...)
    54  		return UnwrapCountPacketReader(upstreamReader, countFunc)
    55  	}
    56  	return reader, countFunc
    57  }
    58  
    59  func UnwrapCountPacketWriter(writer PacketWriter, countFunc []CountFunc) (PacketWriter, []CountFunc) {
    60  	writer = UnwrapPacketWriter(writer)
    61  	if counter, isCounter := writer.(PacketWriteCounter); isCounter {
    62  		upstreamWriter, upstreamCountFunc := counter.UnwrapPacketWriter()
    63  		countFunc = append(countFunc, upstreamCountFunc...)
    64  		return UnwrapCountPacketWriter(upstreamWriter, countFunc)
    65  	}
    66  	return writer, countFunc
    67  }