github.com/codingeasygo/util@v0.0.0-20231206062002-1ce2f004b7d9/xio/prefix.go (about)

     1  package xio
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  	"time"
     7  )
     8  
     9  // PrefixReader provide feature to prefix read data from Reader
    10  type PrefixReader struct {
    11  	io.Reader
    12  	Prefix []byte
    13  }
    14  
    15  // NewPrefixReader will return new PrefixReader
    16  func NewPrefixReader(base io.Reader) (prefix *PrefixReader) {
    17  	prefix = &PrefixReader{Reader: base}
    18  	return
    19  }
    20  
    21  func (p *PrefixReader) Read(b []byte) (n int, err error) {
    22  	if p.Prefix != nil {
    23  		n = copy(b, p.Prefix)
    24  		p.Prefix = p.Prefix[n:]
    25  		if len(p.Prefix) < 1 {
    26  			p.Prefix = nil
    27  		}
    28  		return
    29  	}
    30  	readed, err := p.Reader.Read(b[n:])
    31  	n += readed
    32  	return
    33  }
    34  
    35  // PreRead will read prefix size data before read loop start
    36  func (p *PrefixReader) PreRead(size int) (data []byte, err error) {
    37  	data = make([]byte, size)
    38  	err = FullBuffer(p.Reader, data, uint32(size), nil)
    39  	if err == nil {
    40  		p.Prefix = data
    41  	}
    42  	return
    43  }
    44  
    45  func (p *PrefixReader) String() string {
    46  	return RemoteAddr(p.Reader)
    47  }
    48  
    49  // PrefixReadWriteCloser is prefix read implement
    50  type PrefixReadWriteCloser struct {
    51  	io.ReadWriteCloser
    52  	PrefixReader
    53  }
    54  
    55  // NewPrefixReadWriteCloser will return new PrefixReadWriteCloser
    56  func NewPrefixReadWriteCloser(base io.ReadWriteCloser) (prefix *PrefixReadWriteCloser) {
    57  	prefix = &PrefixReadWriteCloser{}
    58  	prefix.ReadWriteCloser = base
    59  	prefix.PrefixReader.Reader = base
    60  	return
    61  }
    62  
    63  func (p *PrefixReadWriteCloser) Read(b []byte) (n int, err error) {
    64  	n, err = p.PrefixReader.Read(b)
    65  	return
    66  }
    67  
    68  func (p *PrefixReadWriteCloser) String() string {
    69  	return RemoteAddr(p.ReadWriteCloser)
    70  }
    71  
    72  // Network is net.Addr implement
    73  func (p *PrefixReadWriteCloser) Network() string {
    74  	return "prefix"
    75  }
    76  
    77  // LocalAddr returns the local network address.
    78  func (p *PrefixReadWriteCloser) LocalAddr() net.Addr {
    79  	if conn, ok := p.ReadWriteCloser.(net.Conn); ok {
    80  		return conn.LocalAddr()
    81  	}
    82  	return p
    83  }
    84  
    85  // RemoteAddr returns the remote network address.
    86  func (p *PrefixReadWriteCloser) RemoteAddr() net.Addr {
    87  	if conn, ok := p.ReadWriteCloser.(net.Conn); ok {
    88  		return conn.RemoteAddr()
    89  	}
    90  	return p
    91  }
    92  
    93  // SetDeadline sets the read and write deadlines associated
    94  func (p *PrefixReadWriteCloser) SetDeadline(t time.Time) error {
    95  	if conn, ok := p.ReadWriteCloser.(net.Conn); ok {
    96  		return conn.SetDeadline(t)
    97  	}
    98  	return nil
    99  }
   100  
   101  // SetReadDeadline sets the deadline for future Read calls
   102  func (p *PrefixReadWriteCloser) SetReadDeadline(t time.Time) error {
   103  	if conn, ok := p.ReadWriteCloser.(net.Conn); ok {
   104  		return conn.SetReadDeadline(t)
   105  	}
   106  	return nil
   107  }
   108  
   109  // SetWriteDeadline sets the deadline for future Write calls
   110  func (p *PrefixReadWriteCloser) SetWriteDeadline(t time.Time) error {
   111  	if conn, ok := p.ReadWriteCloser.(net.Conn); ok {
   112  		return conn.SetWriteDeadline(t)
   113  	}
   114  	return nil
   115  }
   116  
   117  // PrefixConn is net.Conn implement for prefix read data
   118  type PrefixConn struct {
   119  	net.Conn
   120  	PrefixReader
   121  }
   122  
   123  // NewPrefixConn will return newPrefixConn
   124  func NewPrefixConn(conn net.Conn) (prefix *PrefixConn) {
   125  	prefix = &PrefixConn{}
   126  	prefix.Conn = conn
   127  	prefix.PrefixReader.Reader = conn
   128  	return
   129  }
   130  
   131  func (p *PrefixConn) Read(b []byte) (n int, err error) {
   132  	n, err = p.PrefixReader.Read(b)
   133  	return
   134  }
   135  
   136  func (p *PrefixConn) String() string {
   137  	return RemoteAddr(p.Conn)
   138  }
   139  
   140  // LocalAddr returns the local network address.
   141  func (p *PrefixConn) LocalAddr() net.Addr {
   142  	return p.Conn.LocalAddr()
   143  }
   144  
   145  // RemoteAddr returns the remote network address.
   146  func (p *PrefixConn) RemoteAddr() net.Addr {
   147  	return p.Conn.RemoteAddr()
   148  }
   149  
   150  // SetDeadline sets the read and write deadlines associated
   151  func (p *PrefixConn) SetDeadline(t time.Time) error {
   152  	return p.Conn.SetDeadline(t)
   153  }
   154  
   155  // SetReadDeadline sets the deadline for future Read calls
   156  func (p *PrefixConn) SetReadDeadline(t time.Time) error {
   157  	return p.Conn.SetReadDeadline(t)
   158  }
   159  
   160  // SetWriteDeadline sets the deadline for future Write calls
   161  func (p *PrefixConn) SetWriteDeadline(t time.Time) error {
   162  	return p.Conn.SetWriteDeadline(t)
   163  }